File size: 2,740 Bytes
fd4e724 | 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 | import os
os.environ["OCR_CPU_AUTH_TOKEN"] = "test-token"
import pytest
from fastapi.testclient import TestClient
from ocr.config import Settings
from ocr.worker_api import create_app
class MockBackend:
class Health:
ready = True
async def health(self):
return self.Health()
async def capabilities(self):
return {"models": []}
async def recognize(self, p, opts):
return {}
@pytest.fixture(scope="module")
def client():
settings = Settings()
settings.worker_token = "test-token"
app = create_app(settings)
with TestClient(app) as c:
app.state.backend = MockBackend()
yield c
def test_health_public(client):
response = client.get("/health")
assert response.status_code == 200
assert response.json() == {"status": "ok"}
def test_ready_unauthorized(client):
response = client.get("/ready")
assert response.status_code == 401
def test_v1_models_authorized(client):
response = client.get("/v1/models", headers={"Authorization": "Bearer test-token"})
assert response.status_code == 200
assert response.json()["data"][0]["id"] == "ocr-cpu"
def test_v1_chat_completions_invalid_model(client):
req = {
"model": "gpt-4",
"messages": [
{
"role": "user",
"content": [
{
"type": "image_url",
"image_url": {"url": "data:image/png;base64,iVBORw0KGgo="},
}
],
}
],
}
response = client.post(
"/v1/chat/completions", headers={"Authorization": "Bearer test-token"}, json=req
)
assert response.status_code == 422
assert response.json()["code"] == "invalid_request"
def test_v1_chat_completions_invalid_url(client):
req = {
"model": "ocr-cpu",
"messages": [
{
"role": "user",
"content": [
{"type": "image_url", "image_url": {"url": "https://example.com/image.png"}}
],
}
],
}
response = client.post(
"/v1/chat/completions", headers={"Authorization": "Bearer test-token"}, json=req
)
assert response.status_code == 422
assert response.json()["code"] == "only_image_data_urls_supported"
def test_v1_ocr_page_wrong_sensitivity(client):
files = {"file": ("test.png", b"fake-image-data", "image/png")}
data = {"sensitivity": "D1"}
response = client.post(
"/v1/ocr/page", headers={"Authorization": "Bearer test-token"}, files=files, data=data
)
assert response.status_code == 422
assert response.json()["code"] == "invalid_worker_options"
|