projectAIpopular/tests/test_proxy_upstream.py

"""上游客户端测试（T-P2）：SSE 透传 / usage 注入过滤 / 三家归一化 / failover / aborted。"""
import asyncio
import json

import httpx
import pytest

from gateway.proxy.errors import UpstreamError
from gateway.proxy.upstream import (
    UpstreamAborted,
    filter_usage_chunk,
    normalize_usage,
    stream,
)


def asyncio_run(coro):
    return asyncio.run(coremaybe(coro))


async def coremaybe(coro):
    return await coro


def _sse(lines) -> bytes:
    return ("\n\n".join(lines) + "\n\n").encode("utf-8")


def test_stream_passthrough_and_usage_sink():
    """SSE 逐块 yield；usage chunk 进 sink；始终注入 include_usage。"""
    body = _sse([
        'data: {"choices":[{"delta":{"content":"你"}}]}',
        'data: {"choices":[{"delta":{"content":"好"}}]}',
        'data: {"choices":[{"delta":{},"finish_reason":"stop"}],'
        '"usage":{"prompt_tokens":100,"prompt_cache_hit_tokens":60,"completion_tokens":20}}',
        "data: [DONE]",
    ])
    seen_bodies = []

    def handler(request: httpx.Request) -> httpx.Response:
        seen_bodies.append(json.loads(request.read()))
        return httpx.Response(200, content=body)

    transport = httpx.MockTransport(handler)
    import gateway.proxy.upstream as up
    client = httpx.AsyncClient(transport=transport)
    orig = up._client
    up._client = client
    try:
        entry = {"base_url": "http://up", "api_key": "k", "provider": "deepseek"}

        async def run():
            sink = {}
            chunks = []
            async for b in stream({"model": "deepseek-chat", "messages": []}, entry, sink):
                chunks.append(b)
            return sink, chunks

        sink, chunks = asyncio_run(run())
    finally:
        up._client = orig
        asyncio.get_event_loop_policy()
        client.__asyncio_run = None
        try:
            asyncio.get_event_loop().run_until_complete(client.aclose())
        except Exception:
            pass
    text = b"".join(chunks).decode("utf-8")
    assert "\u4f60" in text and "\u597d" in text and "data: [DONE]" in text
    assert sink["usage"] == {"in_miss": 40, "in_hit": 60, "out": 20}
    assert isinstance(sink["ttfb_ms"], int)
    assert seen_bodies[0]["stream_options"] == {"include_usage": True}


def test_failover_before_first_token():
    """首条目连不上 -> 切换备用 -> 正常输出。"""
    ok_body = _sse(['data: {"choices":[{"delta":{"content":"ok"}}]}', "data: [DONE]"])

    def handler(request: httpx.Request) -> httpx.Response:
        if request.url.host == "bad":
            raise httpx.ConnectError("拒连")
        return httpx.Response(200, content=ok_body)

    transport = httpx.MockTransport(handler)
    import gateway.proxy.upstream as up
    client = httpx.AsyncClient(transport=transport)
    orig = up._client
    up._client = client
    try:
        async def run():
            sink = {}
            return [b async for b in stream(
                {"model": "m", "messages": []},
                {"base_url": "http://bad", "api_key": "", "provider": "openai"},
                sink, [{"base_url": "http://good", "api_key": "", "provider": "openai"}])]

        chunks = asyncio_run(run())
        assert b"ok" in b"".join(chunks)
    finally:
        up._client = orig
        try:
            asyncio.run(client.aclose())
        except Exception:
            pass


def test_abort_after_first_token_no_failover():
    """首 token 已 yield 后断流 -> UpstreamAborted（不切换，D-P4）。"""
    class StreamResp:
        def raise_for_status(self):
            pass

        async def aiter_lines(self):
            yield 'data: {"choices":[{"delta":{"content":"前半"}}]}'
            raise httpx.ReadError("中途断")

    class CM:
        async def __aenter__(self):
            return StreamResp()

        async def __aexit__(self, *a):
            return False

    class FakeClient:
        def stream(self, *a, **k):
            return CM()

    import gateway.proxy.upstream as up
    orig = up._client
    up._client = FakeClient()
    try:
        async def run():
            sink = {}
            gen = stream({"model": "m", "messages": []},
                         {"base_url": "http://x", "api_key": "", "provider": "openai"}, sink)
            first = await gen.__anext__()
            assert "\u524d\u534a".encode("utf-8") in first   # "前半"
            rest = []
            async for b in gen:
                rest.append(b)
            return rest

        try:
            asyncio_run(run())
            raised = False
        except UpstreamAborted:
            raised = True
        assert raised, "应抛 UpstreamAborted"
    finally:
        up._client = orig


def test_all_candidates_fail_raises_upstream_error():
    """全部候选首 token 前失败 -> UpstreamError(502)。"""
    import gateway.proxy.upstream as up

    class FakeClient:
        def stream(self, *a, **k):
            raise httpx.ConnectError("全挂")

    orig = up._client
    up._client = FakeClient()
    try:
        async def run():
            async for _ in stream({"model": "m", "messages": []},
                                  {"base_url": "http://a", "api_key": "", "provider": "openai"},
                                  {}, [{"base_url": "http://b", "api_key": "",
                                        "provider": "openai"}]):
                pass

        try:
            asyncio_run(run())
            raised = False
        except UpstreamError:
            raised = True
        assert raised
    finally:
        up._client = orig


def test_normalize_usage_three_providers():
    """deepseek / openai 兼容 / anthropic 三家字段归一；命中数钳制。"""
    assert normalize_usage("deepseek", {
        "prompt_tokens": 100, "prompt_cache_hit_tokens": 60, "completion_tokens": 20
    }) == {"in_miss": 40, "in_hit": 60, "out": 20}
    assert normalize_usage("openai", {
        "prompt_tokens": 100,
        "prompt_tokens_details": {"cached_tokens": 30},
        "completion_tokens": 20,
    }) == {"in_miss": 70, "in_hit": 30, "out": 20}
    assert normalize_usage("anthropic", {
        "input_tokens": 100, "cache_read_input_tokens": 90, "output_tokens": 5
    }) == {"in_miss": 10, "in_hit": 90, "out": 5}
    assert normalize_usage("deepseek", {
        "prompt_tokens": 10, "prompt_cache_hit_tokens": 99, "completion_tokens": 1
    })["in_miss"] == 0


def test_filter_usage_chunk():
    """客户端未要求 usage 时剥除 usage chunk；要求时透传；[DONE] 恒透传。"""
    usage_line = 'data: {"choices":[],"usage":{"prompt_tokens":1}}'
    normal_line = 'data: {"choices":[{"delta":{"content":"x"}}]}'
    assert filter_usage_chunk(usage_line, client_wants_usage=False) is None
    assert filter_usage_chunk(usage_line, client_wants_usage=True) == usage_line
    assert filter_usage_chunk(normal_line, client_wants_usage=False) == normal_line
    assert filter_usage_chunk("data: [DONE]", False) == "data: [DONE]"