agent_backend/examples/test_main.py

import json

import httpx
import pytest
from fastapi.testclient import TestClient

from main import app


@pytest.fixture
def client():
    app.state.products.clear()
    app.state.next_id = 1
    with TestClient(app) as test_client:
        yield test_client


def test_health_and_docs(client):
    assert client.get("/health").json() == {"status": "ok"}
    assert client.get("/docs").status_code == 200
    paths = client.get("/openapi.json").json()["paths"]
    assert "/products" in paths
    assert paths["/products"]["post"]["responses"]["422"]["content"]["application/json"]["schema"] == {
        "$ref": "#/components/schemas/ErrorResponse"
    }


def test_crud_lifecycle(client):
    created = client.post("/products", json={"name": "陶瓷杯"})
    assert created.status_code == 201
    assert created.json() == {"id": 1, "name": "陶瓷杯", "description": "", "is_active": False, "price_cents": 0}
    assert created.headers["location"] == "/products/1"
    assert client.get(created.headers["location"]).json() == created.json()
    assert client.get("/products/1").json() == created.json()
    replacement = {"name": "蓝色陶瓷杯", "description": "350ml", "is_active": True, "price_cents": 1990}
    updated = client.put("/products/1", json=replacement)
    assert updated.status_code == 200
    assert updated.json() == {"id": 1, **replacement}
    removed = client.delete("/products/1")
    assert removed.status_code == 204
    assert removed.content == b""
    assert client.get("/products").json() == []
    assert client.get("/products/1").status_code == 404
    assert client.post("/products", json={"name": "新商品"}).json()["id"] == 2


def test_list_filter_and_limit(client):
    client.post("/products", json={"name": "Ceramic mug", "description": "350ml"})
    client.post("/products", json={"name": "马克杯", "description": "white ceramic"})
    client.post("/products", json={"name": "帆布袋"})
    assert len(client.get("/products", params={"q": "CERAMIC"}).json()) == 2
    assert len(client.get("/products", params={"q": "ceramic", "limit": 1}).json()) == 1
    assert client.get("/products", params={"q": "unknown"}).json() == []


def test_default_limit_and_combined_is_active_filter(client):
    for index in range(12):
        client.post("/products", json={
            "name": f"product {index}",
            "description": "match" if index in (1, 2, 3) else "",
            "is_active": index % 2 == 0,
        })
    assert len(client.get("/products").json()) == 10
    is_active = client.get("/products", params={"is_active": "true"}).json()
    inactive = client.get("/products", params={"is_active": "false"}).json()
    assert len(is_active) == len(inactive) == 6
    assert all(product["is_active"] is True for product in is_active)
    assert all(product["is_active"] is False for product in inactive)
    combined = client.get("/products", params={"q": "match", "is_active": "false", "limit": 1}).json()
    assert [product["id"] for product in combined] == [2]


@pytest.mark.parametrize("name", ["", "   ", "x" * 121])
def test_invalid_name_reports_safe_validation_errors(client, name):
    response = client.post("/products", json={"name": name})
    assert response.status_code == 422
    error = response.json()["error"]
    assert error["code"] == "validation_error"
    assert error["details"][0]["loc"] == ["body", "name"]
    assert set(error["details"][0]) == {"loc", "msg", "type"}
    assert client.get("/products").json() == []


@pytest.mark.parametrize("missing", ["name", "description", "is_active", "price_cents"])
def test_put_requires_every_mutable_field(client, missing):
    client.post("/products", json={"name": "保留的商品"})
    body = {"name": "替换", "description": "内容", "is_active": False, "price_cents": 0}
    del body[missing]
    response = client.put("/products/1", json=body)
    assert response.status_code == 422
    assert client.get("/products/1").json()["name"] == "保留的商品"


@pytest.mark.parametrize("price", [-1, 1.5, "1990", True, None, 2147483648])
def test_price_rejects_invalid_create_and_replace_without_mutation(client, price):
    original = client.post("/products", json={"name": "陶瓷杯", "price_cents": 1990}).json()
    rejected = client.post("/products", json={"name": "无效价格", "price_cents": price})
    assert rejected.status_code == 422
    assert rejected.json()["error"]["details"][0]["loc"] == ["body", "price_cents"]
    replacement = {"name": "不应保存", "description": "", "is_active": True, "price_cents": price}
    assert client.put("/products/1", json=replacement).status_code == 422
    assert client.get("/products/1").json() == original
    assert client.get("/products").json() == [original]


@pytest.mark.parametrize("price", [0, 1990, 2147483647])
def test_price_integer_boundaries_are_preserved(client, price):
    response = client.post("/products", json={"name": "陶瓷杯", "price_cents": price})
    assert response.status_code == 201
    assert response.json()["price_cents"] == price
    assert type(response.json()["price_cents"]) is int


def test_description_size_matches_database_text_budget(client):
    original = client.post("/products", json={"name": "陶瓷杯", "description": "杯" * 16000})
    assert original.status_code == 201
    assert client.post("/products", json={"name": "过长描述", "description": "杯" * 16001}).status_code == 422
    replacement = {"name": "过长描述", "description": "杯" * 16001, "is_active": False, "price_cents": 0}
    assert client.put("/products/1", json=replacement).status_code == 422
    assert client.get("/products/1").json() == original.json()


