Team Ai
Apppublic

DevelopmentT/background-remover

sourceHugging Faceupdated 6mo agoView on Hugging Face
0likes
test_rest_api.py103 linesDownload Raw Back to tests
1"""2Integration tests for REST endpoints (routes/rest.py).3 4Uses FastAPI TestClient. OCR/rembg calls are mocked to keep tests fast.5"""6 7import io8import cv29import numpy as np10import pytest11from unittest.mock import patch, MagicMock12from fastapi.testclient import TestClient13 14# Patch easyocr before importing main15with patch.dict("sys.modules", {"easyocr": MagicMock()}):16    from main import app17 18from services.ocr_service import OCRResult19 20 21def _png_bytes(w=100, h=50) -> bytes:22    img = np.full((h, w, 3), 200, dtype=np.uint8)23    cv2.putText(img, "Test", (5, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 1)24    _, buf = cv2.imencode(".png", img)25    return buf.tobytes()26 27 28@pytest.fixture29def client():30    # Reset rate limiter between tests31    from routes.rest import _ip_tracker32    _ip_tracker.clear()33    return TestClient(app)34 35 36class TestHealth:37    def test_health_ok(self, client):38        r = client.get("/health")39        assert r.status_code == 20040        body = r.json()41        assert body["status"] == "ok"42        assert "timestamp" in body43 44 45class TestOCRRest:46    @patch("routes.rest.OCRService")47    def test_ocr_success(self, mock_cls, client):48        mock_inst = MagicMock()49        mock_inst.read_all.return_value = [50            OCRResult(text="Paracetamol", confidence=0.95, bbox=[[0, 0], [10, 0], [10, 10], [0, 10]]),51        ]52        mock_cls.get_instance.return_value = mock_inst53 54        r = client.post("/api/ocr", files={"file": ("test.png", io.BytesIO(_png_bytes()), "image/png")})55        assert r.status_code == 20056        body = r.json()57        assert body["status"] == "success"58        assert body["total_lines"] == 159        assert body["lines"][0]["text"] == "Paracetamol"60 61    def test_ocr_invalid_type(self, client):62        r = client.post("/api/ocr", files={"file": ("test.txt", io.BytesIO(b"hello"), "text/plain")})63        assert r.status_code == 40064 65    def test_ocr_empty_file(self, client):66        r = client.post("/api/ocr", files={"file": ("test.png", io.BytesIO(b""), "image/png")})67        assert r.status_code == 40068 69    def test_ocr_oversized(self, client):70        big = b"\x00" * (6 * 1024 * 1024)71        r = client.post("/api/ocr", files={"file": ("big.png", io.BytesIO(big), "image/png")})72        assert r.status_code == 40073 74    def test_ocr_no_file(self, client):75        r = client.post("/api/ocr")76        assert r.status_code == 40077 78 79class TestRateLimit:80    def test_rate_limit_exceeded(self, client):81        from routes.rest import _ip_tracker82        _ip_tracker.clear()83        # Fire 15 requests (limit), then 16th should fail84        for i in range(15):85            r = client.get("/health")86            assert r.status_code == 200, f"Request {i+1} failed unexpectedly"87        # 16th88        r = client.get("/health")89        # Health doesn't have rate limiter, so test on /api/ocr90        # Re-test with OCR endpoint91        _ip_tracker.clear()92        for _ in range(15):93            client.post("/api/ocr", files={"file": ("t.png", io.BytesIO(_png_bytes()), "image/png")})94        r = client.post("/api/ocr", files={"file": ("t.png", io.BytesIO(_png_bytes()), "image/png")})95        assert r.status_code == 42996 97 98class TestBackwardCompat:99    def test_legacy_remove_bg_exists(self, client):100        # Just verify the endpoint accepts the request (will fail at processing since rembg not loaded)101        r = client.post("/remove-bg", files={"file": ("t.txt", io.BytesIO(b"x"), "text/plain")})102        assert r.status_code == 400  # invalid type, not 404103