test: 添加 Gemini 流式解析和格式转换相关测试

- 添加 Gemini JSON-array 兜底解析测试
- 添加 Gemini stream_parser 单元测试
- 添加 provider transport URL 构建测试
This commit is contained in:
fawney19
2026-01-23 23:36:08 +08:00
parent 76ed7e9fcb
commit 6cf251b19a
3 changed files with 187 additions and 0 deletions

View File

@@ -0,0 +1,102 @@
import json
from typing import AsyncIterator, Optional
import pytest
from src.api.handlers.base.parsers import get_parser_for_format
from src.api.handlers.base.stream_context import StreamContext
from src.api.handlers.base.stream_processor import StreamProcessor
from src.core.api_format import GeminiToOpenAIConverter, converter_registry
class _DummyResponseCtx:
async def __aexit__(self, exc_type, exc, tb) -> None: # noqa: ANN001
return None
class _DummyHTTPClient:
async def aclose(self) -> None:
return None
async def _iter_bytes(chunks: list[bytes]) -> AsyncIterator[bytes]:
for c in chunks:
yield c
@pytest.mark.asyncio
async def test_stream_processor_converts_gemini_json_lines_without_data_prefix() -> None:
# Register only what we need, and restore after test.
original_converters = converter_registry._converters.copy()
try:
converter_registry.register("GEMINI", "OPENAI", GeminiToOpenAIConverter())
ctx = StreamContext(model="gemini-test", api_format="OPENAI")
ctx.provider_api_format = "GEMINI"
ctx.client_api_format = "OPENAI"
ctx.needs_conversion = True
ctx.request_id = "req_test"
ctx.mapped_model = "gemini-test"
# Simulate Gemini JSON-array/chunks stream: wrapper lines + two JSON objects.
chunk1 = {
"candidates": [
{
"content": {"parts": [{"text": "Hello"}], "role": "model"},
}
]
}
chunk2 = {
"candidates": [
{
"content": {"parts": [{"text": "Hello world"}], "role": "model"},
"finishReason": "STOP",
}
],
"usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 2, "totalTokenCount": 3},
}
upstream_lines = [
b"[\n",
(json.dumps(chunk1) + ",\n").encode("utf-8"),
(json.dumps(chunk2) + "\n").encode("utf-8"),
b"]\n",
]
processor = StreamProcessor(
request_id="req_test",
default_parser=get_parser_for_format("OPENAI"),
)
out = b""
async for b in processor.create_response_stream(
ctx=ctx,
byte_iterator=_iter_bytes(upstream_lines),
response_ctx=_DummyResponseCtx(),
http_client=_DummyHTTPClient(), # type: ignore[arg-type]
prefetched_chunks=[],
start_time=None,
):
out += b
text = out.decode("utf-8", errors="replace")
data_lines = [ln for ln in text.splitlines() if ln.startswith("data: ")]
# OpenAI termination marker should be present (StreamProcessor will append if upstream doesn't send it).
assert "data: [DONE]" in data_lines
# Parse JSON events (excluding [DONE]) and validate we have expected deltas.
events = [json.loads(ln[6:]) for ln in data_lines if ln != "data: [DONE]"]
delta_contents: list[str] = []
for evt in events:
for choice in evt.get("choices", []) or []:
delta = choice.get("delta") or {}
if "content" in delta and delta["content"]:
delta_contents.append(delta["content"])
assert "Hello" in "".join(delta_contents)
assert " world" in "".join(delta_contents)
finally:
converter_registry._converters = original_converters

View File

@@ -0,0 +1,20 @@
from src.api.handlers.gemini.stream_parser import GeminiStreamParser
def test_is_done_event_true_for_new_finish_reason_values() -> None:
parser = GeminiStreamParser()
event = {"candidates": [{"finishReason": "MALFORMED_FUNCTION_CALL"}]}
assert parser.is_done_event(event) is True
def test_is_done_event_false_for_unspecified() -> None:
parser = GeminiStreamParser()
event = {"candidates": [{"finishReason": "FINISH_REASON_UNSPECIFIED"}]}
assert parser.is_done_event(event) is False
def test_is_done_event_false_when_no_candidates_or_reason() -> None:
parser = GeminiStreamParser()
assert parser.is_done_event({}) is False
assert parser.is_done_event({"candidates": [{}]}) is False

View File

@@ -0,0 +1,65 @@
from dataclasses import dataclass
from src.services.provider.transport import build_provider_url
@dataclass
class _DummyEndpoint:
base_url: str
api_format: str
custom_path: str | None = None
def test_gemini_stream_adds_alt_sse_and_drops_key_query_param() -> None:
endpoint = _DummyEndpoint(
base_url="https://generativelanguage.googleapis.com",
api_format="GEMINI",
)
url = build_provider_url(
endpoint, # type: ignore[arg-type] - test stub
query_params={"key": "SECRET"},
path_params={"model": "gemini-1.5-pro"},
is_stream=True,
)
assert url.startswith(
"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-pro:streamGenerateContent"
)
assert "key=" not in url
assert "alt=sse" in url
def test_gemini_stream_does_not_override_existing_alt() -> None:
endpoint = _DummyEndpoint(
base_url="https://generativelanguage.googleapis.com",
api_format="GEMINI",
)
url = build_provider_url(
endpoint, # type: ignore[arg-type] - test stub
query_params={"alt": "json"},
path_params={"model": "gemini-1.5-pro"},
is_stream=True,
)
assert "alt=json" in url
assert "alt=sse" not in url
def test_gemini_non_stream_does_not_add_alt() -> None:
endpoint = _DummyEndpoint(
base_url="https://generativelanguage.googleapis.com",
api_format="GEMINI",
)
url = build_provider_url(
endpoint, # type: ignore[arg-type] - test stub
path_params={"model": "gemini-1.5-pro"},
is_stream=False,
)
assert url.endswith("/v1beta/models/gemini-1.5-pro:generateContent")
assert "alt=" not in url