Fixes pytorch#772: unit test attribute error - 'bytes' object has no …

…attribute 'metrics'
abishekchiffon · Nov 4, 2020 · 8a8ba8f · 8a8ba8f
1 parent 878d248
commit 8a8ba8f
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 10 deletions.
diff --git a/ts/torch_handler/unit_tests/test_image_classifier.py b/ts/torch_handler/unit_tests/test_image_classifier.py
@@ -14,7 +14,7 @@
 
 @pytest.fixture()
 def model_setup():
-    context = MockContext()
+    context = MockContext(model_name = "mnist")
     with open('ts/torch_handler/unit_tests/models/tmp/images/kitten.jpg', 'rb') as fin:
         image_bytes = fin.read()
     return (context, image_bytes)
@@ -28,9 +28,9 @@ def test_initialize(model_setup):
     return handler
 
 def test_handle(model_setup):
-    _, image_bytes = model_setup
+    context, image_bytes = model_setup
     handler = test_initialize(model_setup)
     test_data = [{'data': image_bytes}] * 2
-    results = handler.handle(test_data, image_bytes)
+    results = handler.handle(test_data, context)
     assert(len(results) == 2)
     assert('tiger_cat' in results[0])
diff --git a/ts/torch_handler/unit_tests/test_image_segmenter.py b/ts/torch_handler/unit_tests/test_image_segmenter.py
@@ -14,7 +14,7 @@
 
 @pytest.fixture()
 def model_setup():
-    context = MockContext()
+    context = MockContext(model_name = "image_segmenter")
     with open('ts/torch_handler/unit_tests/models/tmp/persons.jpg', 'rb') as fin:
         image_bytes = fin.read()
     return (context, image_bytes)
@@ -28,9 +28,9 @@ def test_initialize(model_setup):
     return handler
 
 def test_handle(model_setup):
-    _, image_bytes = model_setup
+    context, image_bytes = model_setup
     handler = test_initialize(model_setup)
     test_data = [{'data': image_bytes}] * 2
-    results = handler.handle(test_data, image_bytes)
+    results = handler.handle(test_data, context)
     assert(len(results) == 2)
     assert(len(results[0]) == 224)
diff --git a/ts/torch_handler/unit_tests/test_object_detector.py b/ts/torch_handler/unit_tests/test_object_detector.py
@@ -14,7 +14,7 @@
 
 @pytest.fixture()
 def model_setup():
-    context = MockContext()
+    context = MockContext(model_name = "object_detector")
     with open('ts/torch_handler/unit_tests/models/tmp/persons.jpg', 'rb') as fin:
         image_bytes = fin.read()
     return (context, image_bytes)
@@ -28,9 +28,9 @@ def test_initialize(model_setup):
     return handler
 
 def test_handle(model_setup):
-    _, image_bytes = model_setup
+    context, image_bytes = model_setup
     handler = test_initialize(model_setup)
     test_data = [{'data': image_bytes}] * 2
-    results = handler.handle(test_data, image_bytes)
+    results = handler.handle(test_data, context)
     assert(len(results) == 2)
     assert('bench' in results[0])
diff --git a/ts/torch_handler/unit_tests/test_utils/mock_context.py b/ts/torch_handler/unit_tests/test_utils/mock_context.py
@@ -2,7 +2,9 @@
 Mocks for adding model context without loading all of Torchserve
 """
 
+import uuid
 import torch
+from ts.metrics.metrics_store import MetricsStore
 
 class MockContext():
     """
@@ -12,7 +14,8 @@ def __init__(self,
                  model_pt_file='model.pt',
                  model_dir='ts/torch_handler/unit_tests/models/tmp',
                  model_file='model.py',
-                 gpu_id='0'):
+                 gpu_id='0',
+                 model_name = "mnist"):
         self.manifest = {
             'model': {
                 'serializedFile': model_pt_file,
@@ -28,3 +31,5 @@ def __init__(self,
 
         if torch.cuda.is_available() and gpu_id:
             self.system_properties['gpu_id'] = gpu_id
+
+        self.metrics = MetricsStore(uuid.uuid4(), model_name)