From 02a30334242d0401de679994a524548c453da24b Mon Sep 17 00:00:00 2001
From: Sneha Maniraju <sbagalur@cisco.com>
Date: Mon, 11 Oct 2021 14:27:07 -0700
Subject: [PATCH] Add support for snat_subnet_only extension

Change-Id: I4a0fe3c91cbc9145bffa46105a2447a7a8ee1d01
(cherry picked from commit cedb8ad42ce0c9c24379bca4634b8f704771fcd1)
---
 gbpclient/gbp/v2_0/subnet.py        | 37 +++++++++++++++++++++++++++++
 gbpclient/tests/unit/test_subnet.py |  8 +++++++
 2 files changed, 45 insertions(+)

diff --git a/gbpclient/gbp/v2_0/subnet.py b/gbpclient/gbp/v2_0/subnet.py
index 68dea95..d13a09d 100644
--- a/gbpclient/gbp/v2_0/subnet.py
+++ b/gbpclient/gbp/v2_0/subnet.py
@@ -38,6 +38,11 @@ def _get_attrs_subnet_extension(client_manager, parsed_args, is_create=True):
     if 'apic_active_active_aap_disable' in parsed_args and \
        parsed_args.apic_active_active_aap_disable:
         attrs['apic:active_active_aap'] = False
+    if parsed_args.apic_snat_subnet_only_enable:
+        attrs['apic:snat_subnet_only'] = True
+    if parsed_args.apic_snat_subnet_only_disable:
+        attrs['apic:snat_subnet_only'] = False
+
     return attrs
 
 
@@ -51,6 +56,8 @@ subnet_sdk.Subnet.apic_snat_host_pool = resource.Body(
     'apic:snat_host_pool')
 subnet_sdk.Subnet.apic_active_active_aap = resource.Body(
     'apic:active_active_aap')
+subnet_sdk.Subnet.apic_snat_subnet_only = resource.Body(
+    'apic:snat_subnet_only')
 
 
 class CreateSubnetExtension(hooks.CommandHook):
@@ -86,6 +93,21 @@ class CreateSubnetExtension(hooks.CommandHook):
             help=_("Set APIC active active aap to false\n"
                    "Default value for apic_active_active_aap is False ")
         )
+        parser.add_argument(
+            '--apic-snat-subnet-only-enable',
+            action='store_true',
+            default=None,
+            dest='apic_snat_subnet_only_enable',
+            help=_("Set APIC snat subnet only to true\n"
+                   "Default value for apic_snat_subnet_only is False ")
+        )
+        parser.add_argument(
+            '--apic-snat-subnet-only-disable',
+            action='store_true',
+            dest='apic_snat_subnet_only_disable',
+            help=_("Set APIC snat subnet only to false\n"
+                   "Default value for apic_snat_subnet_only is False ")
+        )
         return parser
 
     def get_epilog(self):
@@ -116,6 +138,21 @@ class SetSubnetExtension(hooks.CommandHook):
             help=_("Set APIC snat host pool to false\n"
                    "Default value for apic_snat_host_pool is False ")
         )
+        parser.add_argument(
+            '--apic-snat-subnet-only-enable',
+            action='store_true',
+            default=None,
+            dest='apic_snat_subnet_only_enable',
+            help=_("Set APIC snat subnet only to true\n"
+                   "Default value for apic_snat_subnet_only is False ")
+        )
+        parser.add_argument(
+            '--apic-snat-subnet-only-disable',
+            action='store_true',
+            dest='apic_snat_subnet_only_disable',
+            help=_("Set APIC snat subnet only to false\n"
+                   "Default value for apic_snat_subnet_only is False ")
+        )
         return parser
 
     def get_epilog(self):
diff --git a/gbpclient/tests/unit/test_subnet.py b/gbpclient/tests/unit/test_subnet.py
index 4051052..8a1b471 100644
--- a/gbpclient/tests/unit/test_subnet.py
+++ b/gbpclient/tests/unit/test_subnet.py
@@ -51,6 +51,7 @@ class TestSubnetCreate(test_subnet.TestSubnet, test_cli20.CLITestV20Base):
             ('network', self._subnet.network_id),
             ('apic_snat_host_pool_enable', None),
             ('apic_active_active_aap_enable', None),
+            ('apic_snat_subnet_only_enable', None),
         ]
         create_ext = subnet_ext.CreateSubnetExtension(self.app)
         parsed_args = self.check_parser_ext(
@@ -71,12 +72,14 @@ class TestSubnetCreate(test_subnet.TestSubnet, test_cli20.CLITestV20Base):
             self._subnet.name,
             "--apic-snat-host-pool-enable",
             "--apic-active-active-aap-enable",
+            "--apic-snat-subnet-only-enable",
         ]
         verifylist = [
             ('name', self._subnet.name),
             ('network', self._subnet.network_id),
             ('apic_snat_host_pool_enable', True),
             ('apic_active_active_aap_enable', True),
+            ('apic_snat_subnet_only_enable', True),
         ]
         create_ext = subnet_ext.CreateSubnetExtension(self.app)
         parsed_args = self.check_parser_ext(
@@ -90,6 +93,7 @@ class TestSubnetCreate(test_subnet.TestSubnet, test_cli20.CLITestV20Base):
             'network_id': self._subnet.network_id,
             'apic:active_active_aap': True,
             'apic:snat_host_pool': True,
+            'apic:snat_subnet_only': True,
         })
 
 
@@ -112,6 +116,7 @@ class TestSubnetSet(test_subnet.TestSubnet, test_cli20.CLITestV20Base):
         verifylist = [
             ('subnet', self._subnet.name),
             ('apic_snat_host_pool_enable', None),
+            ('apic_snat_subnet_only_enable', None),
         ]
         set_ext = subnet_ext.SetSubnetExtension(self.app)
         parsed_args = self.check_parser_ext(
@@ -125,10 +130,12 @@ class TestSubnetSet(test_subnet.TestSubnet, test_cli20.CLITestV20Base):
         arglist = [
             self._subnet.name,
             "--apic-snat-host-pool-disable",
+            "--apic-snat-subnet-only-disable",
         ]
         verifylist = [
             ('subnet', self._subnet.name),
             ('apic_snat_host_pool_disable', True),
+            ('apic_snat_subnet_only_disable', True),
         ]
         set_ext = subnet_ext.SetSubnetExtension(self.app)
         parsed_args = self.check_parser_ext(
@@ -137,6 +144,7 @@ class TestSubnetSet(test_subnet.TestSubnet, test_cli20.CLITestV20Base):
 
         attrs = {
             'apic:snat_host_pool': False,
+            'apic:snat_subnet_only': False,
         }
         self.network.update_subnet.assert_called_with(self._subnet, **attrs)
         self.assertIsNone(result)