tzt
f24f016e95
- upstream.py：httpx.AsyncClient 模块级单例（keepalive，limits=100， 超时 connect10/read120/write10/pool30）；stream() 始终注入 stream_options.include_usage（计量不依赖客户端）+ filter_usage_chunk （客户端未要求 usage 时剥除该 chunk）；D-P4 首 token 前 failover 链、 流中失败抛 UpstreamAborted（不可切换）；ttfb_ms 记录 - normalize_usage 三家归一：deepseek prompt_cache_hit_tokens / openai prompt_tokens_details.cached_tokens / anthropic cache_read_input_tokens （命中数>总数时钳制） - model_pool：+provider（枚举校验 deepseek/openai/anthropic）+in_hit_price （缺省 = price_in×1/30，D-P3） - 测试 +6：透传+sink/failover/流中 aborted/全挂 502/三家归一/过滤，全量 345 passed
207 lines
6.8 KiB
Python
207 lines
6.8 KiB
Python
"""上游客户端测试（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]"
|