Add SetAttribute support to the client

kmip/pie/client.py

@@ -102,6 +102,7 @@ def __init__(self,
         self.logger = logging.getLogger(__name__)
 
         self.attribute_factory = attributes.AttributeFactory()
+        self.attribute_value_factory = self.attribute_factory.value_factory
         self.object_factory = factory.ObjectFactory()
 
         # TODO (peter-hamilton) Consider adding validation checks for inputs.
@@ -431,6 +432,46 @@ def delete_attribute(self, unique_identifier=None, **kwargs):
 
         return response_payload.unique_identifier, response_payload.attribute
 
+    @is_connected
+    def set_attribute(self, unique_identifier=None, **kwargs):
+        """
+        Set an attribute on a KMIP managed object.
+
+        Args:
+            unique_identifier (string): The ID of the managed object.
+            **kwargs (various): A placeholder for attribute-related fields.
+                Supported parameters include:
+                    attribute_name (string): The name of the attribute being
+                        set. Required.
+                    attribute_value (various): The value of the attribute
+                        being set. Required.
+
+                Here is an example. To set an object's 'sensitive' attribute
+                to True, specify:
+                    attribute_name='Sensitive'
+                    attribute_value=True
+
+                For a list of all supported attributes, see the
+                AttributeValueFactory.
+
+        Returns:
+            string: The ID of the managed object the attribute was set on.
+        """
+        a = self.attribute_value_factory.create_attribute_value_by_enum(
+            enums.convert_attribute_name_to_tag(kwargs.get("attribute_name")),
+            kwargs.get("attribute_value")
+        )
+        request_payload = payloads.SetAttributeRequestPayload(
+            unique_identifier=unique_identifier,
+            new_attribute=cobjects.NewAttribute(attribute=a)
+        )
+        response_payload = self.proxy.send_request_payload(
+            enums.Operation.SET_ATTRIBUTE,
+            request_payload
+        )
+
+        return response_payload.unique_identifier
+
     @is_connected
     def register(self, managed_object):
         """

kmip/services/kmip_client.py

@@ -341,14 +341,20 @@ def send_request_payload(self, operation, payload, credential=None):
             )
 
         # TODO (peterhamilton) For now limit this to the new DeleteAttribute
-        # operation. Migrate over existing operations to use this method
-        # instead.
+        # and SetAttribute operations. Migrate over existing operations to use
+        # this method instead.
         if operation == enums.Operation.DELETE_ATTRIBUTE:
             if not isinstance(payload, payloads.DeleteAttributeRequestPayload):
                 raise TypeError(
                     "The request payload for the DeleteAttribute operation "
                     "must be a DeleteAttributeRequestPayload object."
                 )
+        elif operation == enums.Operation.SET_ATTRIBUTE:
+            if not isinstance(payload, payloads.SetAttributeRequestPayload):
+                raise TypeError(
+                    "The request payload for the SetAttribute operation must "
+                    "be a SetAttributeRequestPayload object."
+                )
 
         batch_item = messages.RequestBatchItem(
             operation=operation,
@@ -388,6 +394,15 @@ def send_request_payload(self, operation, payload, credential=None):
                     "Invalid response payload received for the "
                     "DeleteAttribute operation."
                 )
+        elif batch_item.operation.value == enums.Operation.SET_ATTRIBUTE:
+            if not isinstance(
+                batch_item.response_payload,
+                payloads.SetAttributeRequestPayload
+            ):
+                raise exceptions.InvalidMessage(
+                    "Invalid response payload received for the SetAttribute "
+                    "operation."
+                )
 
         return batch_item.response_payload
 

kmip/tests/unit/pie/test_client.py

@@ -25,6 +25,7 @@
 from kmip.core.factories import attributes
 from kmip.core.messages import contents
 from kmip.core.messages import payloads
+from kmip.core import primitives
 from kmip.core.primitives import DateTime
 
 from kmip.services.kmip_client import KMIPProxy
@@ -794,6 +795,43 @@ def test_delete_attribute(self):
             self.assertEqual("1", unique_identifier)
             self.assertIsNone(attribute)
 
+    @mock.patch(
+        "kmip.pie.client.KMIPProxy",
+        mock.MagicMock(spec_set=KMIPProxy)
+    )
+    def test_set_attribute(self):
+        """
+        Test that the client can set an attribute.
+        """
+        request_payload = payloads.SetAttributeRequestPayload(
+            unique_identifier="1",
+            new_attribute=obj.NewAttribute(
+                attribute=primitives.Boolean(
+                    value=True,
+                    tag=enums.Tags.SENSITIVE
+                )
+            )
+        )
+        response_payload = payloads.SetAttributeResponsePayload(
+            unique_identifier="1"
+        )
+
+        with ProxyKmipClient() as client:
+            client.proxy.send_request_payload.return_value = response_payload
+
+            unique_identifier = client.set_attribute(
+                "1",
+                attribute_name="Sensitive",
+                attribute_value=True
+            )
+
+            args = (
+                enums.Operation.SET_ATTRIBUTE,
+                request_payload
+            )
+            client.proxy.send_request_payload.assert_called_with(*args)
+            self.assertEqual("1", unique_identifier)
+
     @mock.patch(
         'kmip.pie.client.KMIPProxy', mock.MagicMock(spec_set=KMIPProxy)
     )

kmip/tests/unit/services/test_kmip_client.py

@@ -843,6 +843,18 @@ def test_send_request_payload_mismatch_operation_payload(self):
             *args
         )
 
+        args = (
+            OperationEnum.SET_ATTRIBUTE,
+            payloads.CreateRequestPayload()
+        )
+        self.assertRaisesRegex(
+            TypeError,
+            "The request payload for the SetAttribute operation must be a "
+            "SetAttributeRequestPayload object.",
+            self.client.send_request_payload,
+            *args
+        )
+
     @mock.patch(
         "kmip.services.kmip_client.KMIPProxy._build_request_message"
     )
@@ -957,7 +969,6 @@ def test_send_request_payload_mismatch_response_payload(
                 attribute_index=2
             )
         )
-
         self.assertRaisesRegex(
             exceptions.InvalidMessage,
             "Invalid response payload received for the DeleteAttribute "
@@ -966,6 +977,32 @@ def test_send_request_payload_mismatch_response_payload(
             *args
         )
 
+        batch_item = ResponseBatchItem(
+            operation=Operation(OperationEnum.SET_ATTRIBUTE),
+            result_status=ResultStatus(ResultStatusEnum.SUCCESS),
+            response_payload=response_payload
+        )
+        send_mock.return_value = ResponseMessage(batch_items=[batch_item])
+        args = (
+            OperationEnum.SET_ATTRIBUTE,
+            payloads.SetAttributeRequestPayload(
+                unique_identifier="1",
+                new_attribute=objects.NewAttribute(
+                    attribute=primitives.Boolean(
+                        value=True,
+                        tag=enums.Tags.SENSITIVE
+                    )
+                )
+            )
+        )
+        self.assertRaisesRegex(
+            exceptions.InvalidMessage,
+            "Invalid response payload received for the SetAttribute "
+            "operation.",
+            self.client.send_request_payload,
+            *args
+        )
+
     @mock.patch(
         "kmip.services.kmip_client.KMIPProxy._build_request_message"
     )