Spaces:
Running on Zero
Running on Zero
| import io | |
| import pytest | |
| from fastapi.testclient import TestClient | |
| from PIL import Image | |
| import app as studio | |
| 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"] | |
| 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 | |
| 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' | |