84 lines
3.1 KiB
Python
84 lines
3.1 KiB
Python
"""Unit tests verifying stream_chat() captures a trailing usage-only SSE chunk.
|
|
|
|
Exercises DeepSeekClient, QwenClient, and QwenVLClient directly (not through
|
|
TrackedLLMClient) by mocking the underlying httpx.Client.stream() call — none
|
|
of these tests make a real network call.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from unittest.mock import MagicMock
|
|
|
|
from app.services.llm.base_client import LLMConfig, LLMProvider
|
|
from app.services.llm.deepseek_client import DeepSeekClient
|
|
|
|
|
|
def _sse_lines(*chunks: str, usage: dict | None = None) -> list[str]:
|
|
"""Build raw SSE 'data: ...' lines the way an OpenAI-compatible gateway sends them."""
|
|
lines = [
|
|
f'data: {json.dumps({"choices": [{"delta": {"content": c}}]})}'
|
|
for c in chunks
|
|
]
|
|
if usage is not None:
|
|
# Trailing usage-only chunk, as sent when stream_options.include_usage=true.
|
|
lines.append(f'data: {json.dumps({"choices": [], "usage": usage})}')
|
|
lines.append("data: [DONE]")
|
|
return lines
|
|
|
|
|
|
def _mock_streaming_client(lines: list[str]) -> MagicMock:
|
|
"""Build a MagicMock standing in for httpx.Client, configured for .stream()."""
|
|
fake_response = MagicMock()
|
|
fake_response.raise_for_status.return_value = None
|
|
fake_response.iter_lines.return_value = lines
|
|
|
|
stream_cm = MagicMock()
|
|
stream_cm.__enter__.return_value = fake_response
|
|
stream_cm.__exit__.return_value = False
|
|
|
|
client = MagicMock()
|
|
client.stream.return_value = stream_cm
|
|
return client
|
|
|
|
|
|
def _drain(gen):
|
|
"""Manually drive a generator, returning (yielded_chunks, stop_iteration_value)."""
|
|
chunks = []
|
|
value = None
|
|
while True:
|
|
try:
|
|
chunks.append(next(gen))
|
|
except StopIteration as stop:
|
|
value = stop.value
|
|
break
|
|
return chunks, value
|
|
|
|
|
|
def test_deepseek_stream_chat_returns_usage_from_trailing_chunk():
|
|
"""DeepSeekClient.stream_chat() must return the trailing usage dict."""
|
|
config = LLMConfig(provider=LLMProvider.DEEPSEEK, model="deepseek-v4-flash", api_key="k", base_url="http://x/v1")
|
|
client = DeepSeekClient(config)
|
|
usage = {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}
|
|
client._client = _mock_streaming_client(_sse_lines("Hello", " world", usage=usage))
|
|
|
|
chunks, returned_usage = _drain(client.stream_chat([{"role": "user", "content": "hi"}]))
|
|
|
|
assert chunks == ["Hello", " world"]
|
|
assert returned_usage == usage
|
|
# The gateway must actually be asked to include usage in the stream.
|
|
sent_payload = client._client.stream.call_args.kwargs["json"]
|
|
assert sent_payload["stream_options"] == {"include_usage": True}
|
|
|
|
|
|
def test_deepseek_stream_chat_without_usage_chunk_returns_none():
|
|
"""If the gateway never sends a usage chunk, the generator returns None (unchanged behavior)."""
|
|
config = LLMConfig(provider=LLMProvider.DEEPSEEK, model="deepseek-v4-flash", api_key="k", base_url="http://x/v1")
|
|
client = DeepSeekClient(config)
|
|
client._client = _mock_streaming_client(_sse_lines("Hi"))
|
|
|
|
chunks, returned_usage = _drain(client.stream_chat([{"role": "user", "content": "hi"}]))
|
|
|
|
assert chunks == ["Hi"]
|
|
assert returned_usage is None
|