"""
This file contains tests for the API
"""
import os
import pickle
import pytest
import json
from unittest.mock import MagicMock, patch
from model import YOLO
def load_file(path):
# json
if path.endswith(".json"):
with open(path, "r") as f:
return json.load(f)
# pickle
if path.endswith(".pickle"):
with open(path, "rb") as f:
return pickle.load(f)
@pytest.fixture
def client():
from _wsgi import init_app
app = init_app(model_class=YOLO)
app.config["TESTING"] = True
with app.test_client() as client:
yield client
TEST_DIR = os.path.dirname(__file__)
label_configs = [
# test 1: wrong key in task data
"""
""",
# test 2: model_skip=true
"""
""",
]
tasks = [
# test 1: wrong key in task data
{"data": {"wrong_key": "https://some/path"}},
# test 2: model_skip
{"data": {"image": "https://some/path"}},
]
expected = [
# test 1: wrong key in task data
"Can't load path using key",
# test 2: model skip
"No suitable control tags",
]
@pytest.mark.parametrize(
"label_config, task, expect", zip(label_configs, tasks, expected)
)
def test_label_configs(client, label_config, task, expect, capsys):
data = {"schema": label_config, "project": "42"}
response = client.post(
"/setup", data=json.dumps(data), content_type="application/json"
)
assert response.status_code == 200, "Error while setup: " + str(response.content)
data = {"tasks": [task], "label_config": label_config}
response = client.post(
"/predict", data=json.dumps(data), content_type="application/json"
)
assert (
response.status_code == 500
), "Error was expected, but another status code found"
# Capture stdout and stderr
captured = capsys.readouterr()
# Check for specific words in the output
assert expect in captured.out, "Text not found in error string"
@pytest.fixture
def mock_logger(mocker):
return mocker.patch("model.logger")
@pytest.fixture
def mock_ml_backend():
return MagicMock()
@pytest.fixture
def mock_label_interface(mock_ml_backend):
mock_label_interface = MagicMock()
mock_ml_backend.label_interface = mock_label_interface
return mock_label_interface
@pytest.fixture
def mock_model_class():
mock_model_class = MagicMock()
mock_model_class.type = "RectangleLabels"
mock_model_class.is_control_matched = MagicMock(return_value=True)
mock_model_class.create = MagicMock()
return mock_model_class
@pytest.fixture
def available_model_classes(mock_model_class):
return [mock_model_class, MagicMock(), MagicMock()]
@pytest.fixture
def yolo_instance(mock_label_interface):
yolo = YOLO()
yolo.label_interface = mock_label_interface
return yolo
def test_control_no_to_name(yolo_instance, mock_label_interface, mock_logger):
mock_label_interface.controls = [
MagicMock(to_name=None, tag="RectangleLabels", name="test_label")
]
with pytest.raises(ValueError):
yolo_instance.detect_control_models()
mock_logger.warning.assert_called_once()
def test_control_no_label_map(
yolo_instance,
mock_label_interface,
mock_model_class,
available_model_classes,
mock_logger,
):
mock_instance = MagicMock()
mock_instance.label_map = None
mock_model_class.create.return_value = mock_instance
mock_label_interface.controls = [
MagicMock(to_name=["image"], tag="RectangleLabels", name="test_label")
]
with patch("model.available_model_classes", available_model_classes):
result = yolo_instance.detect_control_models()
# Updated expectation to handle control models that should be skipped if label_map is None
assert len(result) == 1
mock_logger.error.assert_called_once()
def test_control_with_valid_label_map(
yolo_instance,
mock_label_interface,
mock_model_class,
available_model_classes,
mock_logger,
):
mock_instance = MagicMock()
mock_instance.label_map = {"car": "Car"}
mock_model_class.create.return_value = mock_instance
mock_label_interface.controls = [
MagicMock(to_name=["image"], tag="RectangleLabels", name="test_label")
]
with patch("model.available_model_classes", available_model_classes):
result = yolo_instance.detect_control_models()
assert len(result) == 1
assert result[0] == mock_instance
mock_logger.debug.assert_called_once()