Fix unit tests

Author: saimidu

test/unit/test_default_inference_handler.py

@@ -14,6 +14,7 @@
 
 import csv
 import json
+import os
 
 import mock
 import numpy as np
@@ -30,8 +31,7 @@
 
 
 class DummyModel(nn.Module):
-
-    def __init__(self, ):
+    def __init__(self):
         super(DummyModel, self).__init__()
 
     def forward(self, x):
@@ -58,7 +58,9 @@ def eia_inference_handler():
 
 
 def test_default_model_fn(inference_handler):
-    with mock.patch("sagemaker_pytorch_serving_container.default_pytorch_inference_handler.os") as mock_os:
+    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
@@ -69,7 +71,9 @@ def test_default_model_fn(inference_handler):
 
 
 def test_default_model_fn_unknown_name(inference_handler):
-    with mock.patch("sagemaker_pytorch_serving_container.default_pytorch_inference_handler.os") as mock_os:
+    with mock.patch(
+        "sagemaker_pytorch_serving_container.default_pytorch_inference_handler.os"
+    ) as mock_os:
         mock_os.getenv.return_value = "false"
         mock_os.path.join.return_value = "model_dir"
         mock_os.path.isfile.return_value = True
@@ -80,18 +84,25 @@ def test_default_model_fn_unknown_name(inference_handler):
     assert model is not None
 
 
-@pytest.mark.parametrize("listdir_return_value", [["abcd.py", "efgh.txt", "ijkl.bin"], ["abcd.pt", "efgh.pth"]])
+@pytest.mark.parametrize(
+    "listdir_return_value", [["abcd.py", "efgh.txt", "ijkl.bin"], ["abcd.pt", "efgh.pth"]]
+)
 def test_default_model_fn_no_model_file(inference_handler, listdir_return_value):
-    with mock.patch("sagemaker_pytorch_serving_container.default_pytorch_inference_handler.os") as mock_os:
+    with mock.patch(
+        "sagemaker_pytorch_serving_container.default_pytorch_inference_handler.os"
+    ) as mock_os:
         mock_os.getenv.return_value = "false"
         mock_os.path.join.return_value = "model_dir"
+        mock_os.path.exists.return_value = False
         mock_os.path.isfile.return_value = True
         mock_os.listdir.return_value = listdir_return_value
+        mock_os.path.splitext = os.path.splitext
         with mock.patch("torch.jit.load") as mock_torch_load:
             mock_torch_load.return_value = DummyModel()
-            with pytest.raises(ValueError):
+            with pytest.raises(
+                ValueError, match=r"Exactly one .pth or .pt file is required for PyTorch models: .*"
+            ):
                 model = inference_handler.default_model_fn("model_dir")
-    assert model is not None
 
 
 def test_default_input_fn_json(inference_handler, tensor):
@@ -220,7 +231,9 @@ def test_default_output_fn_gpu(inference_handler):
 
 
 def test_eia_default_model_fn(eia_inference_handler):
-    with mock.patch("sagemaker_pytorch_serving_container.default_pytorch_inference_handler.os") as mock_os:
+    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
@@ -231,7 +244,9 @@ def test_eia_default_model_fn(eia_inference_handler):
 
 
 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:
+    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
@@ -241,7 +256,9 @@ def test_eia_default_model_fn_error(eia_inference_handler):
 
 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:
+    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"