"""Общие фикстуры: реальный мок-сервер 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": "stop", "message": {"role": "assistant", "content": beh.get("content", "4")}}] 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["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)