Spaces:
Running on Zero
Running on Zero
File size: 5,731 Bytes
c31d06f adca048 c31d06f adca048 f6465fc c31d06f adca048 c31d06f adca048 bf8f4a1 c31d06f d773772 c31d06f 3ad81a7 | 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 | import io
import pytest
from fastapi.testclient import TestClient
from PIL import Image
import app as studio
@pytest.fixture
def client(tmp_path, monkeypatch):
monkeypatch.setattr(studio, "RUNTIME", tmp_path)
monkeypatch.setattr(studio, "SESSIONS", {})
monkeypatch.setattr(studio, "JOBS", {})
return TestClient(studio.app)
def new_session(client, example="english"):
response = client.post("/api/session", json={"example": example})
assert response.status_code == 200
return response.json()["id"]
@pytest.mark.parametrize(("example", "count"), [("english", 4), ("chinese", 13)])
def test_examples_load_complete_ready_to_run_presets(client, example, count):
session = client.post("/api/session", json={"example": example}).json()
sid = session["id"]
targets = session["regions"]
settings = session["settings"]
image = Image.open(io.BytesIO(client.get(f"/api/sessions/{sid}/source.png").content))
assert len(targets) == count
assert settings == {"variant": "meanflow", "steps": 4, "seed": 42,
"guidance": 2.6, "text_threshold": 0.5}
for row in targets:
x0, y0, x1, y1 = row["bbox"]
assert 0 <= x0 < x1 <= image.width and 0 <= y0 < y1 <= image.height
assert row["origin"] == "example" and row["enabled"] and row["text"].strip()
# The first response is already a valid inference request; no OCR/setup request is needed.
request = studio.RefineRequest.model_validate({**settings, "regions": targets})
assert studio.validate_refine(sid, request)["size"] == image.size
restored = client.get(f"/api/sessions/{sid}/targets").json()
assert restored == {"regions": targets, "settings": settings}
def test_upload_rejects_invalid_data_and_normalizes_valid_images(client):
assert client.post("/api/upload", files={"file": ("test.png", b"not a png")}).status_code == 422
buffer = io.BytesIO()
Image.new("RGB", (128, 80), "red").save(buffer, "JPEG")
response = client.post("/api/upload", files={"file": ("test.jpg", buffer.getvalue())})
assert response.status_code == 200
session = response.json()
assert (session["width"], session["height"]) == (128, 80)
assert not session.get("regions") and "settings" not in session
assert client.get(session["image"]).headers["content-type"] == "image/webp"
assert client.get(f"/api/sessions/{session['id']}/source.png").headers["content-type"] == "image/png"
assert client.get(f"/api/sessions/{session['id']}/targets").status_code == 404
def test_invalid_regions_never_enter_the_model_queue(client, monkeypatch):
sid = new_session(client)
calls = []
monkeypatch.setattr(studio.POOL, "submit", lambda *args: calls.append(args))
for box, text in [([-2, 1, 10, 10], "word"), ([0, 0, 99999, 20], "word"), ([5, 1, 3, 4], "word"), ([1, 1, 5, 5], " ")]:
response = client.post(f"/api/sessions/{sid}/refine", json={"regions": [{"id": "x", "bbox": box, "text": text}]})
assert response.status_code == 422
assert not calls
def test_duplicate_clicks_share_one_job_and_sessions_cannot_read_each_other(client, monkeypatch):
first, second = new_session(client), new_session(client, "chinese")
calls = []
monkeypatch.setattr(studio.POOL, "submit", lambda *args: calls.append(args))
path = f"/api/sessions/{first}/detect"
a, b = client.post(path, json={}).json(), client.post(path, json={}).json()
assert a["job_id"] == b["job_id"]
assert len(calls) == 1
assert client.get(f"/api/sessions/{second}/jobs/{a['job_id']}").status_code == 404
assert client.get(f"/api/sessions/{first}/jobs/{a['job_id']}").status_code == 200
assert client.get(f"/api/sessions/{first}/job").json()["id"] == a["job_id"]
assert client.get(f"/api/sessions/{second}/job").status_code == 404
# A second session cannot retrieve a result by guessing its name.
own = studio.SESSIONS[first]["directory"] / "unique.png"
Image.new("RGB", (16, 16)).save(own)
assert client.get(f"/api/sessions/{second}/unique.png").status_code == 404
assert client.get(f"/api/sessions/{first}/unique.png").status_code == 200
def test_sessions_and_jobs_expire_together(client, monkeypatch):
sid = new_session(client)
directory = studio.SESSIONS[sid]["directory"]
studio.SESSIONS[sid]["touched"] -= studio.SESSION_TTL + 1
new_session(client)
assert not directory.exists()
assert client.get(f"/api/sessions/{sid}/source.png").status_code == 404
@pytest.mark.parametrize('example', ['english', 'chinese'])
def test_example_selection_reuses_preview_without_reencoding(client, monkeypatch, example):
def unexpected_encoding(*args, **kwargs):
raise AssertionError('Example selection is re-encoding images on the request path')
with monkeypatch.context() as patch:
patch.setattr(Image.Image, 'save', unexpected_encoding)
first = client.post('/api/session', json={'example': example}).json()
second = client.post('/api/session', json={'example': example}).json()
assert first['id'] != second['id']
assert first['image'] == second['image']
preview = client.get(first['image'])
assert preview.status_code == 200
assert 'public' in preview.headers['cache-control']
assert 'immutable' in preview.headers['cache-control']
assert client.get(first['image'], headers={'If-None-Match': preview.headers['etag']}).status_code == 304
assert first['image'] in client.get('/').text
assert Image.open(io.BytesIO(preview.content)).size == (first['width'], first['height'])
original = client.get(f"/api/sessions/{first['id']}/source.png")
assert original.headers['cache-control'] == 'private, no-store'
|