omnisvg-lite/tests/test_lm_client.py

"""Юнит-тесты для lm_client.py.

Покрывают:
- успешный ответ: парсинг N=4 choice'ов, текст, finish_reasons, model, usage
- таймаут: httpx.TimeoutException → LMStudioUnavailable с упоминанием timeout
- 5xx от LM Studio: status_code 503 → LMStudioUnavailable
- 4xx: status_code 401 / 404 → LMStudioUnavailable (по контракту: ошибка клиента
  или сервера — нам всё равно)
- пустой / битый JSON
- ответ без choices → LMStudioUnavailable
- ответ с tool_use (content=None, tool_calls есть) → raw_texts содержит пустую строку
- ответ с несколькими ```svg блоками: parse_svg идёт в валидаторе, тут проверим,
  что chat() возвращает всю строку модели как есть в raw_text
- generate_svg() высокоуровневая обёртка: проверка сообщений и prompt guard
- encode_pil_to_data_url(): кодирует PIL.Image в data: URL
- validate_image(): PIL.Image проходит / отклоняется по размеру / формату

httpx мокается через unittest.mock — клиент создаётся внутри `chat()`,
поэтому патчим `httpx.Client`.

Запуск: `python -m pytest tests/test_lm_client.py -v`
"""

from __future__ import annotations

import base64
import io
import sys
from pathlib import Path
from typing import Any
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
    DEFAULT_BASE_URL,
    LMStudioUnavailable,
    chat,
    encode_pil_to_data_url,
    generate_svg,
    validate_image,
)


# ---------------------------------------------------------------------------
# Хелперы для построения mock-ответа httpx
# ---------------------------------------------------------------------------


def _make_response(
    *,
    status_code: int = 200,
    json_payload: dict | None = None,
    text: str = "",
) -> MagicMock:
    """Создаёт mock httpx.Response с заданным status_code и JSON-телом."""
    resp = MagicMock(spec=httpx.Response)
    resp.status_code = status_code
    if json_payload is not None:
        resp.json.return_value = json_payload
    else:
        resp.json.side_effect = ValueError("not json")
    resp.text = text
    return resp


def _ok_payload(texts: list[str], *, model: str = "test-model") -> dict:
    """Стандартный OpenAI-style ответ с N choice'ами."""
    return {
        "id": "chatcmpl-test",
        "object": "chat.completion",
        "model": model,
        "choices": [
            {
                "index": i,
                "message": {"role": "assistant", "content": t},
                "finish_reason": "stop",
            }
            for i, t in enumerate(texts)
        ],
        "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
    }


# ---------------------------------------------------------------------------
# Успешный путь
# ---------------------------------------------------------------------------


def test_chat_success_returns_n_texts_and_metadata():
    """N=4: возвращаем 4 текста, finish_reasons, model, usage."""
    payload = _ok_payload(["a", "b", "c", "d"], model="my-model")
    resp = _make_response(json_payload=payload)

    with patch("lm_client.httpx.Client") as MockClient:
        mock_inst = MagicMock()
        mock_inst.post.return_value = resp
        MockClient.return_value.__enter__.return_value = mock_inst

        result = chat(
            messages=[{"role": "user", "content": "hi"}],
            model="my-model",
            n=4,
            temperature=0.5,
            base_url="http://mock:1234/v1",
        )

    assert isinstance(result.raw_texts, list)
    assert len(result.raw_texts) == 4
    assert result.raw_texts == ["a", "b", "c", "d"]
    assert result.finish_reasons == ["stop"] * 4
    assert result.model == "my-model"
    assert result.usage == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
    assert result.elapsed_s >= 0.0

    # Проверяем, что URL и заголовки формируются правильно.
    call = mock_inst.post.call_args
    url = call.args[0] if call.args else call.kwargs["url"]
    assert url == "http://mock:1234/v1/chat/completions"
    headers = call.kwargs["headers"]
    assert headers["Content-Type"] == "application/json"
    assert headers["Authorization"] == "Bearer lm-studio"
    body = call.kwargs["json"]
    assert body["model"] == "my-model"
    assert body["n"] == 4
    assert body["temperature"] == 0.5
    assert body["stream"] is False


