"""Юнит-тесты для 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("icon" 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": [""], "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"