@pytest.mark.parametrize("method,kwargs", [
    ("get", {}),
    ("put", {"json": {"name": "missing", "description": "", "is_active": False, "price_cents": 0}}),
    ("delete", {}),
])
def test_missing_product(client, method, kwargs):
    response = getattr(client, method)("/products/999", **kwargs)
    assert response.status_code == 404
    assert response.json() == {"error": {"code": "http_error", "message": "Product not found", "details": []}}


@pytest.mark.parametrize("url", ["/products?limit=0", "/products?limit=101", "/products?is_active=invalid", "/products/no-number", "/slow?delay=-1", "/slow?delay=6"])
def test_query_and_path_validation(client, url):
    response = client.get(url)
    assert response.status_code == 422
    assert response.json()["error"]["code"] == "validation_error"


def test_private_key_is_required_and_environment_is_read(monkeypatch, client):
    monkeypatch.delenv("DEMO_API_KEY", raising=False)
    assert client.get("/private", headers={"X-API-Key": "classroom-key"}).status_code == 503
    monkeypatch.setenv("DEMO_API_KEY", "classroom-key")
    missing = client.get("/private")
    assert missing.status_code == 401
    assert missing.headers["www-authenticate"] == "APIKey"
    assert client.get("/private", headers={"X-API-Key": "wrong"}).status_code == 401
    assert client.get("/private", headers={"X-API-Key": "classroom-key"}).status_code == 200
    assert client.get("/health").status_code == 200


def test_generate_and_sse_protocol(client):
    prompt = '你好 "world"\n下一行'
    response = client.post("/generate", json={"prompt": prompt})
    assert response.status_code == 200
    expected = response.json()["text"]
    stream = client.get("/stream", params={"prompt": prompt})
    assert stream.status_code == 200
    assert stream.headers["content-type"].startswith("text/event-stream")
    assert stream.text.endswith("\n\n")
    frames = stream.text.rstrip("\n").split("\n\n")
    chunks = []
    for frame in frames[:-1]:
        event, data = frame.split("\n")
        assert event == "event: token"
        assert data.startswith("data: ")
        chunks.append(json.loads(data.removeprefix("data: "))["text"])
    assert "".join(chunks) == expected
    assert frames[-1] == "event: done\ndata: {}"


def test_empty_generate_prompt_is_rejected(client):
    assert client.post("/generate", json={"prompt": " "}).status_code == 422
    assert client.get("/stream", params={"prompt": ""}).status_code == 422


def test_websocket_echo_ping_and_reconnect(client):
    with client.websocket_connect("/ws") as websocket:
        websocket.send_text("你好")
        assert websocket.receive_json() == {"type": "echo", "text": "你好"}
        websocket.send_text("ping")
        assert websocket.receive_json() == {"type": "pong"}
    with client.websocket_connect("/ws") as websocket:
        websocket.send_text("再次连接")
        assert websocket.receive_json()["text"] == "再次连接"


def test_cors_preflight_permits_local_client_only(client):
    headers = {
        "Origin": "http://localhost:5500",
        "Access-Control-Request-Method": "POST",
        "Access-Control-Request-Headers": "content-type",
    }
    accepted = client.options("/products", headers=headers)
    assert accepted.status_code == 200
    assert accepted.headers["access-control-allow-origin"] == "http://localhost:5500"
    headers["Origin"] = "https://example.org"
    rejected = client.options("/products", headers=headers)
    assert rejected.status_code == 400
    assert "access-control-allow-origin" not in rejected.headers


def test_slow_completes(client):
    assert client.get("/slow?delay=0").json() == {"delay": 0, "status": "done"}


def test_lifespan_opens_and_closes_http_client():
    with TestClient(app):
        http_client = app.state.http
        assert not http_client.is_closed
    assert http_client.is_closed


@pytest.mark.parametrize("scenario,status", [
    ("success", 200), ("timeout", 504), ("http_failure", 502), ("connection_failure", 502),
])
def test_upstream_result_and_error_mapping(client, monkeypatch, scenario, status):
    def handler(request):
        assert request.url.path == "/slow"
        assert request.url.params["delay"] == "1.0"
        if scenario == "timeout":
            raise httpx.ReadTimeout("simulated timeout", request=request)
        if scenario == "connection_failure":
            raise httpx.ConnectError("simulated connection failure", request=request)
        if scenario == "http_failure":
            return httpx.Response(503, json={"message": "unavailable"})
        return httpx.Response(200, json={"delay": 1, "status": "done"})

    async def mocked_get(path, *, params):
        async with httpx.AsyncClient(
            transport=httpx.MockTransport(handler), base_url="http://upstream.test",
        ) as upstream:
            return await upstream.get(path, params=params)

    monkeypatch.setattr(app.state.http, "get", mocked_get)
    response = client.get("/upstream-demo?delay=1")
    assert response.status_code == status
    if status == 200:
        assert response.json() == {"upstream": {"delay": 1, "status": "done"}}
    else:
        error = response.json()["error"]
        assert error["code"] == "http_error"
        expected_messages = {
            "timeout": "上游响应超时",
            "http_failure": "上游返回非成功状态",
            "connection_failure": "无法连接上游服务",
        }
        assert error["message"] == expected_messages[scenario]