test

Author: Qian Chen

src/sagemaker_mxnet_serving_container/handler_service.py

@@ -44,7 +44,11 @@ def __init__(self):
 
     @staticmethod
     def _user_module_transformer(model_dir=environment.model_dir):
-        user_module = importlib.import_module(environment.Environment().module_name)
+        try:
+            user_module = importlib.import_module(environment.Environment().module_name)
+        except ModuleNotFoundError as e:
+            logging.error("import_module exception: {}".format(e))
+            raise ValueError('import_module exception: {}'.format(e))
 
         if hasattr(user_module, 'transform_fn'):
             return Transformer(default_inference_handler=DefaultMXNetInferenceHandler())

test/unit/test_handler_service.py

@@ -103,3 +103,14 @@ def test_user_module_unsupported(import_module, env):
 
     import_module.assert_called_once_with(MODULE_NAME)
     e.match('Unsupported model type')
+
+
+@patch('sagemaker_inference.environment.Environment')
+def test_user_module_notfound(env):
+    env.return_value.module_name = MODULE_NAME
+
+    with pytest.raises(ValueError) as e:
+        HandlerService._user_module_transformer()
+
+    importlib.import_module.assert_called_once_with(MODULE_NAME)
+    e.match('import_module exception')