talemate/tests/conftest.py

"""
Shared pytest fixtures and test infrastructure.

Provides MockClient, MockScene, and bootstrap functions used across
multiple test modules (test_graphs, test_layered_history, etc.).
"""

import contextvars
from collections import deque
from pathlib import Path

import pytest
import yaml

import talemate.agents as agents
import talemate.agents.memory
import talemate.agents.tts.voice_library as voice_library
import talemate.config.state as config_state
import talemate.emit.async_signals as async_signals
import talemate.instance as instance
from talemate.client import ClientBase
from talemate.config.schema import Config
from talemate.tale_mate import Scene

# Root of the repository (where config.example.yaml lives)
_REPO_ROOT = Path(__file__).resolve().parent.parent


@pytest.fixture(autouse=True, scope="session")
def _use_example_config():
    """Ensure all tests use config.example.yaml instead of the local config.yaml.

    This prevents local configuration from leaking into test results and
    keeps CI and local runs deterministic.
    """
    example_path = _REPO_ROOT / "config.example.yaml"
    with open(example_path, "r") as f:
        yaml_data = yaml.safe_load(f) or {}
    test_config = Config.model_validate(yaml_data)

    original = config_state.CONFIG
    config_state.CONFIG = test_config
    yield
    config_state.CONFIG = original


# ---------------------------------------------------------------------------
# Contextvar-based response queue for MockClient
# ---------------------------------------------------------------------------

client_responses = contextvars.ContextVar("client_responses", default=deque())


class MockClientContext:
    """Async context manager that provides a fresh response queue."""

    async def __aenter__(self):
        try:
            self.client_responses = client_responses.get()
        except LookupError:
            _client_responses = deque()
            self.token = client_responses.set(_client_responses)
            self.client_responses = _client_responses

        return self.client_responses

    async def __aexit__(self, exc_type, exc_value, traceback):
        if hasattr(self, "token"):
            client_responses.reset(self.token)


# ---------------------------------------------------------------------------
# Async signal helpers
# ---------------------------------------------------------------------------


@pytest.fixture
def isolate_signals():
    """Factory that clears a signal's receivers for the duration of the test
    (so handlers don't leak between tests) and restores them on teardown.
    Returns the isolated AsyncSignal objects for connecting test handlers."""
    restores = []

    def _isolate(*names):
        signals = []
        for name in names:
            sig = async_signals.get(name)
            restores.append((sig, list(sig.receivers)))
            sig.receivers.clear()
            signals.append(sig)
        return signals[0] if len(signals) == 1 else signals

    yield _isolate

    for sig, receivers in restores:
        sig.receivers.clear()
        sig.receivers.extend(receivers)


def connect_recorder(signal) -> list:
    """Connect a recording handler to a signal and return the list that
    received payloads are appended to."""
    received = []

    async def handler(payload):
        received.append(payload)

    signal.connect(handler)
    return received


# ---------------------------------------------------------------------------
# Mock classes
# ---------------------------------------------------------------------------


class MockClient(ClientBase):
    """LLM client stub that pops pre-defined responses from a queue."""

    def __init__(self, name: str):
        self.name = name
        self.remote_model_name = "test-model"
        self.current_status = "idle"
        self.prompt_history = []

    @property
    def enabled(self):
        return True

    async def send_prompt(
        self, prompt, kind="conversation", finalize=lambda x: x, retries=2, **kwargs
    ):
        response_stack = client_responses.get()
        self.prompt_history.append({"prompt": prompt, "kind": kind})
        if not response_stack:
            return ""
        return response_stack.popleft()


class MockMemoryAgent(talemate.agents.memory.MemoryAgent):
    """MemoryAgent with no-op persistence methods."""

    async def add_many(self, items: list[dict]):
        pass

    async def delete(self, filters: dict):
        pass


class MockScene(Scene):
    """Real Scene subclass with auto_progress forced on."""

    @property
    def auto_progress(self):
        return True


# ---------------------------------------------------------------------------
# Bootstrap helpers
# ---------------------------------------------------------------------------


def bootstrap_engine():
    """Instantiate all real agents (using MockMemoryAgent for memory)."""
    voice_library.VOICE_LIBRARY = voice_library.VoiceLibrary(voices={})
    for agent_type in agents.AGENT_CLASSES:
        if agent_type == "memory":
            agent = MockMemoryAgent()
        else:
            agent = agents.AGENT_CLASSES[agent_type]()
        instance.AGENTS[agent_type] = agent


def pytest_addoption(parser):
    """Add custom command-line options."""
    parser.addoption(
        "--update-baselines",
        action="store_true",
        default=False,
        help="Update baseline snapshot files instead of comparing against them.",
    )


@pytest.fixture
def update_baselines(request):
    """Whether to update baseline files instead of comparing."""
    return request.config.getoption("--update-baselines")


def bootstrap_scene(mock_scene):
    """Wire a MockClient and the mock_scene into every agent."""
    bootstrap_engine()
    client = MockClient("test_client")
    for agent in instance.AGENTS.values():
        agent.client = client
        agent.scene = mock_scene

    director = instance.get_agent("director")
    conversation = instance.get_agent("conversation")
    summarizer = instance.get_agent("summarizer")
    editor = instance.get_agent("editor")
    world_state = instance.get_agent("world_state")

    mock_scene.mock_client = client

    return {
        "director": director,
        "conversation": conversation,
        "summarizer": summarizer,
        "editor": editor,
        "world_state": world_state,
    }