pigway-frame-flow/tests/test_ai_images.py

import base64
import json
from io import BytesIO

from PIL import Image


def test_generate_requires_configuration(client):
    assert (
        client.post("/api/v1/ai/images", json={"prompt": "Landscape"}).status_code
        == 409
    )


def test_generation_queued_once_and_never_displays(client, service, monkeypatch):
    assert (
        client.put(
            "/api/v1/settings/ai-images",
            json={
                "endpoint": "https://example.com/v1/images/generations",
                "model": "test-model",
            },
        ).status_code
        == 200
    )
    client.put("/api/v1/settings/ai-images/key", json={"key": "private-ai-key"})
    image = BytesIO()
    Image.new("RGB", (100, 200), "blue").save(image, "PNG")
    calls = []

    def fetch(url, **kwargs):
        calls.append(kwargs)
        return json.dumps(
            {"data": [{"b64_json": base64.b64encode(image.getvalue()).decode()}]}
        ).encode()

    monkeypatch.setattr(service.assets, "fetch_url", fetch)
    headers = {"Idempotency-Key": "art-once"}
    first = client.post(
        "/api/v1/ai/images", json={"prompt": "Landscape"}, headers=headers
    )
    assert first.status_code == 202
    assert (
        client.post(
            "/api/v1/ai/images", json={"prompt": "Landscape"}, headers=headers
        ).json()
        == first.json()
    )
    service.process_one()
    operation = service.operation(first.json()["operation_id"])
    assert operation["status"] == "COMPLETED"
    assert len(calls) == 1
    assert calls[0]["timeout"] == 180
    assert json.loads(calls[0]["body"])["n"] == 1
    assert service.registry.get("living-room").client.uploads == 0
    assert "private-ai-key" not in client.get("/api/v1/settings/ai-images").text


def test_generation_failure_never_retries(client, service, monkeypatch):
    client.put(
        "/api/v1/settings/ai-images",
        json={"endpoint": "https://example.com/v1/images/generations", "model": "test"},
    )
    client.put("/api/v1/settings/ai-images/key", json={"key": "private-ai-key"})
    calls = []

    def fail(*args, **kwargs):
        calls.append(1)
        raise TimeoutError("private-ai-key")

    monkeypatch.setattr(service.assets, "fetch_url", fail)
    result = client.post("/api/v1/ai/images", json={"prompt": "Landscape"})
    service.process_one()
    service.process_one()
    op = service.operation(result.json()["operation_id"])
    assert op["status"] == "FAILED"
    assert len(calls) == 1
    assert "private-ai-key" not in json.dumps(op)


def test_interrupted_generation_is_never_requeued(app, service, settings):
    from pigway_frame_flow.api.app import create_app

    id = service.submit(
        "render-ai",
        {"prompt": "One image"},
        "living-room",
        "admin",
        "interrupted",
        "one",
    )
    service.db.execute(
        "UPDATE operations SET status='RENDERING',write_started=1 WHERE id=?", (id,)
    )
    restarted = create_app(settings, start_worker=False)
    assert restarted.state.service.operation(id)["status"] == "UNKNOWN"
    assert not restarted.state.service.blocked("living-room")
    assert restarted.state.service.process_one() is False
    restarted.state.service.db.close()


def test_extension_script_served(client):
    response = client.get("/ui/extensions.js")
    assert response.status_code == 200
    assert "layoutStudioPage" in response.text