Add unit tests via pytest

tests/test_lambda.py

@@ -0,0 +1,57 @@
+import os
+import sys
+import pytest
+import json
+import base64
+
+
+current_dir = os.path.dirname(os.path.abspath(__file__))
+parent_dir = os.path.dirname(current_dir)
+sys.path.insert(0, os.path.dirname(parent_dir))
+
+from lambda_function import lambda_handler
+
+
+@pytest.fixture
+def event():
+    # Get the directory of the current test file
+    test_dir = os.path.dirname(os.path.abspath(__file__))
+    # Construct the image path relative to the test directory
+    image_path = os.path.join(test_dir, 'data', 'savanna.jpg')
+
+    # Read image data
+    with open(image_path, 'rb') as image_file:
+        image_data = image_file.read()
+
+    # Encode the image data in base64
+    encoded_image = base64.b64encode(image_data).decode('utf-8')
+
+    # Prepare the payload
+    json_event = {
+        'body': encoded_image
+    }
+
+    return json_event
+
+
+@pytest.fixture
+def context():
+    return None
+
+
+def test_lambda_handler(event, context):
+    lambda_response = lambda_handler(event, context)
+    response_data = json.loads(lambda_response["body"])
+
+    print("lambda_response - type",type(lambda_response))
+    print("lambda_response", lambda_response)
+    print("response_data - type", type(response_data))
+    print("response_data", response_data)
+
+    response_keys = list(response_data.keys())
+    gt_keys = ['scores', 'labels', 'boxes']
+
+    assert lambda_response["statusCode"] == 200
+    assert set(response_keys) == set(gt_keys), "Response keys do not match ground truth"
+    assert len(response_data['scores']) == 5
+    assert len(response_data['labels']) == 5

tests/test_ml_detection.py

@@ -0,0 +1,35 @@
+from detection import ml_detection
+import os
+import pytest
+
+
+# Test model loading
+@pytest.mark.parametrize("test_model_uri",[
+    ("facebook/detr-resnet-50"),
+    ("facebook/detr-resnet-101"),
+])
+def test_load_model(test_model_uri):
+    processor, model = ml_detection.load_model(test_model_uri)
+    assert processor is not None
+    assert model is not None
+
+
+# Test image detection
+@pytest.mark.parametrize("test_model_uri", [
+    ("facebook/detr-resnet-50"),
+    ("facebook/detr-resnet-101"),
+])
+def test_object_detection(test_model_uri):
+    processor, model = ml_detection.load_model(test_model_uri)
+
+    # Get the directory of the current test file
+    test_dir = os.path.dirname(os.path.abspath(__file__))
+    # Construct the image path relative to the test directory
+    image_path = os.path.join(test_dir, 'data', 'savanna.jpg')
+
+    with open(image_path, "rb") as f:
+        image_bytes = f.read()
+
+    results = ml_detection.object_detection(processor, model, image_bytes)
+    assert results is not None
+    assert isinstance(results, dict)

tests/test_ml_utils.py

@@ -0,0 +1,11 @@
+from detection import ml_utils
+import torch
+import json
+
+
+# Test dictionary conversion
+def test_convert_tensor_dict_to_json():
+    my_dict = {'scores': torch.tensor([1, 2, 3])}
+    my_list_gt = {'scores': [1, 2, 3]}
+    my_list = ml_utils.convert_tensor_dict_to_json(my_dict)
+    assert my_list == my_list_gt