Allow user-provided update_payload_callback function

Manager API: add update_payload_cbk argument in update method in
             NsxLibIPSet class

Policy API: add update_payload_cbk argument in update_with_conditions
            method in NsxPolicyGroupApi class


Change-Id: I89ae7eb564b09fe15934efa48ab10b69e68de70a

vmware_nsxlib/tests/unit/v3/policy/test_resources.py

@@ -576,6 +576,57 @@ class TestPolicyGroup(NsxPolicyLibTestCase):
             update_call.assert_called_once_with(
                 group_def.get_resource_path(), updated_group)
 
+    def test_update_with_conditions_callback(self):
+
+        def update_payload_cbk(revised_payload, payload):
+            revised_ips = revised_payload["expression"][0]["ip_addresses"]
+            new_ips = payload["conditions"][0].ip_addresses
+            updated_ips = revised_ips + new_ips
+            payload["conditions"] = [core_defs.IPAddressExpression(
+                updated_ips)]
+
+        domain_id = '111'
+        group_id = '222'
+        name = 'name'
+        new_name = 'new name'
+        description = 'desc'
+        new_description = 'new desc'
+
+        ips1 = ["1.1.1.1"]
+        ips2 = ["2.2.2.2"]
+        cond1_def = core_defs.IPAddressExpression(ips1)
+        cond2_def = core_defs.IPAddressExpression(ips2)
+        updated_cond_def = core_defs.IPAddressExpression(ips1 + ips2)
+
+        original_group = {
+            'id': group_id,
+            'resource_type': 'Group',
+            'display_name': name,
+            'description': description,
+            'expression': [cond1_def.get_obj_dict()]}
+        updated_group = {
+            'id': group_id,
+            'resource_type': 'Group',
+            'display_name': new_name,
+            'description': new_description,
+            'expression': [updated_cond_def.get_obj_dict()]}
+        group_def = core_defs.GroupDef(
+            domain_id=domain_id,
+            group_id=group_id,
+            conditions=[cond2_def],
+            tenant=TEST_TENANT)
+        with mock.patch.object(self.policy_api, "get",
+                               return_value=original_group),\
+            mock.patch.object(self.policy_api.client,
+                              "update") as update_call:
+            self.resourceApi.update_with_conditions(
+                domain_id, group_id, name=new_name,
+                description=new_description,
+                conditions=[cond2_def], tenant=TEST_TENANT,
+                update_payload_cbk=update_payload_cbk)
+            update_call.assert_called_once_with(
+                group_def.get_resource_path(), updated_group)
+
     def test_unset(self):
         domain_id = '111'
         group_id = '222'

vmware_nsxlib/tests/unit/v3/test_security.py

@@ -368,6 +368,30 @@ class TestNsxLibIPSet(nsxlib_testcase.NsxClientTestCase):
                 resource = 'ip-sets/%s' % fake_ip_set['id']
                 update.assert_called_with(resource, data, headers=None)
 
+    def test_update_ip_set_callback(self):
+
+        def update_payload_cbk(revised_payload, payload):
+            payload['ip_addresses'] = (revised_payload['ip_addresses'] +
+                                       payload['ip_addresses'])
+
+        fake_ip_set = test_constants.FAKE_IP_SET.copy()
+        new_ip_addresses = ['10.0.0.0']
+        updated_ip_addresses = fake_ip_set['ip_addresses'] + new_ip_addresses
+        data = {
+            'id': fake_ip_set['id'],
+            'display_name': fake_ip_set['display_name'],
+            'ip_addresses': updated_ip_addresses,
+            'resource_type': 'IPSet'
+        }
+        with mock.patch.object(self.nsxlib.client, 'get',
+                               return_value=fake_ip_set):
+            with mock.patch.object(self.nsxlib.client, 'update') as update:
+                self.nsxlib.ip_set.update(
+                    fake_ip_set['id'], ip_addresses=new_ip_addresses,
+                    update_payload_cbk=update_payload_cbk)
+                resource = 'ip-sets/%s' % fake_ip_set['id']
+                update.assert_called_with(resource, data, headers=None)
+
 
 class TestNsxLibNSGroup(nsxlib_testcase.NsxClientTestCase):
     """Tests for vmware_nsxlib.v3.security.NsxLibNSGroup"""

vmware_nsxlib/v3/policy/core_resources.py

@@ -523,7 +523,8 @@ class NsxPolicyGroupApi(NsxPolicyResourceBase):
     def update_with_conditions(
         self, domain_id, group_id,
         name=IGNORE, description=IGNORE, conditions=IGNORE,
-        tags=IGNORE, tenant=constants.POLICY_INFRA_TENANT):
+        tags=IGNORE, tenant=constants.POLICY_INFRA_TENANT,
+        update_payload_cbk=None):
         group_def = self._init_def(domain_id=domain_id,
                                    group_id=group_id,
                                    name=name,
@@ -539,6 +540,17 @@ class NsxPolicyGroupApi(NsxPolicyResourceBase):
         def _update():
             # Get the current data of group
             group = self.policy_api.get(group_def)
+            if update_payload_cbk:
+                # The update_payload_cbk function takes two arguments.
+                # The first one is the result from the internal GET request.
+                # The second one is a dict of user-provided attributes,
+                # which can be changed inside the callback function and
+                # used as the new payload for the following PUT request.
+                # For example, users want to combine the new conditions
+                # passed to update_with_conditions() with the original
+                # conditions retrieved from the internal GET request
+                # instead of overriding the original conditions.
+                update_payload_cbk(group, group_def.attrs)
             group_def.set_obj_dict(group)
             body = group_def.get_obj_dict()
             # Update the entire group at the NSX

vmware_nsxlib/v3/security.py

@@ -648,7 +648,16 @@ class NsxLibIPSet(utils.NsxLibApiBase):
         return self.client.create(self.get_path(), body)
 
     def update(self, ip_set_id, display_name=None, description=None,
-               ip_addresses=None, tags_update=None):
+               ip_addresses=None, tags_update=None, update_payload_cbk=None):
+        # The update_payload_cbk function takes two arguments.
+        # The first one is the result from the internal GET request.
+        # The second one is a dict of user-provided attributes,
+        # which can be changed inside the callback function and
+        # used as the new payload for the following PUT request.
+        # For example, users want to combine the new ip_addresses
+        # passed to update() with the original ip_addresses retrieved
+        # from the internal GET request instead of overriding the
+        # original ip_addresses.
         ip_set = {}
         if tags_update:
             ip_set['tags_update'] = tags_update
@@ -659,7 +668,8 @@ class NsxLibIPSet(utils.NsxLibApiBase):
         if ip_addresses is not None:
             ip_set['ip_addresses'] = ip_addresses
         return self._update_resource(self.get_path(ip_set_id),
-                                     ip_set, retry=True)
+                                     ip_set, retry=True,
+                                     update_payload_cbk=update_payload_cbk)
 
     def read(self, ip_set_id):
         return self.client.get('ip-sets/%s' % ip_set_id)