Change error classes used by _default_input_fn() and _default_output_fn() (#65)

This will make the containers return more appropriate HTTP status codes when
_default_input_fn() receives an unsupported content type or when
_default_output_fn() receives an unsupported accpet type.

Author: laurenyu

src/tf_container/serve.py

@@ -25,7 +25,9 @@
 from tf_container import proxy_client
 from six import StringIO
 import csv
-from container_support.serving import JSON_CONTENT_TYPE, CSV_CONTENT_TYPE, OCTET_STREAM_CONTENT_TYPE, ANY_CONTENT_TYPE
+from container_support.serving import UnsupportedContentTypeError, UnsupportedAcceptTypeError, \
+                                      JSON_CONTENT_TYPE, CSV_CONTENT_TYPE, \
+                                      OCTET_STREAM_CONTENT_TYPE, ANY_CONTENT_TYPE
 from tf_container.run import logger
 import time
 
@@ -227,19 +229,17 @@ def _default_output_fn(data, accepts):
         if accepts == OCTET_STREAM_CONTENT_TYPE:
             return data.SerializeToString()
 
-        raise ValueError('invalid accept type {}'.format(accepts))
+        raise UnsupportedAcceptTypeError('invalid accept type {}'.format(accepts))
 
     def _default_input_fn(self, serialized_data, content_type):
         if content_type == JSON_CONTENT_TYPE:
-            data = self._parse_json_request(serialized_data)
-        elif content_type == CSV_CONTENT_TYPE:
-            data = self._parse_csv_request(serialized_data)
-        elif content_type == OCTET_STREAM_CONTENT_TYPE:
-            data = self.proxy_client.parse_request(serialized_data)
-        else:
-            raise ValueError("Unsupported content-type {}".format(content_type))
+            return self._parse_json_request(serialized_data)
+        if content_type == CSV_CONTENT_TYPE:
+            return self._parse_csv_request(serialized_data)
+        if content_type == OCTET_STREAM_CONTENT_TYPE:
+            return self.proxy_client.parse_request(serialized_data)
 
-        return data
+        raise UnsupportedContentTypeError('Unsupported content-type {}'.format(content_type))
 
     @classmethod
     def from_module(cls, m, grpc_proxy_client):

test/unit/test_serve.py

@@ -1,14 +1,14 @@
 #  Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
-#  
+#
 #  Licensed under the Apache License, Version 2.0 (the "License").
 #  You may not use this file except in compliance with the License.
 #  A copy of the License is located at
-#  
+#
 #      http://www.apache.org/licenses/LICENSE-2.0
-#  
-#  or in the "license" file accompanying this file. This file is distributed 
-#  on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either 
-#  express or implied. See the License for the specific language governing 
+#
+#  or in the "license" file accompanying this file. This file is distributed
+#  on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+#  express or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 
 import json
@@ -18,6 +18,8 @@
 from test.unit.utils import mock_import_modules
 from types import ModuleType
 
+from container_support.serving import UnsupportedAcceptTypeError, UnsupportedContentTypeError
+
 JSON_CONTENT_TYPE = "application/json"
 
 
@@ -279,6 +281,20 @@ def test_wait_model_to_load(proxy_client, serve):
     client.cache_prediction_metadata.assert_called_once_with()
 
 
+def test_transformer_default_output_fn_unsupported_type(serve):
+    accept_type = 'fake/accept-type'
+
+    with pytest.raises(UnsupportedAcceptTypeError):
+        serve.Transformer._default_output_fn(None, accept_type)
+
+
+def test_transformer_default_input_fn_unsupported_type(serve):
+    content_type = 'fake/content-type'
+
+    with pytest.raises(UnsupportedContentTypeError):
+        serve.Transformer(None)._default_input_fn(None, content_type)
+
+
 class DummyTransformer(object):
     def transform(self, content, mimetype):
         if content.startswith('500'):