Update test_predictor.py

Author: Balaji Veeramani

tests/unit/test_predictor.py

@@ -17,6 +17,7 @@
 import pytest
 from mock import Mock, call, patch
 
+from sagemaker.deserializers import CSVDeserializer, StringDeserializer
 from sagemaker.predictor import Predictor
 from sagemaker.serializers import JSONSerializer, CSVSerializer
 
@@ -132,7 +133,7 @@ def json_sagemaker_session():
     response_body.close = Mock("close", return_value=None)
     ims.sagemaker_runtime_client.invoke_endpoint = Mock(
         name="invoke_endpoint",
-        return_value={"Body": response_body, "ContentType": DEFAULT_CONTENT_TYPE},
+        return_value={"Body": response_body, "ContentType": "application/json"},
     )
     return ims
 
@@ -169,7 +170,7 @@ def ret_csv_sagemaker_session():
     ims.sagemaker_client.describe_endpoint_config = Mock(return_value=ENDPOINT_CONFIG_DESC)
 
     response_body = Mock("body")
-    response_body.read = Mock("read", return_value=CSV_RETURN_VALUE)
+    response_body.read = Mock("read", return_value=bytes(CSV_RETURN_VALUE, "utf-8"))
     response_body.close = Mock("close", return_value=None)
     ims.sagemaker_runtime_client.invoke_endpoint = Mock(
         name="invoke_endpoint",
@@ -180,23 +181,44 @@ def ret_csv_sagemaker_session():
 
 def test_predict_call_with_csv():
     sagemaker_session = ret_csv_sagemaker_session()
-    predictor = Predictor(ENDPOINT, sagemaker_session, serializer=CSVSerializer())
+    predictor = Predictor(ENDPOINT, sagemaker_session, serializer=CSVSerializer(), deserializer=CSVDeserializer())
 
     data = [1, 2]
     result = predictor.predict(data)
 
     assert sagemaker_session.sagemaker_runtime_client.invoke_endpoint.called
 
     expected_request_args = {
-        "Accept": DEFAULT_ACCEPT,
+        "Accept": CSV_CONTENT_TYPE,
+        "Body": "1,2",
+        "ContentType": CSV_CONTENT_TYPE,
+        "EndpointName": ENDPOINT,
+    }
+    call_args, kwargs = sagemaker_session.sagemaker_runtime_client.invoke_endpoint.call_args
+    assert kwargs == expected_request_args
+
+    assert result == [["1", "2", "3"]]
+
+
+def test_predict_call_with_multiple_accept_types():
+    sagemaker_session = ret_csv_sagemaker_session()
+    predictor = Predictor(ENDPOINT, sagemaker_session, serializer=CSVSerializer(), deserializer=StringDeserializer())
+
+    data = [1, 2]
+    result = predictor.predict(data)
+
+    assert sagemaker_session.sagemaker_runtime_client.invoke_endpoint.called
+
+    expected_request_args = {
+        "Accept": "application/json, text/csv",
         "Body": "1,2",
         "ContentType": CSV_CONTENT_TYPE,
         "EndpointName": ENDPOINT,
     }
     call_args, kwargs = sagemaker_session.sagemaker_runtime_client.invoke_endpoint.call_args
     assert kwargs == expected_request_args
 
-    assert result == CSV_RETURN_VALUE
+    assert result == "1,2,3\r\n"
 
 
 @patch("sagemaker.predictor.name_from_base")