import httpx

from yuqei_sdk import AiPlatformClient, LegalQaRequest, RequestContext


def test_ai_platform_client_posts_legal_qa_with_context_headers() -> None:
    captured_request: httpx.Request | None = None

    def handler(request: httpx.Request) -> httpx.Response:
        nonlocal captured_request
        captured_request = request
        return httpx.Response(
            200,
            json={
                "answer": "依据劳动合同法第十九条，需要核对合同期限。",
                "citations": [
                    {
                        "source_type": "law",
                        "title": "中华人民共和国劳动合同法",
                        "reference": "第十九条",
                        "quote": "劳动合同期限三个月以上不满一年的，试用期不得超过一个月。",
                    }
                ],
                "confidence": 0.72,
            },
        )

    client = AiPlatformClient(
        "http://ai-platform.test",
        transport=httpx.MockTransport(handler),
    )

    response = client.run_legal_qa(
        LegalQaRequest(question="试用期如何约定？"),
        context=RequestContext(trace_id="trace-1", tenant_id="tenant-1", user_id="user-1"),
    )

    assert response.citations[0].reference == "第十九条"
    assert captured_request is not None
    assert captured_request.url.path == "/api/v1/legal/qa"
    assert captured_request.headers["X-Trace-Id"] == "trace-1"
    assert captured_request.headers["X-Tenant-Id"] == "tenant-1"
    assert captured_request.headers["X-User-Id"] == "user-1"