def test_chat_success_single_candidate_default():
    """Дефолт n=1 — возвращаем ровно один текст."""
    resp = _make_response(json_payload=_ok_payload(["only one"]))

    with patch("lm_client.httpx.Client") as MockClient:
        mock_inst = MagicMock()
        mock_inst.post.return_value = resp
        MockClient.return_value.__enter__.return_value = mock_inst

        result = chat(
            messages=[{"role": "user", "content": "x"}],
            base_url="http://m:1/v1",
        )

    assert len(result.raw_texts) == 1
    assert result.raw_texts[0] == "only one"


def test_chat_server_returns_fewer_choices_pads_with_empty():
    """Сервер вернул 1 из 4 — добиваем пустыми строками и 'missing'."""
    payload = _ok_payload(["one"])  # n=4, но пришёл только 1
    resp = _make_response(json_payload=payload)

    with patch("lm_client.httpx.Client") as MockClient:
        mock_inst = MagicMock()
        mock_inst.post.return_value = resp
        MockClient.return_value.__enter__.return_value = mock_inst

        result = chat(
            messages=[{"role": "user", "content": "x"}],
            n=4,
            base_url="http://m:1/v1",
        )

    assert len(result.raw_texts) == 4
    assert result.raw_texts[0] == "one"
    assert result.raw_texts[1:] == ["", "", ""]
    assert result.finish_reasons == ["stop", "missing", "missing", "missing"]


# ---------------------------------------------------------------------------
# Ошибочные пути
# ---------------------------------------------------------------------------


def test_chat_timeout_raises_lmstudio_unavailable():
    """httpx.TimeoutException → LMStudioUnavailable, в тексте есть 'timeout'."""
    with patch("lm_client.httpx.Client") as MockClient:
        mock_inst = MagicMock()
        mock_inst.post.side_effect = httpx.TimeoutException("timed out")
        MockClient.return_value.__enter__.return_value = mock_inst

        with pytest.raises(LMStudioUnavailable) as excinfo:
            chat(
                messages=[{"role": "user", "content": "x"}],
                n=1,
                base_url="http://m:1/v1",
                timeout_s=5.0,
            )
    msg = str(excinfo.value).lower()
    assert "timeout" in msg
    # Код форматирует "5.0с" / "5с" — проверим, что число таймаута попало в сообщение.
    assert "5" in msg
    assert "с" in str(excinfo.value)  # "превысил 5с"


def test_chat_5xx_raises_lmstudio_unavailable():
    resp = _make_response(status_code=503, text="service unavailable")

    with patch("lm_client.httpx.Client") as MockClient:
        mock_inst = MagicMock()
        mock_inst.post.return_value = resp
        MockClient.return_value.__enter__.return_value = mock_inst

        with pytest.raises(LMStudioUnavailable) as excinfo:
            chat(
                messages=[{"role": "user", "content": "x"}],
                base_url="http://m:1/v1",
            )
    msg = str(excinfo.value)
    assert "503" in msg
    assert "service unavailable"[:30] in msg or "service unav" in msg


def test_chat_500_raises_lmstudio_unavailable():
    resp = _make_response(status_code=500, text="internal error")

    with patch("lm_client.httpx.Client") as MockClient:
        mock_inst = MagicMock()
        mock_inst.post.return_value = resp
        MockClient.return_value.__enter__.return_value = mock_inst

        with pytest.raises(LMStudioUnavailable) as excinfo:
            chat(
                messages=[{"role": "user", "content": "x"}],
                base_url="http://m:1/v1",
            )
    assert "500" in str(excinfo.value)


def test_chat_4xx_raises_lmstudio_unavailable():
    """4xx — наша ошибка, но клиент всё равно бросает LMStudioUnavailable."""
    resp = _make_response(status_code=401, text="unauthorized")

    with patch("lm_client.httpx.Client") as MockClient:
        mock_inst = MagicMock()
        mock_inst.post.return_value = resp
        MockClient.return_value.__enter__.return_value = mock_inst

        with pytest.raises(LMStudioUnavailable) as excinfo:
            chat(
                messages=[{"role": "user", "content": "x"}],
                base_url="http://m:1/v1",
            )
    assert "401" in str(excinfo.value)


