diff --git a/tt-metal-yolov4/requirements-test.txt b/tt-metal-yolov4/requirements-test.txt index 01b78ec..4db54de 100644 --- a/tt-metal-yolov4/requirements-test.txt +++ b/tt-metal-yolov4/requirements-test.txt @@ -1,2 +1,3 @@ pillow==10.3.0 locust==2.25.0 +pytest==7.2.2 diff --git a/tt-metal-yolov4/tests/locustfile.py b/tt-metal-yolov4/tests/locustfile.py index 6ff8e7c..2b05872 100644 --- a/tt-metal-yolov4/tests/locustfile.py +++ b/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) diff --git a/tt-metal-yolov4/tests/test_inference_api.py b/tt-metal-yolov4/tests/test_inference_api.py new file mode 100644 index 0000000..d4f630a --- /dev/null +++ b/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 diff --git a/tt-metal-yolov4/tests/utils.py b/tt-metal-yolov4/tests/utils.py new file mode 100644 index 0000000..e884c07 --- /dev/null +++ b/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