omnisvg-lite/tests/test_lm_streaming.py

"""Юнит-тесты для `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"