Clean code

Author: Balaji Veeramani

src/sagemaker/serializers.py

@@ -196,7 +196,14 @@ def serialize(self, data):
 class IdentitySerializer(BaseSerializer):
     """Serialize data by returning data without modification."""
 
-    CONTENT_TYPE = "application/json"
+    def __init__(self, content_type="application/octet-stream"):
+        """Initialize the ``content_type`` attribute.
+
+        Args:
+            content_type (str): The MIME type of the serialized data (default:
+                "application/octet-stream").
+        """
+        self.content_type = content_type
 
     def serialize(self, data):
         """Return data without modification.
@@ -209,6 +216,11 @@ def serialize(self, data):
         """
         return data
 
+    @property
+    def CONTENT_TYPE(self):
+        """The MIME type of the data sent to the inference endpoint."""
+        return self.content_type
+
 
 class JSONLinesSerializer(BaseSerializer):
     """Serialize data to a JSON Lines formatted string."""

tests/integ/test_tfs.py

@@ -23,8 +23,9 @@
 import sagemaker.utils
 import tests.integ
 import tests.integ.timeout
+from sagemaker.deserializers import JSONLinesDeserializer
 from sagemaker.tensorflow.model import TensorFlowModel, TensorFlowPredictor
-from sagemaker.serializers import CSVSerializer
+from sagemaker.serializers import CSVSerializer, IdentitySerializer, JSONLinesSerializer
 
 
 @pytest.fixture(scope="module")
@@ -190,7 +191,7 @@ def test_predict_jsons_json_content_type(tfs_predictor):
     predictor = sagemaker.Predictor(
         tfs_predictor.endpoint_name,
         tfs_predictor.sagemaker_session,
-        serializer=sagemaker.serializers.IdentitySerializer(),
+        serializer=IdentitySerializer(content_type="application/json"),
         deserializer=sagemaker.deserializers.JSONDeserializer(),
     )
 
@@ -205,7 +206,7 @@ def test_predict_jsons(tfs_predictor):
     predictor = sagemaker.Predictor(
         tfs_predictor.endpoint_name,
         tfs_predictor.sagemaker_session,
-        serializer=sagemaker.serializer.IdentitySerializer(),
+        serializer=IdentitySerializer(content_type="application/json"),
         deserializer=sagemaker.deserializers.JSONDeserializer(),
     )
 
@@ -220,8 +221,8 @@ def test_predict_jsonlines(tfs_predictor):
     predictor = sagemaker.Predictor(
         tfs_predictor.endpoint_name,
         tfs_predictor.sagemaker_session,
-        serializer=sagemaker.serializers.JSONLinesSerializer(),
-        deserializer=sagemaker.deserializers.JSONLinesDeserializer(),
+        serializer=JSONLinesSerializer(),
+        deserializer=JSONLinesDeserializer(),
     )
 
     result = predictor.predict(input_data)

tests/unit/sagemaker/test_serializers.py

@@ -237,15 +237,17 @@ def test_json_serializer_csv_buffer(json_serializer):
         assert result == validation_value
 
 
-@pytest.fixture
-def identity_serializer():
-    return IdentitySerializer()
-
-
 def test_identity_serializer(identity_serializer):
+    identity_serializer = IdentitySerializer()
     assert identity_serializer.serialize(b"{}") == b"{}"
 
 
+def test_identity_serializer_with_custom_content_type():
+    identity_serializer = IdentitySerializer(content_type="text/csv")
+    assert identity_serializer.serialize(b"a,b\n1,2") == b"a,b\n1,2"
+    assert identity_serializer.CONTENT_TYPE == "text/csv"
+
+
 @pytest.fixture
 def json_lines_serializer():
     return JSONLinesSerializer()