omnisvg-lite/scripts/smoke_test.py

"""Smoke-test реального пути данных через локальный mock-LM-Server.

Поднимает HTTP-сервер на 127.0.0.1:<port> (по умолчанию 1234), который
принимает POST /v1/chat/completions и возвращает фиксированный валидный
SVG-ответ в формате OpenAI chat completions. Затем:

1) гоняет lm_client.generate_svg() через этот mock;
2) валидирует результат через validator.validate_svg();
3) рендерит PNG через renderer.render_png() (если cairo есть);
4) пишет запись в history через History() (SQLite в tmp).

Печатает PASS или FAIL с подробностями. Код выхода:
  0 — PASS
  1 — FAIL

Запуск:
    python scripts/smoke_test.py

Переменные окружения (опц.):
    SMOKE_PORT — порт mock-сервера (default 1234; будет увеличен, если занят)
    OMNISVG_DB_PATH — путь к SQLite (default в tmp)
    OMNISVG_PREVIEW_DIR — каталог превью (default в tmp)

ВАЖНО: скрипт НЕ зависит от запущенного LM Studio. Это именно mock-сервер,
который имитирует LM Studio достаточно, чтобы пройти через весь pipeline.
"""

from __future__ import annotations

import json
import os
import socket
import sys
import tempfile
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

# Добавляем корень проекта в sys.path
ROOT = Path(__file__).resolve().parent.parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))


# ---------------------------------------------------------------------------
# Mock LM Studio
# ---------------------------------------------------------------------------

# Валидный SVG, который mock всегда возвращает. Иконка 64x64, простая.
FIXTURE_SVG = (
    '<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 64 64" width="64" height="64">'
    '<circle cx="32" cy="32" r="20" fill="#3B82F6"/>'
    '<rect x="36" y="36" width="6" height="20" rx="3" fill="#3B82F6" '
    'transform="rotate(45 39 46)"/>'
    '</svg>'
)


class MockLMHandler(BaseHTTPRequestHandler):
    """Отвечает 200 OK + OpenAI-style JSON на /v1/chat/completions.

    Возвращает ровно `n` choice'ов (из тела запроса), все с FIXTURE_SVG.
    Если `n` не прислали — 1.
    """

    def log_message(self, format, *args):  # noqa: A002 — шумно молчим
        return  # без stdout-спама

    def do_POST(self):  # noqa: N802 — http.server API
        length = int(self.headers.get("Content-Length", "0") or "0")
        body_bytes = self.rfile.read(length) if length > 0 else b""

        if self.path.endswith("/chat/completions") or self.path == "/v1/chat/completions":
            # Попробуем достать n из payload (иначе 1)
            n = 1
            try:
                req = json.loads(body_bytes.decode("utf-8"))
                n = int(req.get("n", 1))
                if n < 1:
                    n = 1
            except (ValueError, json.JSONDecodeError):
                pass

            choices = [
                {
                    "index": i,
                    "message": {"role": "assistant", "content": FIXTURE_SVG},
                    "finish_reason": "stop",
                }
                for i in range(n)
            ]
            payload = {
                "id": "chatcmpl-smoke",
                "object": "chat.completion",
                "created": int(time.time()),
                "model": "smoke-mock",
                "choices": choices,
                "usage": {
                    "prompt_tokens": 10,
                    "completion_tokens": 5 * n,
                    "total_tokens": 10 + 5 * n,
                },
            }
            body = json.dumps(payload).encode("utf-8")
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
            return

        # Любой другой путь — 404
        self.send_response(404)
        self.send_header("Content-Type", "application/json")
        self.end_headers()
        self.wfile.write(b'{"error":"not found"}')


def _find_free_port(preferred: int) -> int:
    """Возвращает preferred, если свободен; иначе 0 (= ОС выберет)."""
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        try:
            s.bind(("127.0.0.1", preferred))
            return preferred
        except OSError:
            return 0


def _start_mock_server(port: int) -> tuple[ThreadingHTTPServer, int]:
    """Поднимает mock-LM и возвращает (server, actual_port)."""
    server = ThreadingHTTPServer(("127.0.0.1", port), MockLMHandler)
    actual_port = server.server_address[1]
    t = threading.Thread(target=server.serve_forever, name="mock-lm", daemon=True)
    t.start()
    return server, actual_port


# ---------------------------------------------------------------------------
# Smoke-логика
# ---------------------------------------------------------------------------


def _check(step: str, cond: bool, detail: str = "") -> bool:
    """Печатает шаг и его результат. Возвращает cond для chaining."""
    status = "OK" if cond else "FAIL"
    print(f"  [{status}] {step}" + (f" — {detail}" if detail else ""))
    return cond


