Spaces:
Sleeping
Sleeping
File size: 4,112 Bytes
00340de |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 |
import io
from unittest.mock import MagicMock
import pytest
import torch
from PIL import Image
import app as app_module
def test_health(client):
response = client.get("/health")
assert response.status_code == 200
assert response.json() == {"status": "ok"}
def test_detect_returns_detections(client):
img = Image.new("RGB", (640, 480), color="red")
buffer = io.BytesIO()
img.save(buffer, format="PNG")
buffer.seek(0)
response = client.post("/detect", files={"file": ("test.png", buffer, "image/png")})
assert response.status_code == 200
data = response.json()
assert "num_detections" in data
assert "detections" in data
assert isinstance(data["detections"], list)
def test_detect_with_jpeg(client):
img = Image.new("RGB", (640, 480), color="blue")
buffer = io.BytesIO()
img.save(buffer, format="JPEG")
buffer.seek(0)
response = client.post(
"/detect", files={"file": ("test.jpg", buffer, "image/jpeg")}
)
assert response.status_code == 200
data = response.json()
assert "num_detections" in data
assert "detections" in data
def test_detect_missing_file(client):
response = client.post("/detect")
assert response.status_code == 422
def test_detect_response_structure(client):
img = Image.new("RGB", (100, 100), color="green")
buffer = io.BytesIO()
img.save(buffer, format="PNG")
buffer.seek(0)
response = client.post("/detect", files={"file": ("test.png", buffer, "image/png")})
assert response.status_code == 200
data = response.json()
assert isinstance(data["num_detections"], int)
assert data["num_detections"] >= 0
assert data["num_detections"] == len(data["detections"])
for detection in data["detections"]:
assert "class_id" in detection
assert "class_name" in detection
assert "confidence" in detection
assert "bbox_xyxy" in detection
assert isinstance(detection["class_id"], int)
assert isinstance(detection["class_name"], str)
assert 0 <= detection["confidence"] <= 1
assert len(detection["bbox_xyxy"]) == 4
def test_detect_model_not_loaded(client):
"""Test that detect raises error when model is not loaded."""
original_model = app_module.model
app_module.model = None
try:
img = Image.new("RGB", (100, 100), color="red")
buffer = io.BytesIO()
img.save(buffer, format="PNG")
buffer.seek(0)
with pytest.raises(RuntimeError, match="Model not loaded"):
client.post("/detect", files={"file": ("test.png", buffer, "image/png")})
finally:
app_module.model = original_model
def test_detect_with_detections(client):
"""Test detection with mocked YOLO results containing detections."""
original_model = app_module.model
# Create mock model
mock_model = MagicMock()
mock_model.names = {0: "person", 1: "car"}
# Create mock box
mock_box = MagicMock()
mock_box.cls = torch.tensor([0])
mock_box.conf = torch.tensor([0.95])
mock_box.xyxy = torch.tensor([[100.0, 150.0, 300.0, 400.0]])
# Create mock result
mock_result = MagicMock()
mock_result.boxes = [mock_box]
mock_model.return_value = [mock_result]
app_module.model = mock_model
try:
img = Image.new("RGB", (640, 480), color="red")
buffer = io.BytesIO()
img.save(buffer, format="PNG")
buffer.seek(0)
response = client.post(
"/detect", files={"file": ("test.png", buffer, "image/png")}
)
assert response.status_code == 200
data = response.json()
assert data["num_detections"] == 1
assert len(data["detections"]) == 1
assert data["detections"][0]["class_id"] == 0
assert data["detections"][0]["class_name"] == "person"
assert data["detections"][0]["confidence"] == pytest.approx(0.95)
assert data["detections"][0]["bbox_xyxy"] == pytest.approx(
[100.0, 150.0, 300.0, 400.0]
)
finally:
app_module.model = original_model
|