arena-dashboard/tests/conftest.py

"""Общие фикстуры: реальный мок-сервер OmniRoute на localhost + чистые модули."""
from __future__ import annotations

import json
import os
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from types import SimpleNamespace

import pytest
from fastapi.testclient import TestClient

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
APP_DIR = os.path.join(REPO, "app")
if APP_DIR not in sys.path:
    sys.path.insert(0, APP_DIR)

MODULES = ("main", "probe", "usage_utils")
SAVED_ENV = {}


class OmniMock:
    """Управляемый мок OmniRoute: реальные HTTP-запросы на 127.0.0.1."""

    def __init__(self):
        self.state = {
            # поведение /v1/chat/completions по умолчанию; переопределяется
            # per-model: status, usage, no_usage, omit_choices, content, delay
            "completions": {"status": 200, "usage": None, "content": "4"},
            "models": {"data": [
                {"id": "zai/test-model-a"},
                {"id": "agy/test-model-b"},
            ]},
            "global_delay": 0.0,
            "per_model": {},
        }
        self.requests = []
        self.lock = threading.Lock()

        outer = self

        class Handler(BaseHTTPRequestHandler):
            def log_message(self, *a):
                pass

            def _send_json(self, status, payload):
                data = json.dumps(payload).encode()
                try:
                    self.send_response(status)
                    self.send_header("Content-Type", "application/json")
                    self.send_header("Content-Length", str(len(data)))
                    self.end_headers()
                    self.wfile.write(data)
                except (BrokenPipeError, ConnectionResetError):
                    pass  # клиент отвалился по таймауту — это ожидаемо

            def _read_body(self):
                n = int(self.headers.get("Content-Length") or 0)
                raw = self.rfile.read(n) if n else b""
                try:
                    return json.loads(raw)
                except ValueError:
                    return {}

            def do_POST(self):
                body = self._read_body()
                with outer.lock:
                    outer.requests.append({"path": self.path, "body": body})
                    beh = dict(outer.state["completions"])
                    beh.update(outer.state["per_model"].get(body.get("model"), {}))
                    delay = beh.get("delay", 0.0) or outer.state["global_delay"]
                if delay:
                    time.sleep(delay)
                if beh.get("status", 200) != 200:
                    self._send_json(beh["status"], {"error": {"message": "mock error"}})
                    return
                resp = {"id": "chatcmpl-mock", "object": "chat.completion",
                        "created": int(time.time()), "model": body.get("model", "?")}
                if not beh.get("no_usage"):
                    resp["usage"] = beh.get("usage")
                if not beh.get("omit_choices"):
                    resp["choices"] = [{"index": 0, "finish_reason": beh.get("finish_reason", "stop"),
                                        "message": {"role": "assistant",
                                                    "content": beh.get("content", "4"),
                                                    "reasoning_content": beh.get("reasoning_content")}}]
                self._send_json(200, resp)

            def do_GET(self):
                if self.path == "/v1/models":
                    with outer.lock:
                        payload = dict(outer.state["models"])
                    self._send_json(200, payload)
                    return
                self._send_json(404, {"detail": "not found"})

        self.server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
        self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
        self.thread.start()
        self.url = f"http://127.0.0.1:{self.server.server_address[1]}"

    def set_completions(self, **kw):
        with self.lock:
            self.state["completions"].update(kw)

    def set_per_model(self, model_id, beh):
        with self.lock:
            self.state["per_model"][model_id] = beh

    def close(self):
        self.server.shutdown()
        self.server.server_close()
        self.thread.join(timeout=5)


@pytest.fixture(scope="session")
def omni_server():
    mock = OmniMock()
    yield mock
    mock.close()


@pytest.fixture(autouse=True)
def _reset_omni(omni_server):
    """Изолировать тесты: состояние мока не должно протекать между тестами."""
    with omni_server.lock:
        omni_server.state["per_model"] = {}
        omni_server.state["models"] = {"data": [
            {"id": "zai/test-model-a"}, {"id": "agy/test-model-b"},
        ]}
        omni_server.state["global_delay"] = 0.0
        omni_server.state["completions"] = {"status": 200, "usage": None, "content": "4"}
        omni_server.requests.clear()
    yield


def _set_env(db_path: str, omni_url: str):
    env = {
        "DB_PATH": db_path,
        "OMNIROUTE_URL": omni_url,
        "OMNIROUTE_API_KEY": "test-key",
        "CHECK_TIMEOUT": "2",
        "CHECK_CONCURRENCY": "4",
        "ARENA_CACHE_TEXT": "/nonexistent/arena_text.json",
        "ARENA_CACHE_AGENT": "/nonexistent/arena_agent.json",
    }
    for k, v in env.items():
        SAVED_ENV[k] = os.environ.get(k)
        os.environ[k] = v


def restore_env():
    for k, v in SAVED_ENV.items():
        if v is None:
            os.environ.pop(k, None)
        else:
            os.environ[k] = v


@pytest.fixture
def fresh_main(tmp_path, omni_server):
    """Модуль main с env, привязанным к tmp-БД и мок-серверу (без запуска app).

    init_db НЕ вызывается — тест миграции сам решает, что лежит в БД до него.
    """
    db_path = str(tmp_path / "arena.db")
    _set_env(db_path, omni_server.url)
    for mod in MODULES:
        sys.modules.pop(mod, None)
    import main as m
    try:
        yield m
    finally:
        restore_env()


@pytest.fixture
def core(fresh_main):
    """Приложение с запущенным lifespan (init_db + sync_models) + модуль probe."""
    import probe as probe_mod
    with TestClient(fresh_main.app) as client:
        yield SimpleNamespace(main=fresh_main, probe=probe_mod,
                              client=client, db_path=fresh_main.DB_PATH)