Add authentication header to locust testing and add API demonstration…

… test

tt-metal-yolov4/requirements-test.txt

@@ -1,2 +1,3 @@
 pillow==10.3.0
 locust==2.25.0
+pytest==7.2.2

tt-metal-yolov4/tests/locustfile.py

@@ -2,26 +2,16 @@
 #
 # SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
 
-import io
-import requests
-from PIL import Image
 from locust import HttpUser, task
+from utils import get_auth_header, sample_file
 
-# Save image as JPEG in-memory for load testing
-# Load sample image
-url = "http://images.cocodataset.org/val2017/000000039769.jpg"
-pil_image = Image.open(requests.get(url, stream=True).raw)
-pil_image = pil_image.resize((320, 320))  # Resize to target dimensions
-buf = io.BytesIO()
-pil_image.save(
-    buf,
-    format="JPEG",
-)
-byte_im = buf.getvalue()
-file = {"file": byte_im}
+
+# load sample file in memory
+file = sample_file()
 
 
 class HelloWorldUser(HttpUser):
     @task
     def hello_world(self):
-        self.client.post("/objdetection_v2", files=file)
+        headers = get_auth_header()
+        self.client.post("/objdetection_v2", files=file, headers=headers)

tt-metal-yolov4/tests/test_inference_api.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
+
+from http import HTTPStatus
+import os
+import pytest
+import requests
+from utils import get_auth_header, sample_file
+
+
+DEPLOY_URL = "http://127.0.0.1"
+SERVICE_PORT = int(os.getenv("SERVICE_PORT", 7000))
+API_BASE_URL = f"{DEPLOY_URL}:{SERVICE_PORT}"
+API_URL = f"{API_BASE_URL}/objdetection_v2"
+HEALTH_URL = f"{API_BASE_URL}/health"
+
+
+def test_valid_api_call():
+    # get sample image file
+    file = sample_file()
+    # make request with auth headers
+    headers = get_auth_header()
+    response = requests.post(API_URL, files=file, headers=headers)
+    # perform status and value checking
+    assert response.status_code == HTTPStatus.OK
+    assert isinstance(response.json(), list)
+
+
+@pytest.mark.skip(
+    reason="Not implemented, see https://github.com/tenstorrent/tt-inference-server/issues/63"
+)
+def test_get_health():
+    headers = {}
+    response = requests.get(HEALTH_URL, headers=headers, timeout=35)
+    assert response.status_code == 200

tt-metal-yolov4/tests/utils.py

@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
+
+import io
+import os
+from PIL import Image
+import requests
+
+
+def get_auth_header():
+    if authorization_header := os.getenv("AUTHORIZATION", None):
+        headers = {"Authorization": authorization_header}
+        return headers
+    else:
+        raise RuntimeError("AUTHORIZATION environment variable is undefined.")
+
+
+# save image as JPEG in-memory
+def sample_file():
+    # load sample image
+    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+    pil_image = Image.open(requests.get(url, stream=True).raw)
+    pil_image = pil_image.resize((320, 320))  # Resize to target dimensions
+    # convert to bytes
+    buf = io.BytesIO()
+    # format as JPEG
+    pil_image.save(
+        buf,
+        format="JPEG",
+    )
+    byte_im = buf.getvalue()
+    file = {"file": byte_im}
+    return file