Finalize live-streaming feature: docs and tests
- docs/live_streaming.md: feature description, perf, limitations - 183 tests passing (was 157; added 26+ for streaming + live UI) - All previous regressions fixed Owner-action: completed final-integration myself after tester session got stuck on the e2e attempt (likely trying to spawn a real Gradio on an already-busy port). Manual verification: 183 passed, 1 skipped, 0 failed; feature works end-to-end via Gradio UI on 127.0.0.1:8788.
This commit is contained in:
Mavis
@@ -0,0 +1,458 @@
|
||||
"""Юнит-тесты для `stream_chat()` и `StreamEvent` в lm_client.py.
|
||||
|
||||
Покрывают:
|
||||
- Mock SSE-ответ: 3 data-чанка + [DONE] → 3 delta-события + 1 end
|
||||
- HTTP 500: stream_chat() бросает LMStudioUnavailable
|
||||
- HTTP 4xx: stream_chat() бросает LMStudioUnavailable
|
||||
- Timeout: stream_chat() бросает LMStudioUnavailable (с упоминанием timeout)
|
||||
- Пустой стрим (только [DONE]): 0 delta'ов, 1 end
|
||||
- n>1 в стриме: WARNING в логах, payload уходит с n=1
|
||||
- n=0 валидация
|
||||
- Модель/usage/finish_reason приходят в end из последнего чанка
|
||||
- Skip не-SSE строк (комментариев/heartbeat)
|
||||
- Битый JSON в SSE → LMStudioUnavailable
|
||||
- reasoning-токены из content НЕ отделяются (проходят как обычный delta)
|
||||
|
||||
httpx мокается через unittest.mock: подменяем `httpx.Client`,
|
||||
а у экземпляра — `.stream(...)` (это context manager, отдаёт response
|
||||
с iter_lines).
|
||||
|
||||
Запуск: `python -m pytest tests/test_lm_streaming.py -v`
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import sys
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterator
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
ROOT = Path(__file__).resolve().parent.parent
|
||||
if str(ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(ROOT))
|
||||
|
||||
import pytest # noqa: E402
|
||||
|
||||
import httpx # noqa: E402
|
||||
|
||||
from lm_client import ( # noqa: E402
|
||||
LMStudioUnavailable,
|
||||
StreamEvent,
|
||||
stream_chat,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Хелперы для построения mock-стрима
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _sse_chunk(*deltas: str, model: str = "qwen3.5", finish_reason: str | None = None) -> list[dict]:
|
||||
"""Собирает фиктивный JSON-чанк в стиле OpenAI-SSE.
|
||||
|
||||
Каждый delta — это content, который LM Studio кладёт в
|
||||
`choices[0].delta.content`. Если задан finish_reason, он попадает
|
||||
в этот же чанк (как в финальном чанке OpenAI-стрима).
|
||||
"""
|
||||
choice: dict[str, Any] = {"index": 0, "delta": {"role": "assistant"}}
|
||||
if deltas:
|
||||
choice["delta"]["content"] = "".join(deltas) if len(deltas) > 1 else deltas[0]
|
||||
if finish_reason:
|
||||
choice["finish_reason"] = finish_reason
|
||||
choice["delta"]["content"] = "".join(deltas) if deltas else ""
|
||||
return [{"id": "cmpl-x", "object": "chat.completion.chunk", "model": model, "choices": [choice]}]
|
||||
|
||||
|
||||
def _make_stream_response(
|
||||
*,
|
||||
sse_lines: list[str] | None = None,
|
||||
sse_iter: Iterator[str] | None = None,
|
||||
status_code: int = 200,
|
||||
error_body: str = "",
|
||||
) -> MagicMock:
|
||||
"""Создаёт mock-объект, имитирующий httpx.Response внутри `client.stream(...)`.
|
||||
|
||||
Args:
|
||||
sse_lines: фиксированный список строк (для простых случаев) — имитация
|
||||
тела SSE. Каждая строка — это ровно одна строка из `iter_lines()`.
|
||||
sse_iter: кастомный итератор (если хотим имитировать ошибку посреди стрима).
|
||||
status_code: HTTP status.
|
||||
error_body: тело для случая status_code >= 400 (читается через .read()).
|
||||
"""
|
||||
resp = MagicMock(spec=httpx.Response)
|
||||
resp.status_code = status_code
|
||||
|
||||
if sse_iter is not None:
|
||||
resp.iter_lines.return_value = sse_iter
|
||||
else:
|
||||
resp.iter_lines.return_value = sse_lines or []
|
||||
|
||||
if status_code >= 400:
|
||||
# resp.read() вызывается в stream_chat для превью ошибки.
|
||||
resp.read.return_value = error_body.encode("utf-8")
|
||||
else:
|
||||
resp.read.return_value = b""
|
||||
|
||||
return resp
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patched_stream(response: MagicMock | None = None, side_effect: Exception | None = None):
|
||||
"""Патчит `httpx.Client` так, что `client.stream(...)` отдаёт заданный response.
|
||||
|
||||
Args:
|
||||
response: mock-Response (то, что отдаёт __enter__ контекст-менеджера).
|
||||
side_effect: если задан — `client.stream(...)` бросает это исключение
|
||||
ДО входа в context (имитация timeout на этапе open).
|
||||
"""
|
||||
with patch("lm_client.httpx.Client") as MockClient:
|
||||
mock_inst = MagicMock()
|
||||
MockClient.return_value.__enter__.return_value = mock_inst
|
||||
|
||||
if side_effect is not None:
|
||||
mock_inst.stream.side_effect = side_effect
|
||||
else:
|
||||
# .stream(...) возвращает context manager, у которого __enter__
|
||||
# возвращает наш response.
|
||||
cm = MagicMock()
|
||||
cm.__enter__.return_value = response
|
||||
cm.__exit__.return_value = False
|
||||
mock_inst.stream.return_value = cm
|
||||
yield mock_inst
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Успешный стрим: 3 чанка + [DONE]
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_stream_chat_parses_three_sse_chunks():
|
||||
"""SSE из 3 чанков + [DONE] → 3 StreamEvent(type='delta') + 1 end."""
|
||||
sse_lines = [
|
||||
'data: {"id":"1","choices":[{"index":0,"delta":{"content":"<"}}]}',
|
||||
"",
|
||||
'data: {"id":"2","choices":[{"index":0,"delta":{"content":"svg "}}]}',
|
||||
"",
|
||||
'data: {"id":"3","choices":[{"index":0,"delta":{"content":"xmlns=..."}}]}',
|
||||
"",
|
||||
"data: [DONE]",
|
||||
"",
|
||||
]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
with _patched_stream(response=resp) as mock_inst:
|
||||
events = list(stream_chat(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
base_url="http://m:1/v1",
|
||||
))
|
||||
|
||||
# 3 delta + 1 end = 4 события.
|
||||
assert len(events) == 4
|
||||
assert all(e.type == "delta" for e in events[:3])
|
||||
assert events[3].type == "end"
|
||||
|
||||
# Содержимое каждого delta.
|
||||
assert events[0].content == "<"
|
||||
assert events[1].content == "svg "
|
||||
assert events[2].content == "xmlns=..."
|
||||
|
||||
# Конкатенация = полный текст.
|
||||
full = "".join(e.content for e in events if e.type == "delta")
|
||||
assert full == "<svg xmlns=..."
|
||||
|
||||
# End-event: пустой content, но type='end'.
|
||||
assert events[3].content == ""
|
||||
assert events[3].finish_reason == ""
|
||||
|
||||
|
||||
def test_stream_chat_payload_uses_stream_true_and_n_one():
|
||||
"""Payload должен содержать stream=True и n=1."""
|
||||
sse_lines = ['data: {"choices":[{"delta":{"content":"x"}}]}', "", "data: [DONE]", ""]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
with _patched_stream(response=resp) as mock_inst:
|
||||
list(stream_chat(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
n=1,
|
||||
temperature=0.7,
|
||||
max_tokens=1024,
|
||||
base_url="http://m:1/v1",
|
||||
))
|
||||
|
||||
call = mock_inst.stream.call_args
|
||||
body = call.kwargs["json"]
|
||||
assert body["stream"] is True
|
||||
assert body["n"] == 1
|
||||
assert body["temperature"] == 0.7
|
||||
assert body["max_tokens"] == 1024
|
||||
assert body["model"] # non-empty
|
||||
|
||||
# URL сформирован правильно.
|
||||
method, url = call.args[0], call.args[1]
|
||||
assert method == "POST"
|
||||
assert url == "http://m:1/v1/chat/completions"
|
||||
headers = call.kwargs["headers"]
|
||||
assert headers["Authorization"] == "Bearer lm-studio"
|
||||
assert headers["Content-Type"] == "application/json"
|
||||
|
||||
|
||||
def test_stream_chat_collects_model_and_usage_in_end_event():
|
||||
"""Последний чанк содержит model/usage/finish_reason — это попадает в end."""
|
||||
sse_lines = [
|
||||
'data: {"model":"qwen3.5","choices":[{"delta":{"content":"Hel"}}]}',
|
||||
"",
|
||||
'data: {"model":"qwen3.5","choices":[{"delta":{"content":"lo"}}],"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}}',
|
||||
"",
|
||||
'data: {"choices":[{"delta":{},"finish_reason":"stop"}]}',
|
||||
"",
|
||||
"data: [DONE]",
|
||||
"",
|
||||
]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
with _patched_stream(response=resp):
|
||||
events = list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
|
||||
assert len(events) == 3 # 2 delta + 1 end
|
||||
end = events[-1]
|
||||
assert end.type == "end"
|
||||
assert end.model == "qwen3.5"
|
||||
assert end.usage == {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}
|
||||
assert end.finish_reason == "stop"
|
||||
|
||||
|
||||
def test_stream_chat_handles_empty_stream_only_done():
|
||||
"""Стрим сразу [DONE] → 0 delta'ов, 1 end."""
|
||||
sse_lines = ["data: [DONE]", ""]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
with _patched_stream(response=resp):
|
||||
events = list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
|
||||
assert len(events) == 1
|
||||
assert events[0].type == "end"
|
||||
assert events[0].content == ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Ошибки HTTP / network / timeout
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_stream_chat_500_raises_lmstudio_unavailable():
|
||||
"""HTTP 500 → LMStudioUnavailable, в сообщении есть 500 и body preview."""
|
||||
resp = _make_stream_response(status_code=500, error_body="internal error")
|
||||
with _patched_stream(response=resp):
|
||||
with pytest.raises(LMStudioUnavailable) as excinfo:
|
||||
list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
msg = str(excinfo.value)
|
||||
assert "500" in msg
|
||||
assert "internal error" in msg
|
||||
|
||||
|
||||
def test_stream_chat_4xx_raises_lmstudio_unavailable():
|
||||
"""HTTP 401 → LMStudioUnavailable (4xx — тоже клиентская ошибка)."""
|
||||
resp = _make_stream_response(status_code=401, error_body="unauthorized")
|
||||
with _patched_stream(response=resp):
|
||||
with pytest.raises(LMStudioUnavailable) as excinfo:
|
||||
list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
msg = str(excinfo.value)
|
||||
assert "401" in msg
|
||||
assert "unauthorized" in msg
|
||||
|
||||
|
||||
def test_stream_chat_timeout_raises_lmstudio_unavailable():
|
||||
"""httpx.TimeoutException на open стрима → LMStudioUnavailable с 'timeout'."""
|
||||
with _patched_stream(side_effect=httpx.TimeoutException("stream timed out")):
|
||||
with pytest.raises(LMStudioUnavailable) as excinfo:
|
||||
list(stream_chat(
|
||||
messages=[{"role": "user", "content": "x"}],
|
||||
base_url="http://m:1/v1",
|
||||
timeout_s=5.0,
|
||||
))
|
||||
msg = str(excinfo.value).lower()
|
||||
assert "timeout" in msg
|
||||
|
||||
|
||||
def test_stream_chat_network_error_raises_lmstudio_unavailable():
|
||||
"""Любой httpx.HTTPError → LMStudioUnavailable."""
|
||||
with _patched_stream(side_effect=httpx.ConnectError("connection refused")):
|
||||
with pytest.raises(LMStudioUnavailable) as excinfo:
|
||||
list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
msg = str(excinfo.value).lower()
|
||||
assert "недоступен" in msg or "unavailable" in msg
|
||||
assert "connection refused" in msg
|
||||
|
||||
|
||||
def test_stream_chat_broken_sse_json_raises_lmstudio_unavailable():
|
||||
"""Битый JSON в SSE-чанке → LMStudioUnavailable."""
|
||||
sse_lines = [
|
||||
'data: {"choices":[{"delta":{"content":"Hel"}}]}',
|
||||
"",
|
||||
'data: this is not json',
|
||||
"",
|
||||
]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
with _patched_stream(response=resp):
|
||||
with pytest.raises(LMStudioUnavailable) as excinfo:
|
||||
list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
msg = str(excinfo.value).lower()
|
||||
assert "не-json" in msg or "sse" in msg
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# n>1 в стриме → WARNING + n=1 в payload
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_stream_chat_n_greater_than_one_warns_and_uses_n_one(caplog):
|
||||
"""n>1 в стриме: логируется WARNING, payload уходит с n=1."""
|
||||
sse_lines = ['data: {"choices":[{"delta":{"content":"x"}}]}', "", "data: [DONE]", ""]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="lm_client"):
|
||||
with _patched_stream(response=resp) as mock_inst:
|
||||
events = list(stream_chat(
|
||||
messages=[{"role": "user", "content": "x"}],
|
||||
n=4,
|
||||
base_url="http://m:1/v1",
|
||||
))
|
||||
|
||||
# Payload ушёл с n=1.
|
||||
body = mock_inst.stream.call_args.kwargs["json"]
|
||||
assert body["n"] == 1
|
||||
|
||||
# В логах было WARNING.
|
||||
assert any(
|
||||
"n=" in rec.message and "stream" in rec.message.lower()
|
||||
for rec in caplog.records if rec.levelno == logging.WARNING
|
||||
)
|
||||
|
||||
# Стрим всё равно отработал.
|
||||
assert any(e.type == "delta" for e in events)
|
||||
assert events[-1].type == "end"
|
||||
|
||||
|
||||
def test_stream_chat_rejects_n_less_than_one():
|
||||
"""n=0 → ValueError до обращения к сети."""
|
||||
with pytest.raises(ValueError, match="n должно быть"):
|
||||
list(stream_chat(
|
||||
messages=[{"role": "user", "content": "x"}],
|
||||
n=0,
|
||||
base_url="http://m:1/v1",
|
||||
))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Малые кейсы: пропуск мусорных строк, reasoning в content
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_stream_chat_skips_non_data_lines():
|
||||
"""Строки без префикса `data:` (комментарии, event:, id:) — пропускаются."""
|
||||
sse_lines = [
|
||||
": this is a comment", # SSE-комментарий
|
||||
"event: message", # event-поле
|
||||
"id: 42", # id-поле
|
||||
'data: {"choices":[{"delta":{"content":"Hi"}}]}',
|
||||
"",
|
||||
"data: [DONE]",
|
||||
"",
|
||||
]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
with _patched_stream(response=resp):
|
||||
events = list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
|
||||
deltas = [e for e in events if e.type == "delta"]
|
||||
assert len(deltas) == 1
|
||||
assert deltas[0].content == "Hi"
|
||||
assert events[-1].type == "end"
|
||||
|
||||
|
||||
def test_stream_chat_reasoning_tokens_pass_through_unchanged():
|
||||
"""qwen3.5 кладёт reasoning в основной content — мы НЕ отделяем."""
|
||||
reasoning_then_svg = (
|
||||
"Let me think about this. I need an icon of a fox. The viewBox is 64x64. "
|
||||
"<svg viewBox='0 0 64 64'><circle cx='32' cy='32' r='10'/></svg>"
|
||||
)
|
||||
# Один большой delta с reasoning+svg внутри.
|
||||
sse_lines = [
|
||||
f'data: {{"choices":[{{"delta":{{"content":{reasoning_then_svg!r}}}}}]}}',
|
||||
"",
|
||||
"data: [DONE]",
|
||||
"",
|
||||
]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
with _patched_stream(response=resp):
|
||||
events = list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
|
||||
full = "".join(e.content for e in events if e.type == "delta")
|
||||
assert "Let me think" in full
|
||||
assert "<svg" in full
|
||||
assert "</svg>" in full
|
||||
|
||||
|
||||
def test_stream_chat_handles_content_as_list_of_parts():
|
||||
"""content может быть list[dict] (мультимодальный стрим) — склеиваем в строку."""
|
||||
sse_lines = [
|
||||
'data: {"choices":[{"delta":{"content":[{"type":"text","text":"Hel"},{"type":"text","text":"lo"}]}}]}',
|
||||
"",
|
||||
"data: [DONE]",
|
||||
"",
|
||||
]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
with _patched_stream(response=resp):
|
||||
events = list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
|
||||
deltas = [e for e in events if e.type == "delta"]
|
||||
assert len(deltas) == 1
|
||||
assert deltas[0].content == "Hello"
|
||||
|
||||
|
||||
def test_stream_chat_emits_end_even_without_done_marker():
|
||||
"""Стрим без [DONE] всё равно получает end-event (с WARNING в логе)."""
|
||||
sse_lines = [
|
||||
'data: {"choices":[{"delta":{"content":"x"},"finish_reason":"stop"}]}',
|
||||
"",
|
||||
# нет [DONE] — обрыв по исчерпанию итератора
|
||||
]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
with _patched_stream(response=resp):
|
||||
events = list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
|
||||
assert events[-1].type == "end"
|
||||
assert events[-1].finish_reason == "stop"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sanity: StreamEvent — frozen dataclass с нужными полями
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_stream_event_is_frozen_dataclass():
|
||||
"""StreamEvent — frozen: попытка изменения атрибута → FrozenInstanceError."""
|
||||
ev = StreamEvent(type="delta", content="x")
|
||||
assert ev.type == "delta"
|
||||
assert ev.content == "x"
|
||||
assert ev.usage is None
|
||||
assert ev.model == ""
|
||||
assert ev.finish_reason == ""
|
||||
with pytest.raises(Exception): # FrozenInstanceError
|
||||
ev.type = "end" # type: ignore[misc]
|
||||
|
||||
|
||||
def test_stream_chat_strip_data_prefix_with_or_without_space():
|
||||
"""SSE-префикс `data:` допускает опциональный пробел после двоеточия."""
|
||||
sse_lines = [
|
||||
'data:{"choices":[{"delta":{"content":"A"}}]}', # без пробела
|
||||
"",
|
||||
'data: {"choices":[{"delta":{"content":"B"}}]}', # с пробелом
|
||||
"",
|
||||
"data:[DONE]", # без пробела
|
||||
"",
|
||||
]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
with _patched_stream(response=resp):
|
||||
events = list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
|
||||
deltas = [e for e in events if e.type == "delta"]
|
||||
assert [d.content for d in deltas] == ["A", "B"]
|
||||
assert events[-1].type == "end"
|
||||
Reference in New Issue
Block a user