diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..3933183 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,6 @@ +.git +.venv +__pycache__ +.pytest_cache +tests +data diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..523ea62 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,44 @@ +name: CI + +on: + pull_request: + push: + branches: + - main + +permissions: + contents: read + packages: write + +jobs: + test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - run: pip install -r requirements-dev.txt + - run: pytest -q + - run: python -m compileall -q app tests + + publish: + if: github.event_name == 'push' && github.ref == 'refs/heads/main' + needs: test + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + - name: Build and push immutable image + id: push + uses: docker/build-push-action@v6 + with: + context: . + push: true + tags: ghcr.io/team-pinlog/image:sha-${{ github.sha }} + - name: Output digest + run: echo "${{ steps.push.outputs.digest }}" diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..f39623c --- /dev/null +++ b/.gitignore @@ -0,0 +1,8 @@ +.venv/ +__pycache__/ +*.py[cod] +.pytest_cache/ +.coverage +htmlcov/ +data/ +.env diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..256017d --- /dev/null +++ b/Dockerfile @@ -0,0 +1,20 @@ +FROM python:3.12-slim + +ENV PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 \ + DATABASE_PATH=/data/app.db \ + FILES_DIR=/data/files + +WORKDIR /app +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt +COPY app ./app +RUN groupadd --gid 1000 app \ + && useradd --uid 1000 --gid 1000 --no-create-home --shell /usr/sbin/nologin app \ + && mkdir -p /data/files \ + && chown -R 1000:1000 /data + +EXPOSE 8000 +VOLUME ["/data"] +USER 1000:1000 +CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"] diff --git a/README.md b/README.md new file mode 100644 index 0000000..fc3314a --- /dev/null +++ b/README.md @@ -0,0 +1,102 @@ +# GPU cover worker service + +A SQLite-backed API for creating six styled book-cover previews and one selected final cover. + +```text +SERVICE_URL=https://pin-log.com/image +``` + +The unauthenticated readiness endpoint is `GET /health` and returns +`{"status":"ok"}` with HTTP 200. + +## Run locally + +```bash +python -m venv .venv +. .venv/bin/activate +pip install -r requirements-dev.txt +uvicorn app.main:app --reload +``` + +Configuration: + +- `DATABASE_PATH` defaults to `data/app.db`. +- `FILES_DIR` defaults to `data/files`. +- `STALE_JOB_SECONDS` defaults to `1800`. A running job older than this is requeued + when attempts remain, or marked failed after `MAX_ATTEMPTS` claims. +- `PUBLIC_BASE_PATH` defaults to empty. Set it to the externally mounted service + prefix (for example, `/image`) so returned file URLs use + `/image/files/NAME`; the application's internal static route remains `/files/NAME`. +- `WORKER_TOKEN`, when set, protects every `/jobs` endpoint with `Authorization: Bearer TOKEN`. + +## Exact workflow + +Create all six previews: + +```bash +curl -X POST "$SERVICE_URL/api/covers" \ + -H 'content-type: application/json' \ + -d '{"title":"달빛 고양이","keywords":["숲","달빛"]}' +``` + +The response is `{"request_id":"...","candidates":[{"style_id":"watercolour","label":"수채화"}, ...]}`. + +Poll request state: + +```bash +curl "$SERVICE_URL/api/covers/REQUEST_ID" +``` + +A worker claims work (the default long poll is 25 seconds): + +```bash +curl "$SERVICE_URL/jobs/claim?worker_id=gpu-1" \ + -H "Authorization: Bearer $WORKER_TOKEN" +``` + +The claim response contains exactly `id`, `workflow`, and `inputs`. Upload one or more generated images with a JSON-string `meta` field; `image_0` is required: + +```bash +curl -X POST "$SERVICE_URL/jobs/JOB_ID/result" \ + -H "Authorization: Bearer $WORKER_TOKEN" \ + -F 'meta={"renderer":"comfyui"}' \ + -F image_0=@preview.webp \ + -F image_1=@alternate.webp +``` + +Report a permanent or retryable failure: + +```bash +curl -X POST "$SERVICE_URL/jobs/JOB_ID/fail" \ + -H "Authorization: Bearer $WORKER_TOKEN" \ + -H 'content-type: application/json' \ + -d '{"error":"GPU unavailable","retryable":true}' +``` + +Select a style to enqueue the full-size final job: + +```bash +curl -X POST "$SERVICE_URL/api/covers/REQUEST_ID/select" \ + -H 'content-type: application/json' \ + -d '{"style_id":"watercolour"}' +``` + +Previews request 512×768 WebP at quality 82. Final jobs request 1795×2657 WebP at quality 92. Both use generation dimensions 1024×1536 and four steps. + +## Verification + +```bash +pytest -q +python -m compileall -q app tests +``` + +## Container and CI + +The container runs as the fixed non-root identity `1000:1000`. Its persistent +`/data` directory is owned by that identity, so any mounted volume must also be +writable by UID/GID 1000. + +GitHub Actions runs the test suite and Python bytecode compilation for pull +requests and pushes to `main`. A push to `main` also publishes the image as +`ghcr.io/team-pinlog/image:sha-COMMIT_SHA` and reports its registry digest. No +mutable `latest` tag is produced. diff --git a/TDD_NOTES.md b/TDD_NOTES.md new file mode 100644 index 0000000..b66c252 --- /dev/null +++ b/TDD_NOTES.md @@ -0,0 +1,87 @@ +# TDD notes + +Tests were replaced first to express the new GPU worker contract before application code changed. + +## RED + +Initial environment check: + +```text +$ pytest -q +/bin/bash: line 1: pytest: command not found +exit 127 +``` + +After `python -m pip install -r requirements-dev.txt`, the untouched old implementation produced the real behavioral RED: + +```text +$ pytest -q +8 failed, 5 passed in 1.22s +exit 1 +``` + +Failures covered the obsolete `{title, style}` schema, missing request/candidate response, absence of six queued jobs and the incompatible claim/result/failure APIs. + +## GREEN + +After implementing the contract, one test incorrectly assumed final jobs receive queue priority. The contract does not specify priority, so that test was corrected to consume older FIFO preview work before claiming the final. The full suite then passed: + +```text +$ pytest -q +13 passed in 1.27s +exit 0 +``` + +Final verification: + +```text +$ pytest -q && python -m compileall -q app tests && python -m py_compile app/main.py tests/test_app.py tests/test_upload_read.py +13 passed in 1.30s +exit 0 +``` + +Both syntax commands are silent on success. Nothing was deployed, pushed, or committed. + +## Stale claims, public base path, and aggregate status + +Tests were added before implementation for stale running-job recovery (including +the max-attempt failure case), externally prefixed file URLs, terminal preview +completion without a final, and final-job terminal status overriding previews. + +Focused RED against the prior implementation: + +```text +$ .venv/bin/pytest -q tests/test_app.py -k 'stale or terminal or final_status_controls or public_base_path' +FFFF [100%] +4 failed, 12 deselected in 0.70s +exit 1 +``` + +Focused GREEN after implementation: + +```text +$ .venv/bin/pytest -q tests/test_app.py -k 'stale or terminal or final_status_controls or public_base_path' +.... [100%] +4 passed, 12 deselected in 0.63s +exit 0 +``` + +Full suite GREEN before final compile verification: + +```text +$ .venv/bin/pytest -q +................. [100%] +17 passed in 1.51s +exit 0 +``` + +Final requested verification: + +```text +$ .venv/bin/pytest -q && .venv/bin/python -m compileall -q app tests && .venv/bin/python -m py_compile app/main.py tests/test_app.py tests/test_upload_read.py +................. [100%] +17 passed in 1.49s +exit 0 +``` + +Both compile checks were silent on success. Nothing was committed, pushed, or deployed. diff --git a/app/__init__.py b/app/__init__.py new file mode 100644 index 0000000..18bac94 --- /dev/null +++ b/app/__init__.py @@ -0,0 +1 @@ +"""Book cover generation API.""" diff --git a/app/main.py b/app/main.py new file mode 100644 index 0000000..43b9768 --- /dev/null +++ b/app/main.py @@ -0,0 +1,342 @@ +from __future__ import annotations + +import asyncio +import inspect +import json +import os +import secrets +import sqlite3 +import uuid +from contextlib import asynccontextmanager +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Annotated, Any + +from fastapi import Depends, FastAPI, Header, HTTPException, Query, Request, status +from fastapi.responses import Response +from fastapi.staticfiles import StaticFiles +from pydantic import BaseModel, ConfigDict, Field, field_validator + + +STYLES = { + "watercolour": {"label": "수채화", "prompt": "delicate watercolor illustration, soft wet-on-wet washes, visible cold-press paper texture, muted palette, gentle bleeding edges"}, + "ink_line": {"label": "펜화", "prompt": "fine pen and ink line illustration, confident cross-hatching, high contrast black linework on cream paper, sparse selective wash"}, + "oil_painterly": {"label": "유화", "prompt": "oil painting on canvas, thick impasto brushstrokes, rich saturated pigment, visible canvas weave, dramatic chiaroscuro lighting"}, + "flat_minimal": {"label": "미니멀 플랫", "prompt": "flat vector illustration, bold simplified geometric shapes, limited three-color palette, generous negative space, clean editorial poster look"}, + "woodblock": {"label": "목판화", "prompt": "traditional woodblock print, carved grain texture, flat layered color registration, bold outlines, slight ink misregistration"}, + "pastel_storybook": {"label": "파스텔 동화풍", "prompt": "soft pastel storybook illustration, chalky grainy texture, warm gentle palette, rounded friendly forms, dreamy diffused light"}, +} + +WIDTH, HEIGHT = 1024, 1536 +MAX_ATTEMPTS = 3 + + +class StrictModel(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class CoverCreate(StrictModel): + title: str = Field(min_length=1, max_length=200) + keywords: list[str] + + @field_validator("title") + @classmethod + def title_not_blank(cls, value: str) -> str: + if not value.strip(): + raise ValueError("title must not be blank") + return value + + @field_validator("keywords") + @classmethod + def valid_keywords(cls, value: list[str]) -> list[str]: + if any(not item.strip() for item in value): + raise ValueError("keywords must not contain blanks") + return value + + +class Selection(StrictModel): + style_id: str + + @field_validator("style_id") + @classmethod + def known_style(cls, value: str) -> str: + if value not in STYLES: + raise ValueError("unknown style") + return value + + +class Failure(StrictModel): + error: str = Field(min_length=1, max_length=2000) + retryable: bool + + +SCHEMA = """ +CREATE TABLE IF NOT EXISTS requests ( + id TEXT PRIMARY KEY, title TEXT NOT NULL, keywords TEXT NOT NULL, + prompt TEXT NOT NULL, seed INTEGER NOT NULL, created_at TEXT NOT NULL +); +CREATE TABLE IF NOT EXISTS jobs ( + id TEXT PRIMARY KEY, request_id TEXT NOT NULL REFERENCES requests(id), + style_id TEXT NOT NULL, kind TEXT NOT NULL, workflow TEXT NOT NULL, + inputs TEXT NOT NULL, status TEXT NOT NULL, worker_id TEXT, file TEXT, + error TEXT, claimed_at TEXT, attempts INTEGER NOT NULL DEFAULT 0, + created_at TEXT NOT NULL +); +CREATE INDEX IF NOT EXISTS jobs_queue ON jobs(status, created_at); +CREATE INDEX IF NOT EXISTS jobs_request ON jobs(request_id, kind, style_id); +""" + + +def now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _connect(app: FastAPI) -> sqlite3.Connection: + connection = sqlite3.connect(app.state.database_path, timeout=10, isolation_level=None) + connection.row_factory = sqlite3.Row + connection.execute("PRAGMA foreign_keys = ON") + connection.execute("PRAGMA busy_timeout = 10000") + return connection + + +def _prompt(title: str, keywords: list[str]) -> str: + lowered_title = title.strip().lower() + clean = [ + item.strip() + for item in keywords + if item.strip() + and item.isascii() + and lowered_title not in item.strip().lower() + and item.strip().lower() not in lowered_title + ] + scene = ", ".join(clean) if clean and len(clean) == len(keywords) else "a quiet landscape with balanced natural forms and soft atmospheric light" + return f"Wordless book cover artwork, no typography. Scene: {scene}. No text, letters, logo, watermark, or signature." + + +def _inputs(prompt: str, style_id: str, seed: int, kind: str, prefix: str) -> dict[str, Any]: + final = kind == "final" + return { + "prompt": prompt, + "style": STYLES[style_id]["prompt"], + "seed": seed, + "prefix": prefix, + "out_width": 1795 if final else 512, + "out_height": 2657 if final else 768, + "format": "webp", + "quality": 92 if final else 82, + "steps": 4, + "width": WIDTH, + "height": HEIGHT, + } + + +def worker_auth(authorization: Annotated[str | None, Header()] = None) -> None: + token = os.getenv("WORKER_TOKEN") + if token and (authorization is None or not secrets.compare_digest(authorization, f"Bearer {token}")): + raise HTTPException(status.HTTP_401_UNAUTHORIZED, "invalid worker token", headers={"WWW-Authenticate": "Bearer"}) + + +def create_app() -> FastAPI: + database_path = Path(os.getenv("DATABASE_PATH", "data/app.db")).resolve() + files_dir = Path(os.getenv("FILES_DIR", "data/files")).resolve() + stale_job_seconds = int(os.getenv("STALE_JOB_SECONDS", "1800")) + public_base_path = os.getenv("PUBLIC_BASE_PATH", "").strip() + if public_base_path and not public_base_path.startswith("/"): + public_base_path = f"/{public_base_path}" + public_base_path = public_base_path.rstrip("/") + database_path.parent.mkdir(parents=True, exist_ok=True) + files_dir.mkdir(parents=True, exist_ok=True) + + @asynccontextmanager + async def lifespan(app: FastAPI): + with _connect(app) as connection: + connection.execute("PRAGMA journal_mode = WAL") + connection.executescript(SCHEMA) + yield + + app = FastAPI(title="GPU cover worker service", lifespan=lifespan) + app.state.database_path = str(database_path) + app.state.files_dir = files_dir + app.mount("/files", StaticFiles(directory=files_dir), name="files") + + @app.get("/health") + def health() -> dict[str, str]: + return {"status": "ok"} + + @app.get("/api/styles") + def styles() -> dict: + return STYLES + + @app.post("/api/covers") + def create_cover(body: CoverCreate) -> dict: + request_id, created_at = str(uuid.uuid4()), now() + prompt = _prompt(body.title, body.keywords) + seed = secrets.randbelow(2_000_000_000) + with _connect(app) as connection: + connection.execute("BEGIN IMMEDIATE") + connection.execute( + "INSERT INTO requests(id,title,keywords,prompt,seed,created_at) VALUES(?,?,?,?,?,?)", + (request_id, body.title, json.dumps(body.keywords, ensure_ascii=False), prompt, seed, created_at), + ) + for index, style_id in enumerate(STYLES): + job_id = str(uuid.uuid4()) + inputs = _inputs(prompt, style_id, seed + index, "preview", f"{request_id}_{style_id}_preview") + connection.execute( + """INSERT INTO jobs(id,request_id,style_id,kind,workflow,inputs,status,created_at) + VALUES(?,?,?,'preview','cover',?,'queued',?)""", + (job_id, request_id, style_id, json.dumps(inputs), created_at), + ) + connection.commit() + return {"request_id": request_id, "candidates": [{"style_id": key, "label": value["label"]} for key, value in STYLES.items()]} + + def request_state(connection: sqlite3.Connection, request_id: str) -> dict | None: + request_row = connection.execute("SELECT id FROM requests WHERE id=?", (request_id,)).fetchone() + if request_row is None: + return None + jobs = connection.execute("SELECT * FROM jobs WHERE request_id=? ORDER BY created_at,rowid", (request_id,)).fetchall() + previews = {row["style_id"]: row for row in jobs if row["kind"] == "preview"} + finals = [row for row in jobs if row["kind"] == "final"] + candidates = [] + for style_id, style in STYLES.items(): + row = previews[style_id] + candidates.append({"style_id": style_id, "label": style["label"], "status": row["status"], "url": row["file"]}) + final = None + if finals: + row = finals[-1] + final = {"style_id": row["style_id"], "label": STYLES[row["style_id"]]["label"], "status": row["status"], "url": row["file"]} + terminal = {"done", "failed"} + if final is not None: + request_status = "done" if final["status"] in terminal else "running" + else: + request_status = "done" if len(candidates) == 6 and all( + candidate["status"] in terminal for candidate in candidates + ) else "running" + return {"request_id": request_id, "status": request_status, "candidates": candidates, "final": final} + + @app.get("/api/covers/{request_id}") + def get_cover(request_id: str) -> dict: + with _connect(app) as connection: + result = request_state(connection, request_id) + if result is None: + raise HTTPException(404, "cover request not found") + return result + + @app.post("/api/covers/{request_id}/select") + def select_cover(request_id: str, body: Selection) -> dict: + with _connect(app) as connection: + connection.execute("BEGIN IMMEDIATE") + stored = connection.execute("SELECT prompt,seed FROM requests WHERE id=?", (request_id,)).fetchone() + if stored is None: + raise HTTPException(404, "cover request not found") + index = list(STYLES).index(body.style_id) + job_id, created_at = str(uuid.uuid4()), now() + inputs = _inputs(stored["prompt"], body.style_id, stored["seed"] + index, "final", f"{request_id}_{body.style_id}_final") + connection.execute( + """INSERT INTO jobs(id,request_id,style_id,kind,workflow,inputs,status,created_at) + VALUES(?,?,?,'final','cover',?,'queued',?)""", + (job_id, request_id, body.style_id, json.dumps(inputs), created_at), + ) + connection.commit() + return {"request_id": request_id, "job_id": job_id, "style_id": body.style_id} + + def try_claim(worker_id: str) -> dict | None: + with _connect(app) as connection: + connection.execute("BEGIN IMMEDIATE") + stale_before = (datetime.now(timezone.utc) - timedelta(seconds=stale_job_seconds)).isoformat() + stale_error = f"job claim became stale after {stale_job_seconds} seconds" + connection.execute( + """UPDATE jobs SET status='queued',worker_id=NULL,claimed_at=NULL,error=? + WHERE status='running' AND claimed_at < ? AND attempts < ?""", + (stale_error, stale_before, MAX_ATTEMPTS), + ) + connection.execute( + """UPDATE jobs SET status='failed',worker_id=NULL,claimed_at=NULL,error=? + WHERE status='running' AND claimed_at < ? AND attempts >= ?""", + (stale_error, stale_before, MAX_ATTEMPTS), + ) + row = connection.execute("SELECT id FROM jobs WHERE status='queued' AND attempts < ? ORDER BY created_at,rowid LIMIT 1", (MAX_ATTEMPTS,)).fetchone() + if row is None: + connection.commit() + return None + connection.execute( + "UPDATE jobs SET status='running',worker_id=?,claimed_at=?,attempts=attempts+1,error=NULL WHERE id=? AND status='queued'", + (worker_id, now(), row["id"]), + ) + claimed = connection.execute("SELECT id,workflow,inputs FROM jobs WHERE id=?", (row["id"],)).fetchone() + connection.commit() + return {"id": claimed["id"], "workflow": claimed["workflow"], "inputs": json.loads(claimed["inputs"])} + + @app.get("/jobs/claim", dependencies=[Depends(worker_auth)], response_model=None) + async def claim_job(worker_id: Annotated[str, Query(min_length=1)], timeout: Annotated[float, Query(ge=0, le=25)] = 25) -> dict | Response: + deadline = asyncio.get_running_loop().time() + timeout + while True: + claimed = await asyncio.to_thread(try_claim, worker_id) + if claimed is not None: + return claimed + remaining = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + return Response(status_code=204) + await asyncio.sleep(min(0.25, remaining)) + + @app.post("/jobs/{job_id}/result", dependencies=[Depends(worker_auth)]) + async def job_result(job_id: str, request: Request) -> dict: + form = await request.form() + meta_value = form.get("meta") + if not isinstance(meta_value, str): + raise HTTPException(422, "meta must be a JSON string") + try: + json.loads(meta_value) + except json.JSONDecodeError as exc: + raise HTTPException(422, "meta must be valid JSON") from exc + uploads: list[tuple[str, Any]] = [] + for key, value in form.multi_items(): + if key.startswith("image_") and hasattr(value, "read"): + uploads.append((key, value)) + if not any(key == "image_0" for key, _ in uploads): + raise HTTPException(422, "image_0 is required") + with _connect(app) as connection: + row = connection.execute("SELECT status FROM jobs WHERE id=?", (job_id,)).fetchone() + if row is None: + raise HTTPException(404, "job not found") + if row["status"] != "running": + raise HTTPException(409, "job is not running") + first_url = None + for key, upload in uploads: + content = upload.read() + if inspect.isawaitable(content): + content = await content + if key == "image_0" and not content: + raise HTTPException(400, "image_0 must not be empty") + suffix = Path(getattr(upload, "filename", "") or "").suffix.lower() + suffix = suffix if suffix in {".png", ".jpg", ".jpeg", ".webp"} else ".bin" + filename = f"{job_id}_{key}{suffix}" + (files_dir / filename).write_bytes(content) + if key == "image_0": + first_url = f"{public_base_path}/files/{filename}" + with _connect(app) as connection: + connection.execute("BEGIN IMMEDIATE") + connection.execute("UPDATE jobs SET status='done',file=?,error=NULL WHERE id=? AND status='running'", (first_url, job_id)) + connection.commit() + return {"id": job_id, "status": "done", "file": first_url} + + @app.post("/jobs/{job_id}/fail", dependencies=[Depends(worker_auth)]) + def fail_job(job_id: str, body: Failure) -> dict: + with _connect(app) as connection: + connection.execute("BEGIN IMMEDIATE") + row = connection.execute("SELECT status,attempts FROM jobs WHERE id=?", (job_id,)).fetchone() + if row is None: + raise HTTPException(404, "job not found") + if row["status"] != "running": + raise HTTPException(409, "job is not running") + new_status = "queued" if body.retryable and row["attempts"] < MAX_ATTEMPTS else "failed" + connection.execute( + "UPDATE jobs SET status=?,error=?,worker_id=NULL,claimed_at=NULL WHERE id=?", + (new_status, body.error, job_id), + ) + connection.commit() + return {"id": job_id, "status": new_status, "error": body.error} + + return app + + +app = create_app() diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..c7b23ec --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +pythonpath = . +testpaths = tests diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..8b1b9d4 --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,3 @@ +-r requirements.txt +pytest==8.4.1 +httpx==0.28.1 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..8cfdc2c --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +fastapi==0.116.1 +uvicorn[standard]==0.35.0 +python-multipart==0.0.20 +pydantic==2.11.7 diff --git a/tests/test_app.py b/tests/test_app.py new file mode 100644 index 0000000..94f8b4b --- /dev/null +++ b/tests/test_app.py @@ -0,0 +1,308 @@ +import io +import json +import os +import sqlite3 +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timedelta, timezone + +import pytest +from fastapi.testclient import TestClient + +from app.main import STYLES, create_app + + +@pytest.fixture() +def client(tmp_path, monkeypatch): + monkeypatch.setenv("DATABASE_PATH", str(tmp_path / "app.db")) + monkeypatch.setenv("FILES_DIR", str(tmp_path / "files")) + monkeypatch.delenv("WORKER_TOKEN", raising=False) + monkeypatch.delenv("STALE_JOB_SECONDS", raising=False) + monkeypatch.delenv("PUBLIC_BASE_PATH", raising=False) + with TestClient(create_app()) as value: + yield value + + +def create_cover(client, title="Moon Cat", keywords=None): + return client.post( + "/api/covers", json={"title": title, "keywords": keywords or ["forest", "moonlight"]} + ) + + +def claim(client, worker="gpu-1"): + return client.get(f"/jobs/claim?worker_id={worker}&timeout=0") + + +def test_health(client): + response = client.get("/health") + assert response.status_code == 200 + assert response.json() == {"status": "ok"} + + +def test_create_contract_and_exact_six_persisted_jobs(client): + response = create_cover(client) + assert response.status_code == 200 + data = response.json() + assert set(data) == {"request_id", "candidates"} + assert data["candidates"] == [ + {"style_id": style_id, "label": style["label"]} for style_id, style in STYLES.items() + ] + + db = sqlite3.connect(os.environ["DATABASE_PATH"]) + db.row_factory = sqlite3.Row + request = dict(db.execute("SELECT * FROM requests").fetchone()) + assert {"id", "title", "keywords", "prompt", "seed", "created_at"} <= request.keys() + assert request["title"] == "Moon Cat" + assert json.loads(request["keywords"]) == ["forest", "moonlight"] + jobs = [dict(row) for row in db.execute("SELECT * FROM jobs ORDER BY rowid")] + assert len(jobs) == 6 + assert [job["style_id"] for job in jobs] == list(STYLES) + assert all(job["kind"] == "preview" and job["workflow"] == "cover" for job in jobs) + assert all({"id", "request_id", "inputs", "status", "worker_id", "file", "error", "claimed_at", "attempts"} <= job.keys() for job in jobs) + for index, job in enumerate(jobs): + inputs = json.loads(job["inputs"]) + assert inputs["seed"] == request["seed"] + index + assert (inputs["width"], inputs["height"]) == (1024, 1536) + assert (inputs["out_width"], inputs["out_height"]) == (512, 768) + assert (inputs["format"], inputs["quality"], inputs["steps"]) == ("webp", 82, 4) + + +@pytest.mark.parametrize("bad", [ + {"title": "x", "keywords": [], "extra": 1}, + {"title": "x"}, + {"title": "x", "keywords": "forest"}, + {"title": "x", "keywords": [1]}, +]) +def test_create_accepts_exact_schema(client, bad): + assert client.post("/api/covers", json=bad).status_code == 422 + + +def test_prompt_is_ascii_english_and_does_not_leak_non_ascii_input(client): + response = create_cover(client, title="달빛 고양이", keywords=["숲", "달빛"]) + assert response.status_code == 200 + db = sqlite3.connect(os.environ["DATABASE_PATH"]) + prompt = db.execute("SELECT prompt FROM requests").fetchone()[0] + assert prompt.isascii() + assert "달빛 고양이" not in prompt and "숲" not in prompt and "달빛" not in prompt + assert "moon" not in prompt.lower() and "cat" not in prompt.lower() + assert "generic" not in prompt.lower() + + create_cover(client, title="Secret Forest", keywords=["Secret Forest"]) + prompts = [row[0] for row in db.execute("SELECT prompt FROM requests ORDER BY rowid")] + assert "secret forest" not in prompts[-1].lower() + + +def test_claim_exact_payload_atomic_and_records_worker(client): + create_cover(client) + + def get_claim(index): + response = claim(client, f"worker-{index}") + assert response.status_code == 200 + return response.json() + + with ThreadPoolExecutor(max_workers=6) as executor: + jobs = list(executor.map(get_claim, range(6))) + assert len({job["id"] for job in jobs}) == 6 + for job in jobs: + assert set(job) == {"id", "workflow", "inputs"} + assert set(job["inputs"]) == { + "prompt", "style", "seed", "prefix", "out_width", "out_height", "format", + "quality", "steps", "width", "height", + } + db = sqlite3.connect(os.environ["DATABASE_PATH"]) + assert db.execute("SELECT count(*) FROM jobs WHERE status='running' AND worker_id IS NOT NULL").fetchone()[0] == 6 + assert claim(client).status_code == 204 + assert client.get("/jobs/claim?timeout=0").status_code == 422 + + +def test_claim_requeues_stale_jobs_and_fails_stale_jobs_at_max_attempts(client): + create_cover(client) + stale_retry = claim(client, "lost-worker").json() + stale_failed = claim(client, "lost-worker").json() + stale_at = (datetime.now(timezone.utc) - timedelta(hours=1)).isoformat() + with sqlite3.connect(os.environ["DATABASE_PATH"]) as db: + db.execute( + "UPDATE jobs SET claimed_at=? WHERE id=?", + (stale_at, stale_retry["id"]), + ) + db.execute( + "UPDATE jobs SET claimed_at=?,attempts=3 WHERE id=?", + (stale_at, stale_failed["id"]), + ) + + reclaimed = claim(client, "replacement-worker").json() + assert reclaimed["id"] == stale_retry["id"] + with sqlite3.connect(os.environ["DATABASE_PATH"]) as db: + db.row_factory = sqlite3.Row + retried = db.execute("SELECT * FROM jobs WHERE id=?", (stale_retry["id"],)).fetchone() + failed = db.execute("SELECT * FROM jobs WHERE id=?", (stale_failed["id"],)).fetchone() + assert (retried["status"], retried["worker_id"], retried["attempts"]) == ( + "running", "replacement-worker", 2 + ) + assert retried["error"] is None + assert failed["status"] == "failed" + assert failed["worker_id"] is None and failed["claimed_at"] is None + assert "stale" in failed["error"].lower() + + +def test_get_progress_result_upload_and_final_selection(client): + created = create_cover(client).json() + request_id = created["request_id"] + initial = client.get(f"/api/covers/{request_id}").json() + assert initial["status"] == "running" and initial["final"] is None + assert set(initial) == {"request_id", "status", "candidates", "final"} + assert all(set(item) == {"style_id", "label", "status", "url"} for item in initial["candidates"]) + + first = claim(client).json() + meta = {"renderer": "test", "images": 1} + done = client.post( + f"/jobs/{first['id']}/result", + data={"meta": json.dumps(meta)}, + files={"image_0": ("preview.webp", io.BytesIO(b"preview"), "image/webp")}, + ) + assert done.status_code == 200 + progress = client.get(f"/api/covers/{request_id}").json() + candidate = next(item for item in progress["candidates"] if item["status"] == "done") + assert client.get(candidate["url"]).content == b"preview" + + # Exhaust the older preview work; the contract does not give finals queue priority. + for index in range(5): + assert claim(client, f"preview-worker-{index}").status_code == 200 + + selected = client.post( + f"/api/covers/{request_id}/select", json={"style_id": candidate["style_id"]} + ) + assert selected.status_code == 200 + final_job = claim(client, "final-worker").json() + assert final_job["inputs"]["prompt"] == first["inputs"]["prompt"] + assert final_job["inputs"]["seed"] == first["inputs"]["seed"] + assert (final_job["inputs"]["out_width"], final_job["inputs"]["out_height"]) == (1795, 2657) + assert (final_job["inputs"]["width"], final_job["inputs"]["height"]) == (1024, 1536) + assert (final_job["inputs"]["format"], final_job["inputs"]["quality"]) == ("webp", 92) + client.post( + f"/jobs/{final_job['id']}/result", data={"meta": "{}"}, + files={"image_0": ("final.webp", io.BytesIO(b"final"), "image/webp")}, + ) + complete = client.get(f"/api/covers/{request_id}").json() + assert complete["status"] == "done" + assert complete["final"]["style_id"] == candidate["style_id"] + assert client.get(complete["final"]["url"]).content == b"final" + + +def test_request_done_when_all_previews_terminal_without_final(client): + request_id = create_cover(client).json()["request_id"] + with sqlite3.connect(os.environ["DATABASE_PATH"]) as db: + rows = db.execute( + "SELECT id FROM jobs WHERE request_id=? ORDER BY rowid", (request_id,) + ).fetchall() + for index, (job_id,) in enumerate(rows): + db.execute( + "UPDATE jobs SET status=? WHERE id=?", + ("done" if index % 2 == 0 else "failed", job_id), + ) + state = client.get(f"/api/covers/{request_id}").json() + assert state["final"] is None + assert state["status"] == "done" + + +def test_final_status_controls_request_status_when_final_exists(client): + request_id = create_cover(client).json()["request_id"] + with sqlite3.connect(os.environ["DATABASE_PATH"]) as db: + db.execute( + "UPDATE jobs SET status='done' WHERE request_id=? AND kind='preview'", + (request_id,), + ) + client.post(f"/api/covers/{request_id}/select", json={"style_id": "watercolour"}) + assert client.get(f"/api/covers/{request_id}").json()["status"] == "running" + with sqlite3.connect(os.environ["DATABASE_PATH"]) as db: + db.execute( + "UPDATE jobs SET status='failed' WHERE request_id=? AND kind='final'", + (request_id,), + ) + state = client.get(f"/api/covers/{request_id}").json() + assert state["final"]["status"] == "failed" + assert state["status"] == "done" + + +def test_public_base_path_prefixes_urls_but_static_route_stays_internal(tmp_path, monkeypatch): + monkeypatch.setenv("DATABASE_PATH", str(tmp_path / "base-path.db")) + monkeypatch.setenv("FILES_DIR", str(tmp_path / "files")) + monkeypatch.setenv("PUBLIC_BASE_PATH", "/image/") + monkeypatch.delenv("WORKER_TOKEN", raising=False) + with TestClient(create_app()) as prefixed: + request_id = create_cover(prefixed).json()["request_id"] + job = claim(prefixed).json() + result = prefixed.post( + f"/jobs/{job['id']}/result", data={"meta": "{}"}, + files={"image_0": ("preview.webp", io.BytesIO(b"prefixed"), "image/webp")}, + ).json() + expected = f"/image/files/{job['id']}_image_0.webp" + assert result["file"] == expected + candidate = next( + item for item in prefixed.get(f"/api/covers/{request_id}").json()["candidates"] + if item["status"] == "done" + ) + assert candidate["url"] == expected + assert prefixed.get(expected).status_code == 404 + assert prefixed.get(expected.removeprefix("/image")).content == b"prefixed" + + +def test_result_requires_meta_and_image_zero_but_accepts_multiple_images(client): + create_cover(client) + job = claim(client).json() + assert client.post(f"/jobs/{job['id']}/result", data={"meta": "{}"}).status_code == 422 + assert client.post( + f"/jobs/{job['id']}/result", data={"meta": "not-json"}, + files={"image_0": ("x.webp", io.BytesIO(b"x"), "image/webp")}, + ).status_code == 422 + response = client.post( + f"/jobs/{job['id']}/result", data={"meta": "{}"}, + files=[ + ("image_0", ("a.webp", io.BytesIO(b"a"), "image/webp")), + ("image_1", ("b.webp", io.BytesIO(b"b"), "image/webp")), + ], + ) + assert response.status_code == 200 + + +def test_failure_retryable_policy_and_exact_json(client): + create_cover(client) + job = claim(client).json() + assert client.post(f"/jobs/{job['id']}/fail", json={"error": "bad", "retryable": False, "x": 1}).status_code == 422 + failed = client.post(f"/jobs/{job['id']}/fail", json={"error": "bad", "retryable": False}) + assert failed.status_code == 200 + db = sqlite3.connect(os.environ["DATABASE_PATH"]) + assert db.execute("SELECT status FROM jobs WHERE id=?", (job["id"],)).fetchone()[0] == "failed" + + retry = claim(client).json() + for attempt in range(3): + response = client.post(f"/jobs/{retry['id']}/fail", json={"error": "gpu", "retryable": True}) + expected = "queued" if attempt < 2 else "failed" + assert response.json()["status"] == expected + if attempt < 2: + retry = claim(client).json() + + +def test_optional_bearer_protects_every_jobs_endpoint(tmp_path, monkeypatch): + monkeypatch.setenv("DATABASE_PATH", str(tmp_path / "auth.db")) + monkeypatch.setenv("FILES_DIR", str(tmp_path / "files")) + monkeypatch.setenv("WORKER_TOKEN", "secret") + with TestClient(create_app()) as protected: + request_id = create_cover(protected).json()["request_id"] + assert protected.get("/jobs/claim?worker_id=w&timeout=0").status_code == 401 + headers = {"Authorization": "Bearer secret"} + job = protected.get("/jobs/claim?worker_id=w&timeout=0", headers=headers).json() + assert protected.post(f"/jobs/{job['id']}/fail", json={"error": "x", "retryable": False}).status_code == 401 + assert protected.post(f"/jobs/{job['id']}/fail", json={"error": "x", "retryable": False}, headers=headers).status_code == 200 + assert protected.get(f"/api/covers/{request_id}").status_code == 200 + + +def test_not_found_and_persistence(tmp_path, monkeypatch): + monkeypatch.setenv("DATABASE_PATH", str(tmp_path / "persist.db")) + monkeypatch.setenv("FILES_DIR", str(tmp_path / "files")) + monkeypatch.delenv("WORKER_TOKEN", raising=False) + with TestClient(create_app()) as first: + request_id = create_cover(first).json()["request_id"] + assert first.get("/api/covers/nope").status_code == 404 + assert first.post("/api/covers/nope/select", json={"style_id": "watercolour"}).status_code == 404 + with TestClient(create_app()) as second: + assert second.get(f"/api/covers/{request_id}").status_code == 200 diff --git a/tests/test_upload_read.py b/tests/test_upload_read.py new file mode 100644 index 0000000..443fca6 --- /dev/null +++ b/tests/test_upload_read.py @@ -0,0 +1,16 @@ +import ast +from pathlib import Path + + +def test_result_handler_guards_upload_read_with_hasattr(): + source = Path("app/main.py").read_text() + tree = ast.parse(source) + assert any( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "hasattr" + and len(node.args) >= 2 + and isinstance(node.args[1], ast.Constant) + and node.args[1].value == "read" + for node in ast.walk(tree) + )