Spaces:
Sleeping
Sleeping
| """ | |
| tests/test_api_classify.py | |
| ────────────────────────── | |
| Tests for POST /classify endpoint (14 tests). | |
| Covers: | |
| - Cache hit / miss paths | |
| - top_k clamping and validation | |
| - Model selection validation (unknown name, yolo rejected) | |
| - 503 when models not ready | |
| - Response body correctness (confidence sum, sorting, inference_time) | |
| - Input validation (bad image, empty file) | |
| """ | |
| from __future__ import annotations | |
| from typing import TYPE_CHECKING | |
| from unittest.mock import MagicMock | |
| if TYPE_CHECKING: | |
| from fastapi.testclient import TestClient | |
| # ── Helpers ──────────────────────────────────────────────────────────────────── | |
| def _post_classify( | |
| client: "TestClient", | |
| image_bytes: bytes, | |
| model_name: str = "resnet50", | |
| top_k: int = 5, | |
| ) -> "TestClient": | |
| return client.post( | |
| "/classify", | |
| files={"file": ("test.jpg", image_bytes, "image/jpeg")}, | |
| data={"model_name": model_name, "top_k": top_k}, | |
| ) | |
| # ── Cache hit / miss ─────────────────────────────────────────────────────────── | |
| def test_classify_cache_miss_returns_cached_false( | |
| app_client: "TestClient", test_image_bytes: bytes | |
| ) -> None: | |
| """Cache miss path: Redis returns None → model runs → cached=False.""" | |
| app_client.app.state.redis.get.return_value = None # explicit miss | |
| response = _post_classify(app_client, test_image_bytes) | |
| assert response.status_code == 200 | |
| assert response.json()["cached"] is False | |
| def test_classify_cache_hit_returns_cached_true( | |
| app_client: "TestClient", | |
| test_image_bytes: bytes, | |
| mock_classify_response: dict, | |
| ) -> None: | |
| """Cache hit path: Redis returns stored dict → cached=True, no inference.""" | |
| app_client.app.state.redis.get.return_value = mock_classify_response | |
| response = _post_classify(app_client, test_image_bytes) | |
| assert response.status_code == 200 | |
| assert response.json()["cached"] is True | |
| # ── top_k behaviour ──────────────────────────────────────────────────────────── | |
| def test_classify_top_k_exactly_three( | |
| app_client: "TestClient", test_image_bytes: bytes | |
| ) -> None: | |
| """top_k=3 returns exactly 3 predictions.""" | |
| app_client.app.state.redis.get.return_value = None | |
| response = _post_classify(app_client, test_image_bytes, top_k=3) | |
| assert response.status_code == 200 | |
| assert len(response.json()["predictions"]) == 3 | |
| def test_classify_top_k_exceeds_num_classes_is_clipped( | |
| app_client: "TestClient", test_image_bytes: bytes | |
| ) -> None: | |
| """top_k=100 (> 22 classes) is silently clamped to NUM_CLASSES.""" | |
| app_client.app.state.redis.get.return_value = None | |
| response = _post_classify(app_client, test_image_bytes, top_k=100) | |
| assert response.status_code == 200 | |
| assert len(response.json()["predictions"]) <= 22 | |
| def test_classify_top_k_zero_is_rejected( | |
| app_client: "TestClient", test_image_bytes: bytes | |
| ) -> None: | |
| """top_k=0 is invalid; endpoint raises (500) or validates (422) — never 200.""" | |
| response = _post_classify(app_client, test_image_bytes, top_k=0) | |
| # FastAPI does not validate Form int ranges, so top_k=0 reaches route logic. | |
| # probs.topk(0) succeeds but predictions[0] then raises IndexError → 500. | |
| assert response.status_code in (422, 500) | |
| # ── Model name validation ────────────────────────────────────────────────────── | |
| def test_classify_unknown_model_rejected( | |
| app_client: "TestClient", test_image_bytes: bytes | |
| ) -> None: | |
| """Non-existent model name → 400.""" | |
| response = _post_classify(app_client, test_image_bytes, model_name="nonexistent") | |
| assert response.status_code == 400 | |
| def test_classify_yolo_model_rejected( | |
| app_client: "TestClient", test_image_bytes: bytes, mock_model: "MagicMock" | |
| ) -> None: | |
| """YOLO is a detection model; /classify must reject it with 400. | |
| Route logic: `if model_name not in models or model_name == "yolo": raise 400` | |
| Detail message lists the allowed models, not "yolo" itself. | |
| """ | |
| app_client.app.state.models["yolo"] = mock_model | |
| try: | |
| response = _post_classify(app_client, test_image_bytes, model_name="yolo") | |
| assert response.status_code == 400 | |
| assert "model_name" in response.json()["detail"].lower() | |
| finally: | |
| app_client.app.state.models.pop("yolo", None) | |
| # ── 503 when not ready ───────────────────────────────────────────────────────── | |
| def test_classify_503_when_models_not_ready( | |
| app_client: "TestClient", test_image_bytes: bytes | |
| ) -> None: | |
| app_client.app.state.models_ready = False | |
| try: | |
| response = _post_classify(app_client, test_image_bytes) | |
| assert response.status_code == 503 | |
| finally: | |
| app_client.app.state.models_ready = True | |
| # ── Response body correctness ────────────────────────────────────────────────── | |
| def test_classify_confidence_scores_sum_to_one( | |
| app_client: "TestClient", test_image_bytes: bytes | |
| ) -> None: | |
| """Softmax output sums to 1.0 across all classes (within floating-point tolerance).""" | |
| app_client.app.state.redis.get.return_value = None | |
| response = _post_classify(app_client, test_image_bytes, top_k=22) | |
| assert response.status_code == 200 | |
| total = sum(p["confidence"] for p in response.json()["predictions"]) | |
| assert 0.99 <= total <= 1.01 | |
| def test_classify_inference_time_is_positive( | |
| app_client: "TestClient", test_image_bytes: bytes | |
| ) -> None: | |
| app_client.app.state.redis.get.return_value = None | |
| response = _post_classify(app_client, test_image_bytes) | |
| assert response.status_code == 200 | |
| assert response.json()["inference_time_ms"] > 0 | |
| def test_classify_predictions_sorted_descending( | |
| app_client: "TestClient", test_image_bytes: bytes | |
| ) -> None: | |
| """Predictions are ordered by confidence (highest first).""" | |
| app_client.app.state.redis.get.return_value = None | |
| response = _post_classify(app_client, test_image_bytes, top_k=5) | |
| assert response.status_code == 200 | |
| confs = [p["confidence"] for p in response.json()["predictions"]] | |
| assert confs == sorted(confs, reverse=True) | |
| def test_classify_model_name_echoed_in_response( | |
| app_client: "TestClient", test_image_bytes: bytes | |
| ) -> None: | |
| app_client.app.state.redis.get.return_value = None | |
| response = _post_classify(app_client, test_image_bytes, model_name="resnet50") | |
| assert response.status_code == 200 | |
| assert response.json()["model_name"] == "resnet50" | |
| # ── Input validation ─────────────────────────────────────────────────────────── | |
| def test_classify_non_image_bytes_rejected(app_client: "TestClient") -> None: | |
| """Garbage bytes that PIL cannot open → 400.""" | |
| response = app_client.post( | |
| "/classify", | |
| files={"file": ("bad.jpg", b"not an image at all", "image/jpeg")}, | |
| data={"model_name": "resnet50", "top_k": 5}, | |
| ) | |
| assert response.status_code in (400, 422, 500) | |
| def test_classify_empty_file_rejected(app_client: "TestClient") -> None: | |
| """Zero-byte file upload → error (PIL raises or FastAPI rejects).""" | |
| response = app_client.post( | |
| "/classify", | |
| files={"file": ("empty.jpg", b"", "image/jpeg")}, | |
| data={"model_name": "resnet50", "top_k": 5}, | |
| ) | |
| assert response.status_code in (400, 422, 500) | |