def test_chat_network_error_raises_lmstudio_unavailable():
    """Любой httpx.HTTPError (ConnectionError и пр.) → LMStudioUnavailable."""
    with patch("lm_client.httpx.Client") as MockClient:
        mock_inst = MagicMock()
        mock_inst.post.side_effect = httpx.ConnectError("connection refused")
        MockClient.return_value.__enter__.return_value = mock_inst

        with pytest.raises(LMStudioUnavailable) as excinfo:
            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_chat_invalid_json_raises_lmstudio_unavailable():
    """200 OK, но тело — не JSON."""
    resp = _make_response(status_code=200, text="<html>not json</html>")

    with patch("lm_client.httpx.Client") as MockClient:
        mock_inst = MagicMock()
        mock_inst.post.return_value = resp
        MockClient.return_value.__enter__.return_value = mock_inst

        with pytest.raises(LMStudioUnavailable) as excinfo:
            chat(
                messages=[{"role": "user", "content": "x"}],
                base_url="http://m:1/v1",
            )
    assert "не-JSON" in str(excinfo.value) or "json" in str(excinfo.value).lower()


def test_chat_empty_choices_raises_lmstudio_unavailable():
    """200 OK, choices пустой → ошибка."""
    payload = {"choices": []}
    resp = _make_response(json_payload=payload)

    with patch("lm_client.httpx.Client") as MockClient:
        mock_inst = MagicMock()
        mock_inst.post.return_value = resp
        MockClient.return_value.__enter__.return_value = mock_inst

        with pytest.raises(LMStudioUnavailable) as excinfo:
            chat(
                messages=[{"role": "user", "content": "x"}],
                base_url="http://m:1/v1",
            )
    assert "choice" in str(excinfo.value).lower()


# ---------------------------------------------------------------------------
# Специфические кейсы: tool_use, контент None, несколько svg-блоков
# ---------------------------------------------------------------------------


def test_chat_tool_use_response_yields_empty_string_candidate():
    """Модель вернула tool_calls без content → raw_texts содержит '' для этого кандидата."""
    payload = {
        "choices": [
            {
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [
                        {
                            "id": "1",
                            "type": "function",
                            "function": {"name": "x", "arguments": "{}"},
                        }
                    ],
                },
                "finish_reason": "tool_calls",
            },
            {
                "index": 1,
                "message": {
                    "role": "assistant",
                    "content": "<svg viewBox='0 0 64 64'/>",
                },
                "finish_reason": "stop",
            },
        ]
    }
    resp = _make_response(json_payload=payload)

    with patch("lm_client.httpx.Client") as MockClient:
        mock_inst = MagicMock()
        mock_inst.post.return_value = resp
        MockClient.return_value.__enter__.return_value = mock_inst

        result = chat(
            messages=[{"role": "user", "content": "x"}],
            n=2,
            base_url="http://m:1/v1",
        )

    # 1-й кандидат — пусто, 2-й — реальный svg.
    assert result.raw_texts[0] == ""
    assert result.raw_texts[1] == "<svg viewBox='0 0 64 64'/>"
    assert result.finish_reasons[0] == "tool_calls"


def test_chat_response_with_multiple_svg_blocks_kept_as_raw_text():
    """Модель вернула ответ с несколькими ```svg блоками внутри — chat() должен
    сохранить текст as-is. Парсинг — забота validator.extract_svg()."""
    multi_svg_text = (
        "Here are some variants:\n"
        "```svg\n<svg viewBox='0 0 64 64'><circle cx='10' cy='10' r='5'/></svg>\n```\n"
        "And another:\n"
        "```svg\n<svg viewBox='0 0 64 64'><rect width='10' height='10'/></svg>\n```"
    )
    resp = _make_response(json_payload=_ok_payload([multi_svg_text] * 2))

    with patch("lm_client.httpx.Client") as MockClient:
        mock_inst = MagicMock()
        mock_inst.post.return_value = resp
        MockClient.return_value.__enter__.return_value = mock_inst

        result = chat(
            messages=[{"role": "user", "content": "x"}],
            n=2,
            base_url="http://m:1/v1",
        )

    # chat() НЕ парсит SVG — оба кандидата идентичны.
    assert len(result.raw_texts) == 2
    for raw in result.raw_texts:
        assert raw.count("<svg") == 2
        assert "Here are some variants" in raw


def test_chat_content_as_list_of_parts_joined():
    """content может быть list[dict] (мультимодальный ответ) — склеиваем текстовые части."""
    payload = {
        "choices": [
            {
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": [
                        {"type": "text", "text": "Hello "},
                        {"type": "text", "text": "world"},
                    ],
                },
                "finish_reason": "stop",
            }
        ]
    }
    resp = _make_response(json_payload=payload)

    with patch("lm_client.httpx.Client") as MockClient:
        mock_inst = MagicMock()
        mock_inst.post.return_value = resp
        MockClient.return_value.__enter__.return_value = mock_inst

        result = chat(
            messages=[{"role": "user", "content": "x"}],
            n=1,
            base_url="http://m:1/v1",
        )

    assert result.raw_texts == ["Hello world"]


