"""Юнит-тесты для `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 == "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. " "" ) # Один большой 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 "" 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"