agent-platform/tests/conftest.py

from __future__ import annotations

import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any

REPO_ROOT = Path(__file__).resolve().parents[1]


@dataclass(frozen=True)
class ServiceImportConfig:
    service_name: str
    libs: tuple[str, ...]


SERVICE_IMPORT_CONFIGS: dict[str, ServiceImportConfig] = {
    "agent-service": ServiceImportConfig(
        service_name="agent-service",
        libs=("core-domain", "core-shared", "core-db", "core-events")),
    "auth-service": ServiceImportConfig(
        service_name="auth-service",
        libs=("core-domain", "core-shared", "core-db")),
    "knowledge-service": ServiceImportConfig(
        service_name="knowledge-service",
        libs=("core-domain", "core-shared", "core-db")),
    "model-gateway-service": ServiceImportConfig(
        service_name="model-gateway-service",
        libs=("core-domain", "core-shared", "core-db")),
    "runtime-service": ServiceImportConfig(
        service_name="runtime-service",
        libs=("core-domain", "core-shared", "core-db", "core-events", "core-dsl")),
    "tool-service": ServiceImportConfig(
        service_name="tool-service",
        libs=("core-domain", "core-shared", "core-db")),
    "team-service": ServiceImportConfig(
        service_name="team-service",
        libs=("core-domain", "core-shared", "core-db", "core-events")),
    "workflow-service": ServiceImportConfig(
        service_name="workflow-service",
        libs=("core-domain", "core-shared", "core-db", "core-dsl")),
}


def prepare_service_import(
    service_name: str,
    *,
    libs: tuple[str, ...]) -> None:
    for module_name in list(sys.modules):
        if module_name == "app" or module_name.startswith("app."):
            del sys.modules[module_name]
    _clear_shared_sqlalchemy_metadata()

    for lib_name in libs:
        lib_path = REPO_ROOT / "libs" / lib_name / "src"
        _prepend_sys_path(lib_path)
    _prepend_sys_path(REPO_ROOT / "services" / service_name)


def prepare_known_service_import(service_name: str) -> None:
    config = SERVICE_IMPORT_CONFIGS[service_name]
    prepare_service_import(config.service_name, libs=config.libs)


def build_sqlite_database_url(tmp_path: Path, filename: str) -> str:
    return f"sqlite:///{tmp_path / filename}"


def _prepend_sys_path(path: Path) -> None:
    path_text = str(path)
    if path_text in sys.path:
        sys.path.remove(path_text)
    sys.path.insert(0, path_text)


def _clear_shared_sqlalchemy_metadata() -> None:
    try:
        from core_db import Base
    except ImportError:
        return
    Base.registry.dispose()
    Base.metadata.clear()


def build_fastapi_test_client(app: Any) -> Any:
    _patch_httpx_testclient_compatibility()
    from fastapi.testclient import TestClient

    return TestClient(app)


def _patch_httpx_testclient_compatibility() -> None:
    import inspect

    import httpx

    if "app" in inspect.signature(httpx.Client.__init__).parameters:
        return
    if getattr(httpx.Client.__init__, "_agent_platform_patched", False):
        return

    original_init = httpx.Client.__init__

    def patched_init(self: httpx.Client, *args: Any, **kwargs: Any) -> None:
        kwargs.pop("app", None)
        original_init(self, *args, **kwargs)

    patched_init._agent_platform_patched = True
    httpx.Client.__init__ = patched_init