"""Юнит-тесты для 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="not json")
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": "",
},
"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] == ""
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\n```\n"
"And another:\n"
"```svg\n\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("