LLeMbas/tests/conftest.py

"""Test fixtures.

Every test runs against a throwaway SQLite file in a tmp_path, never the real
data directory. The environment has to be set before lembas.config is imported,
because Settings is a cached singleton read at import time.
"""

from __future__ import annotations

import os
import tempfile
from collections.abc import Iterator
from pathlib import Path

import pytest

# Must happen before any lembas import.
_TMP = tempfile.mkdtemp(prefix="lembas-tests-")
os.environ.update(
    {
        "LEMBAS_SECRET_KEY": "test-secret-key-not-for-real-use",
        "LEMBAS_DATA_DIR": _TMP,
        "LEMBAS_ALLOW_SIGNUP": "true",
        "LEMBAS_LOG_LEVEL": "warning",
    }
)

from fastapi.testclient import TestClient  # noqa: E402
from sqlalchemy.orm import Session  # noqa: E402

from lembas.config import settings  # noqa: E402
from lembas.db.base import Base  # noqa: E402
from lembas.db.session import get_engine, get_session_factory, reset_engine  # noqa: E402


@pytest.fixture(autouse=True)
def fresh_database(tmp_path: Path) -> Iterator[None]:
    """Point the engine at a per-test database and build the schema.

    reset_engine() is essential: the engine is a module-level singleton, so
    without it every test after the first would share the first one's file.
    """
    settings.data_dir = tmp_path
    reset_engine()
    settings.ensure_dirs()

    import lembas.db.models  # noqa: F401  (registers the tables)

    # sync_schema rather than create_all: it is what startup runs, and it also
    # builds the full-text indexes, which are not SQLAlchemy models and so are
    # invisible to create_all. Tests were otherwise running against a schema
    # production does not have.
    from lembas.db.migrations import sync_schema

    Base.metadata.create_all(bind=get_engine())
    sync_schema(get_engine())
    yield
    reset_engine()


@pytest.fixture(autouse=True)
def fresh_generation_registry() -> Iterator[None]:
    """Empty the in-flight reply registry between tests.

    `_RUNNING` and `_TASKS` are module-level dicts, so a test that starts a
    reply and does not wait for it leaves an entry behind for the rest of the
    session -- holding a Generation, and a Task belonging to an event loop that
    has since closed. `_prune()` will not clear it either: it only drops
    generations that have finished, and it runs on every `ensure()`.

    Cheap, and it keeps a test that posts a message from meeting the leftovers
    of one that asked a question.
    """
    from lembas.services import generation as generation_service

    generation_service._RUNNING.clear()
    generation_service._TASKS.clear()
    yield
    generation_service._RUNNING.clear()
    generation_service._TASKS.clear()


@pytest.fixture(autouse=True)
def fresh_terminal_registry() -> Iterator[None]:
    """Empty the open-shell registry between tests, for the same reason.

    A leaked entry holds an asyncssh connection belonging to an event loop that
    has since closed, and the reaper task is module-level too -- one left
    running would wake up inside the next test's loop.
    """
    from lembas.services.agent import terminal as terminal_service

    def _clear() -> None:
        reaper = terminal_service._REAPER
        if reaper is not None:
            reaper.cancel()
        terminal_service._REAPER = None
        terminal_service._SESSIONS.clear()

    _clear()
    yield
    _clear()


@pytest.fixture
def db() -> Iterator[Session]:
    session = get_session_factory()()
    try:
        yield session
    finally:
        session.close()


@pytest.fixture
def client() -> Iterator[TestClient]:
    from lembas.main import app

    # raise_server_exceptions=False so error-handler behaviour is exercised
    # rather than the exception propagating into the test.
    with TestClient(app, raise_server_exceptions=False) as test_client:
        yield test_client


@pytest.fixture
def registered(client: TestClient) -> dict[str, str]:
    """Register the first account. It becomes the administrator."""
    credentials = {
        "name": "Frodo",
        "email": "frodo@shire.test",
        "password": "speak-friend-and-enter",
    }
    response = client.post("/auth/register", data=credentials, follow_redirects=False)
    assert response.status_code == 303, response.text
    return credentials


@pytest.fixture
def make_chat(db: Session):
    """Create a chat row directly, as scaffolding for other tests.

    Chats are normally created by POST /api/chats/start along with their first
    exchange -- there is deliberately no endpoint that makes an empty one. Most
    tests want a chat to act on, not that flow, so they get one straight from
    the database rather than having to subtract an opening turn from every
    assertion. The flow itself is covered in test_chat.py.
    """
    from sqlalchemy import select

    from lembas.db.models import Chat, Model, User

    def _create(email: str | None = None, model_id: str | None = None) -> str:
        user = (
            db.scalar(select(User).where(User.email == email))
            if email
            else db.scalars(select(User).order_by(User.created_at)).first()
        )
        model = (
            db.scalar(select(Model).where(Model.model_id == model_id))
            if model_id
            else db.scalars(select(Model).order_by(Model.position)).first()
        )
        chat = Chat(
            user_id=user.id,
            model_id=model.model_id if model else "",
            connection_id=model.connection_id if model else None,
        )
        db.add(chat)
        db.commit()
        return chat.id

    return _create


@pytest.fixture
def mock_http():
    """Answer every outgoing httpx request with a handler of the test's choosing.

    The services build their own AsyncClient because each needs its own timeout,
    so there is no client to inject; patching the class is what reaches them.
    Returns a callable that installs a handler and is undone on teardown.
    """
    import httpx

    original = httpx.AsyncClient

    def install(handler):
        class Patched(original):
            def __init__(self, **kwargs):
                super().__init__(transport=httpx.MockTransport(handler), **kwargs)

        httpx.AsyncClient = Patched

    yield install
    httpx.AsyncClient = original


@pytest.fixture
def user_id(db: Session, registered: dict[str, str]) -> str:
    """The registered user's id.

    Chats have a real foreign key to users and SQLite enforces it (the
    connect-time PRAGMA in db/session.py turns that on), so tests that build a
    Chat directly need a user that actually exists.
    """
    from sqlalchemy import select

    from lembas.db.models import User

    return db.scalar(select(User).where(User.email == registered["email"])).id