File size: 1,519 Bytes
17d2f7c | 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 | import pytest
from fastapi.testclient import TestClient
from app.main import app
from PIL import Image
import io
client = TestClient(app)
# --- Helper ---
def get_test_image() -> bytes:
"""Create a dummy test image"""
img = Image.new("RGB", (224, 224), color=(128, 64, 32))
buf = io.BytesIO()
img.save(buf, format="JPEG")
return buf.getvalue()
# --- Tests ---
def test_health_endpoint():
res = client.get("/health")
assert res.status_code == 200
assert res.json() == {"status": "ok"}
def test_predict_returns_valid_json():
img_bytes = get_test_image()
res = client.post(
"/predict",
files={"file": ("test.jpg", img_bytes, "image/jpeg")}
)
assert res.status_code == 200
data = res.json()
assert "label" in data
assert "score" in data
assert isinstance(data["score"], float)
assert 0.0 <= data["score"] <= 1.0
def test_predict_rejects_non_image():
res = client.post(
"/predict",
files={"file": ("test.txt", b"not an image", "text/plain")}
)
assert res.status_code == 415
def test_predict_rejects_corrupted_file():
res = client.post(
"/predict",
files={"file": ("bad.jpg", b"\xff\xd8corrupted", "image/jpeg")}
)
assert res.status_code == 400
def test_predict_rejects_oversized_file():
huge = b"A" * (11 * 1024 * 1024) # 11MB
res = client.post(
"/predict",
files={"file": ("big.jpg", huge, "image/jpeg")}
)
assert res.status_code == 413
|