tests/unit/test_api_contract.py

from __future__ import annotations

import base64
import io

from fastapi.testclient import TestClient
from PIL import Image

from app.main import app


def png_bytes(color=(12, 80, 160)) -> bytes:
    buffer = io.BytesIO()
    Image.new("RGB", (64, 48), color).save(buffer, format="PNG")
    return buffer.getvalue()


def data_url(data: bytes) -> str:
    return "data:image/png;base64," + base64.b64encode(data).decode()


def test_health_ready():
    client = TestClient(app)
    assert client.get("/health").json() == {"status": "ok"}
    assert client.get("/ready").json() == {"status": "ready"}


def test_upload_then_followup_text_only_references_image():
    client = TestClient(app)
    upload = client.post(
        "/v1/files",
        files={"file": ("design.png", png_bytes(), "image/png")},
        data={"purpose": "vision"},
    )
    assert upload.status_code == 200
    file_id = upload.json()["id"]

    first = client.post(
        "/v1/responses",
        json={
            "model": "glm-5.2-visual-runtime",
            "input": [
                {
                    "role": "user",
                    "content": [
                        {"type": "input_text", "text": "Remember the visual style of this image."},
                        {"type": "input_image", "file_id": file_id},
                    ],
                }
            ],
        },
    )
    assert first.status_code == 200
    response_id = first.json()["id"]

    second = client.post(
        "/v1/responses",
        json={
            "model": "glm-5.2-visual-runtime",
            "previous_response_id": response_id,
            "input": [{"role": "user", "content": [{"type": "input_text", "text": "Create slides using those colors."}]}],
        },
    )
    assert second.status_code == 200
    body = second.json()
    assert body["conversation_id"] == first.json()["conversation_id"]
    assert "presentation" in body["output_text"].lower()


def test_base64_image_and_chart_query_uses_chart_ir():
    client = TestClient(app)
    response = client.post(
        "/v1/responses",
        json={
            "model": "glm-5.2-visual-runtime",
            "input": [
                {
                    "role": "user",
                    "content": [
                        {"type": "input_text", "text": "What was the Q3 value in this chart?"},
                        {"type": "input_image", "image_url": data_url(png_bytes((230, 230, 240)))},
                    ],
                }
            ],
        },
    )
    assert response.status_code == 200
    assert "ChartIR" in response.json()["output_text"]
    assert "37" in response.json()["output_text"]


def test_rejects_localhost_image_url():
    client = TestClient(app)
    response = client.post(
        "/v1/responses",
        json={
            "model": "glm-5.2-visual-runtime",
            "input": [
                {
                    "role": "user",
                    "content": [
                        {"type": "input_text", "text": "Describe this image."},
                        {"type": "input_image", "image_url": "http://localhost/private.png"},
                    ],
                }
            ],
        },
    )
    assert response.status_code == 400
    assert "private" in response.json()["detail"].lower()