# ---------------------------------------------------------------------------
# generate_svg() — высокоуровневая обёртка
# ---------------------------------------------------------------------------


def test_generate_svg_rejects_empty_prompt():
    with pytest.raises(ValueError, match="1 до 1000"):
        generate_svg("", model="x")


def test_generate_svg_rejects_too_long_prompt():
    with pytest.raises(ValueError, match="1 до 1000"):
        generate_svg("x" * 1001, model="x")


def test_generate_svg_with_image_sends_multimodal_payload():
    """С image_b64 идёт мультимодальный user-message: text + image_url."""
    svg_text = "<svg viewBox='0 0 64 64'/>"
    resp = _make_response(json_payload=_ok_payload([svg_text]))

    with patch("lm_client.chat") as mock_chat:
        mock_chat.return_value = type("R", (), {
            "raw_texts": [svg_text],
            "elapsed_s": 0.05,
            "model": "x",
            "usage": None,
            "finish_reasons": ["stop"],
        })()

        results = generate_svg(
            "a fox",
            image_b64="data:image/png;base64,AAA",
            mode="icon",
            n=1,
            temperature=0.4,
            model="x",
            system_prompt="<system>icon</system>",
        )

    assert len(results) == 1
    assert results[0]["svg_text"] == svg_text
    assert results[0]["raw"] == svg_text
    assert results[0]["finish_reason"] == "stop"

    # Что ушло в chat()
    call = mock_chat.call_args
    msgs = call.kwargs["messages"]
    assert msgs[0]["role"] == "system"
    assert msgs[0]["content"] == "<system>icon</system>"
    assert msgs[1]["role"] == "user"
    user_content = msgs[1]["content"]
    assert isinstance(user_content, list)
    assert user_content[0]["type"] == "text"
    assert user_content[0]["text"] == "a fox"
    assert user_content[1]["type"] == "image_url"
    assert user_content[1]["image_url"]["url"] == "data:image/png;base64,AAA"
    assert call.kwargs["model"] == "x"
    assert call.kwargs["n"] == 1
    assert call.kwargs["temperature"] == 0.4


def test_generate_svg_text_only_sends_string_content():
    """Без image_b64 user.content — просто строка."""
    with patch("lm_client.chat") as mock_chat:
        mock_chat.return_value = type("R", (), {
            "raw_texts": ["<svg/>"],
            "elapsed_s": 0.0,
            "model": "x",
            "usage": None,
            "finish_reasons": ["stop"],
        })()

        generate_svg("hello", n=1, model="x")

    msgs = mock_chat.call_args.kwargs["messages"]
    user_msg = msgs[0]  # без system_prompt единственное user-сообщение
    assert user_msg["role"] == "user"
    assert user_msg["content"] == "hello"


def test_generate_svg_propagates_lmstudio_error():
    """Если chat() упал, generate_svg пробрасывает LMStudioUnavailable."""
    with patch("lm_client.chat") as mock_chat:
        mock_chat.side_effect = LMStudioUnavailable("upstream timeout")

        with pytest.raises(LMStudioUnavailable, match="upstream timeout"):
            generate_svg("hello", n=1, model="x")


def test_chat_rejects_n_less_than_one():
    with pytest.raises(ValueError, match="n должно быть"):
        chat(messages=[{"role": "user", "content": "x"}], n=0, base_url="http://m:1/v1")


# ---------------------------------------------------------------------------
# encode_pil_to_data_url
# ---------------------------------------------------------------------------


def test_encode_pil_to_data_url_produces_data_url_with_png_mime():
    from PIL import Image

    img = Image.new("RGB", (10, 10), color=(255, 0, 0))
    url = encode_pil_to_data_url(img, mime="image/png")
    assert url.startswith("data:image/png;base64,")
    # base64 должен корректно декодироваться обратно
    payload = url.split(",", 1)[1]
    decoded = base64.b64decode(payload)
    assert decoded.startswith(b"\x89PNG")
    # и это валидный PNG
    Image.open(io.BytesIO(decoded))