def main() -> int:
    print("=" * 60)
    print("OmniSVG-Lite smoke test (mock LM Studio)")
    print("=" * 60)

    # tmp-каталоги для изоляции
    tmpdir = Path(tempfile.mkdtemp(prefix="omnisvg_smoke_"))
    db_path = tmpdir / "smoke.sqlite"
    previews_dir = tmpdir / "previews"
    previews_dir.mkdir(parents=True, exist_ok=True)

    # env-инжекция, чтобы lm_client/History/app пошли в наши пути
    os.environ["OMNISVG_DB_PATH"] = str(db_path)
    os.environ["OMNISVG_PREVIEW_DIR"] = str(previews_dir)

    # Поднимаем mock-сервер
    preferred = int(os.environ.get("SMOKE_PORT", "1234"))
    port = _find_free_port(preferred)
    server, actual_port = _start_mock_server(port)
    base_url = f"http://127.0.0.1:{actual_port}/v1"
    print(f"  mock server up at {base_url} (preferred port={preferred})")
    all_ok = True
    all_ok &= _check("mock server started", server is not None)

    # Дать серверу подняться
    time.sleep(0.05)

    results: list[dict] = []
    validated_svgs: list[str] = []
    previews: list[str] = []

    # 1) LM client: generate_svg()
    try:
        from lm_client import generate_svg, chat

        # Сначала сходим напрямую через chat() — чтобы видеть transport
        result = chat(
            messages=[{"role": "user", "content": "a fox"}],
            n=1,
            temperature=0.4,
            base_url=base_url,
            timeout_s=5.0,
        )
        all_ok &= _check(
            "chat() returned LMTurnResult",
            len(result.raw_texts) == 1,
            f"raw_texts={result.raw_texts}",
        )
        all_ok &= _check(
            "chat() payload contains expected SVG",
            FIXTURE_SVG in result.raw_texts[0],
            f"first 60 chars: {result.raw_texts[0][:60]!r}",
        )

        # Теперь — высокоуровневый generate_svg()
        results = generate_svg(
            "a fox",
            n=2,
            temperature=0.4,
            model="smoke-mock",
            system_prompt="<system>icon</system>",
        )
        all_ok &= _check(
            "generate_svg() returned 2 candidates",
            len(results) == 2,
            f"got {len(results)}",
        )
        all_ok &= _check(
            "each candidate has 'svg_text' and 'finish_reason'",
            all("svg_text" in r and "finish_reason" in r for r in results),
        )
    except Exception as exc:  # noqa: BLE001
        all_ok = False
        print(f"  [FAIL] lm_client raised: {type(exc).__name__}: {exc}")

    # 2) Validator
    try:
        from validator import validate_svg

        for r in results:
            ok, reason, cleaned = validate_svg(r["svg_text"], mode="icon")
            all_ok &= _check(
                f"validate_svg(icon) ok={ok}",
                ok,
                f"reason={reason!r}, cleaned_len={len(cleaned)}",
            )
            if ok:
                validated_svgs.append(cleaned)
    except Exception as exc:  # noqa: BLE001
        all_ok = False
        print(f"  [FAIL] validator raised: {type(exc).__name__}: {exc}")

    # 3) Renderer (graceful, если cairo недоступен)
    try:
        from renderer import render_png, save_png

        for i, svg in enumerate(validated_svgs):
            png = render_png(svg, size=(64, 64))
            if png is None:
                # cairo недоступен — это допустимо, просто отметим.
                print(f"  [SKIP] render_png(#{i}) → None (cairo unavailable)")
                continue
            path = save_png(
                png,
                previews_dir=previews_dir,
                record_id=0,
                candidate_index=i,
            )
            previews.append(str(path))
            all_ok &= _check(
                f"render_png(#{i}) → {path.name}",
                path.exists() and path.stat().st_size > 0,
            )
    except Exception as exc:  # noqa: BLE001
        all_ok = False
        print(f"  [FAIL] renderer raised: {type(exc).__name__}: {exc}")

    # 4) History
    try:
        from history import History, Record

        with History(db_path) as h:
            record_id = h.add(
                Record(
                    prompt="a fox",
                    mode="icon",
                    model="smoke-mock",
                    n_requested=2,
                    n_returned=len(results),
                    temperature=0.4,
                    status="ok" if validated_svgs else "all_invalid",
                    error_reason=None,
                    raw_outputs=[r["raw"] for r in results],
                    validated_outputs=validated_svgs,
                    previews=previews,
                    best_index=0 if validated_svgs else None,
                )
            )
        all_ok &= _check(
            "History.add() returned id",
            isinstance(record_id, int) and record_id > 0,
            f"id={record_id}",
        )

        with History(db_path) as h:
            rec = h.get(record_id)
        all_ok &= _check(
            "History.get() roundtrip",
            rec is not None and rec["prompt"] == "a fox",
            f"prompt={rec['prompt'] if rec else None!r}",
        )
    except Exception as exc:  # noqa: BLE001
        all_ok = False
        print(f"  [FAIL] history raised: {type(exc).__name__}: {exc}")

    # 5) Cleanup
    server.shutdown()
    server.server_close()

    print("=" * 60)
    if all_ok:
        print("PASS — все шаги успешны")
        return 0
    print("FAIL — см. подробности выше")
    return 1


if __name__ == "__main__":
    sys.exit(main())