tests/test_api.py

from contextlib import contextmanager

from fastapi.testclient import TestClient

from token_trail_backend.api import create_app
from token_trail_backend.config import MAX_PROMPT_CHARS, Settings


class FakeEngine:
    ready = True

    def close(self):
        pass

    def trace(self, request, cancel=None):
        return {"prompt": request.prompt, "model": request.model, "steps": []}


@contextmanager
def client(engine=None):
    with TestClient(
        create_app(Settings(rate_limit_per_minute=100), engine or FakeEngine(), load_model=False)
    ) as test_client:
        yield test_client


def test_valid_trace():
    with client() as api:
        response = api.post("/v1/traces", json={"prompt": "Hello"})
        assert response.status_code == 200
        assert response.json()["prompt"] == "Hello"


def test_request_bounds():
    with client() as api:
        for payload in (
            {"prompt": "x" * (MAX_PROMPT_CHARS + 1)},
            {"prompt": "x", "max_new_tokens": 33},
            {"prompt": "x", "top_k": 11},
            {"prompt": "x", "model": "other"},
            {"prompt": "x", "unexpected": 1},
            {"prompt": "x", "max_new_tokens": "2"},
        ):
            assert api.post("/v1/traces", json=payload).status_code in (413, 422)


def test_payload_limit_and_malformed_json():
    with client() as api:
        assert api.post("/v1/traces", content=b"x" * 4097).status_code == 413
        assert (
            api.post(
                "/v1/traces",
                content=b'{"prompt":"\\ud800"}',
                headers={"Content-Type": "application/json"},
            ).status_code
            == 422
        )
        assert (
            api.post(
                "/v1/traces", content=b"{", headers={"Content-Type": "application/json"}
            ).status_code
            == 422
        )
        invalid_utf8 = api.post(
            "/v1/traces", content=b"\xff", headers={"Content-Type": "application/json"}
        )
        assert invalid_utf8.status_code == 400
        assert invalid_utf8.json()["error"]["code"] == "invalid_request"


def test_health_and_model_allowlist():
    with client() as api:
        assert api.get("/health").json() == {"status": "ready", "ready": True}
        models = api.get("/v1/models").json()["models"]
        assert len(models) == 1 and models[0]["alias"] == "qwen2.5-0.5b"
    engine = FakeEngine()
    engine.ready = False
    with client(engine) as api:
        assert api.get("/health").json()["ready"] is False
        assert api.post("/v1/traces", json={"prompt": "x"}).status_code == 503


def test_model_load_failure_does_not_fall_back():
    class BrokenEngine(FakeEngine):
        ready = False

        def load(self):
            raise RuntimeError("pinned model unavailable")

    try:
        with TestClient(create_app(engine=BrokenEngine())):
            pass
    except RuntimeError as exc:
        assert "pinned model unavailable" in str(exc)
    else:
        raise AssertionError("startup silently substituted a model")


def test_rate_limit():
    with TestClient(
        create_app(Settings(rate_limit_per_minute=1), FakeEngine(), load_model=False)
    ) as api:
        assert api.post("/v1/traces", json={"prompt": "a"}).status_code == 200
        assert api.post("/v1/traces", json={"prompt": "b"}).status_code == 429


def test_cors():
    with client() as api:
        allowed = api.options(
            "/v1/traces",
            headers={"Origin": "https://crunchycodes.net", "Access-Control-Request-Method": "POST"},
        )
        denied = api.options(
            "/v1/traces",
            headers={"Origin": "https://other.example", "Access-Control-Request-Method": "POST"},
        )
        assert allowed.headers.get("access-control-allow-origin") == "https://crunchycodes.net"
        assert "access-control-allow-origin" not in denied.headers


def test_trace_origin_is_enforced():
    with client() as api:
        allowed = api.post(
            "/v1/traces",
            json={"prompt": "x"},
            headers={"Origin": "https://crunchycodes.net"},
        )
        denied = api.post(
            "/v1/traces",
            json={"prompt": "x"},
            headers={"Origin": "https://unrelated.example"},
        )
        assert allowed.status_code == 200
        assert denied.status_code == 403
        assert denied.json()["error"]["code"] == "origin_not_allowed"