def test_encode_pil_to_data_url_jpeg_mime_converts_rgba():
    """При mime=jpeg и RGBA → конвертируем в RGB (иначе JPEG не съест)."""
    from PIL import Image

    img = Image.new("RGBA", (8, 8), color=(0, 255, 0, 128))
    url = encode_pil_to_data_url(img, mime="image/jpeg")
    assert url.startswith("data:image/jpeg;base64,")
    payload = url.split(",", 1)[1]
    decoded = base64.b64decode(payload)
    # JPEG стартует с FFD8
    assert decoded.startswith(b"\xff\xd8")


def test_encode_pil_to_data_url_png_mime_adds_alpha_if_other_mode():
    """PNG с mime=png и mode=L (grayscale) → конвертируем в RGBA."""
    from PIL import Image

    img = Image.new("L", (4, 4), color=128)
    url = encode_pil_to_data_url(img, mime="image/png")
    assert url.startswith("data:image/png;base64,")
    # Просто проверим, что получили валидный PNG
    payload = url.split(",", 1)[1]
    decoded = base64.b64decode(payload)
    assert decoded.startswith(b"\x89PNG")


# ---------------------------------------------------------------------------
# validate_image
# ---------------------------------------------------------------------------


def test_validate_image_passes_normal_png():
    from PIL import Image

    img = Image.new("RGB", (100, 100), color=(0, 0, 0))
    validate_image(img)  # не бросает


def test_validate_image_rejects_oversized_dimensions():
    from PIL import Image

    img = Image.new("RGB", (5000, 100), color=(0, 0, 0))
    with pytest.raises(ValueError, match="слишком большое"):
        validate_image(img, max_side=4096)


def test_validate_image_rejects_too_many_bytes():
    """Картинка проходит по стороне, но approx-байты > max_bytes.

    Берём 4000x4000 (под max_side=4096) и max_bytes=1MB:
    4000*4000*4 + 64KB = 64MB+ >> 1MB → должно сработать байтовое ограничение.
    """
    from PIL import Image

    img = Image.new("RGB", (4000, 4000), color=(0, 0, 0))
    with pytest.raises(ValueError, match="МБ"):
        validate_image(img, max_bytes=1 * 1024 * 1024)


def test_validate_image_rejects_unsupported_format():
    from PIL import Image

    img = Image.new("RGB", (10, 10), color=(0, 0, 0))
    img.format = "BMP"  # притворяемся BMP
    with pytest.raises(ValueError, match="неподдерживаемый формат"):
        validate_image(img)


def test_validate_image_rejects_none():
    with pytest.raises(ValueError, match="не передано"):
        validate_image(None)  # type: ignore[arg-type]


def test_validate_image_rejects_non_pil():
    with pytest.raises(ValueError, match="ожидался PIL.Image"):
        validate_image("not-an-image")  # type: ignore[arg-type]


# ---------------------------------------------------------------------------
# Минорный: payload включает n/temperature/max_tokens/stream=False
# ---------------------------------------------------------------------------


def test_chat_payload_contains_required_fields():
    resp = _make_response(json_payload=_ok_payload(["x"]))
    with patch("lm_client.httpx.Client") as MockClient:
        mock_inst = MagicMock()
        mock_inst.post.return_value = resp
        MockClient.return_value.__enter__.return_value = mock_inst

        chat(
            messages=[{"role": "user", "content": "x"}],
            n=2,
            temperature=0.7,
            max_tokens=2048,
            base_url="http://m:1/v1",
        )

    body = mock_inst.post.call_args.kwargs["json"]
    assert body["n"] == 2
    assert body["temperature"] == 0.7
    assert body["max_tokens"] == 2048
    assert body["stream"] is False
    assert body["model"]  # non-empty


def test_chat_uses_default_base_url_when_env_unset(monkeypatch):
    """Если base_url=None и env не задан, идём на DEFAULT_BASE_URL."""
    monkeypatch.delenv("LM_STUDIO_BASE_URL", raising=False)
    resp = _make_response(json_payload=_ok_payload(["x"]))
    with patch("lm_client.httpx.Client") as MockClient:
        mock_inst = MagicMock()
        mock_inst.post.return_value = resp
        MockClient.return_value.__enter__.return_value = mock_inst

        chat(messages=[{"role": "user", "content": "x"}], n=1)

    url = mock_inst.post.call_args.args[0]
    assert url == f"{DEFAULT_BASE_URL}/chat/completions"