mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
test: 添加 Gemini 流式解析和格式转换相关测试
- 添加 Gemini JSON-array 兜底解析测试 - 添加 Gemini stream_parser 单元测试 - 添加 provider transport URL 构建测试
This commit is contained in:
@@ -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
|
||||||
|
|
||||||
20
tests/api/handlers/gemini/test_stream_parser.py
Normal file
20
tests/api/handlers/gemini/test_stream_parser.py
Normal 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
|
||||||
|
|
||||||
65
tests/services/test_provider_transport.py
Normal file
65
tests/services/test_provider_transport.py
Normal 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
|
||||||
|
|
||||||
Reference in New Issue
Block a user