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.notes.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 "/notes" in paths
    assert paths["/notes"]["post"]["responses"]["422"]["content"]["application/json"]["schema"] == {
        "$ref": "#/components/schemas/ErrorResponse"
    }


def test_crud_lifecycle(client):
    created = client.post("/notes", json={"title": "REST 笔记"})
    assert created.status_code == 201
    assert created.json() == {"id": 1, "title": "REST 笔记", "content": "", "pinned": False}
    assert created.headers["location"] == "/notes/1"
    assert client.get(created.headers["location"]).json() == created.json()
    assert client.get("/notes/1").json() == created.json()
    replacement = {"title": "修改后的标题", "content": "PUT 完整替换", "pinned": True}
    updated = client.put("/notes/1", json=replacement)
    assert updated.status_code == 200
    assert updated.json() == {"id": 1, **replacement}
    removed = client.delete("/notes/1")
    assert removed.status_code == 204
    assert removed.content == b""
    assert client.get("/notes").json() == []
    assert client.get("/notes/1").status_code == 404
    assert client.post("/notes", json={"title": "新笔记"}).json()["id"] == 2


def test_list_filter_and_limit(client):
    client.post("/notes", json={"title": "FastAPI", "content": "request"})
    client.post("/notes", json={"title": "第二篇", "content": "fastapi response"})
    client.post("/notes", json={"title": "Docker"})
    assert len(client.get("/notes", params={"q": "FASTAPI"}).json()) == 2
    assert len(client.get("/notes", params={"q": "fastapi", "limit": 1}).json()) == 1
    assert client.get("/notes", params={"q": "unknown"}).json() == []


def test_default_limit_and_combined_pinned_filter(client):
    for index in range(12):
        client.post("/notes", json={
            "title": f"note {index}",
            "content": "match" if index in (1, 2, 3) else "",
            "pinned": index % 2 == 0,
        })
    assert len(client.get("/notes").json()) == 10
    pinned = client.get("/notes", params={"pinned": "true"}).json()
    unpinned = client.get("/notes", params={"pinned": "false"}).json()
    assert len(pinned) == len(unpinned) == 6
    assert all(note["pinned"] is True for note in pinned)
    assert all(note["pinned"] is False for note in unpinned)
    combined = client.get("/notes", params={"q": "match", "pinned": "false", "limit": 1}).json()
    assert [note["id"] for note in combined] == [2]


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


@pytest.mark.parametrize("missing", ["title", "content", "pinned"])
def test_put_requires_every_mutable_field(client, missing):
    client.post("/notes", json={"title": "保留的笔记"})
    body = {"title": "替换", "content": "内容", "pinned": False}
    del body[missing]
    response = client.put("/notes/1", json=body)
    assert response.status_code == 422
    assert client.get("/notes/1").json()["title"] == "保留的笔记"


@pytest.mark.parametrize("method,kwargs", [
    ("get", {}),
    ("put", {"json": {"title": "missing", "content": "", "pinned": False}}),
    ("delete", {}),
])
def test_missing_note(client, method, kwargs):
    response = getattr(client, method)("/notes/999", **kwargs)
    assert response.status_code == 404
    assert response.json() == {"error": {"code": "http_error", "message": "Note not found", "details": []}}


@pytest.mark.parametrize("url", ["/notes?limit=0", "/notes?limit=101", "/notes?pinned=invalid", "/notes/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("/notes", 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("/notes", 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]