PR Feedback

Author: dk19y

test/container/1.5.0/Dockerfile.pytorch

@@ -18,7 +18,7 @@ RUN apt-get update \
     libsm6 \
     libxext6 \
     libxrender-dev \
-    openjdk-11-jdk \
+    openjdk-11-jdk-headless \
     && rm -rf /var/lib/apt/lists/*
 
 RUN conda install -c conda-forge opencv==4.0.1 \

test/unit/test_default_inference_handler.py

@@ -15,6 +15,7 @@
 import csv
 import json
 
+import mock
 import numpy as np
 import pytest
 import torch
@@ -177,3 +178,32 @@ def test_default_output_fn_gpu(inference_handler):
     output = inference_handler.default_output_fn(tensor_gpu, content_types.CSV)
 
     assert "1,2,3\n4,5,6\n".encode("utf-8") == output
+
+def test_eia_default_model_fn(eia_inference_handler):
+    with mock.patch("sagemaker_pytorch_serving_container.default_pytorch_inference_handler.os") as mock_os:
+        mock_os.getenv.return_value = "true"
+        mock_os.path.join.return_value = "model_dir"
+        mock_os.path.exists.return_value = True
+        with mock.patch("torch.jit.load") as mock_torch:
+            mock_torch.return_value = DummyModel()
+            model = eia_inference_handler.default_model_fn("model_dir")
+    assert model is not None
+
+
+def test_eia_default_model_fn_error(eia_inference_handler):
+    with mock.patch("sagemaker_pytorch_serving_container.default_pytorch_inference_handler.os") as mock_os:
+        mock_os.getenv.return_value = "true"
+        mock_os.path.join.return_value = "model_dir"
+        mock_os.path.exists.return_value = False
+        with pytest.raises(FileNotFoundError):
+            eia_inference_handler.default_model_fn("model_dir")
+
+
+def test_eia_default_predict_fn(eia_inference_handler, tensor):
+    model = DummyModel()
+    with mock.patch("sagemaker_pytorch_serving_container.default_pytorch_inference_handler.os") as mock_os:
+        mock_os.getenv.return_value = "true"
+        with mock.patch("torch.jit.optimized_execution") as mock_torch:
+            mock_torch.__enter__.return_value = "dummy"
+            eia_inference_handler.default_predict_fn(tensor, model)
+        mock_torch.assert_called_once()