haiku.rag/tests/conftest.py

import json
import logging
import os
import tempfile
from collections.abc import Iterator
from contextlib import contextmanager
from pathlib import Path
from typing import TYPE_CHECKING, Any

# Prevent tests from loading user's local haiku.rag.yaml by setting env var
# to a test config file BEFORE any haiku.rag imports.
# Uses Ollama for embeddings - HTTP calls are recorded/replayed via VCR.
_test_config_dir = tempfile.mkdtemp()
_test_config_path = Path(_test_config_dir) / "test-defaults.yaml"
_test_config_path.write_text("""
embeddings:
  model:
    provider: ollama
    name: qwen3-embedding:4b
    vector_dim: 2560
""")
os.environ["HAIKU_RAG_CONFIG_PATH"] = str(_test_config_path)

import pydantic_ai.models  # noqa: E402
import pytest  # noqa: E402
import yaml  # noqa: E402

from .services import reachable  # noqa: E402

if TYPE_CHECKING:
    from vcr import VCR

    from haiku.rag.client import HaikuRAG
    from haiku.rag.client.scope import DatabaseScope
    from haiku.rag.client.session import SingleDatabaseSession
    from haiku.rag.config.models import AppConfig

setattr(pydantic_ai.models, "ALLOW_MODEL_REQUESTS", False)
logging.getLogger("vcr.cassette").setLevel(logging.WARNING)


@contextmanager
def capture_logs(
    logger: logging.Logger, level: int
) -> Iterator[list[logging.LogRecord]]:
    """Collect records emitted by ``logger`` at or above ``level``.

    Attaches directly to the given logger instead of using ``caplog``:
    ``haiku.rag.logging.get_logger()`` sets ``propagate=False`` on the
    ``haiku.rag`` logger, so records never reach caplog's root handler once
    any test in the session has called it.
    """
    records: list[logging.LogRecord] = []

    class _ListHandler(logging.Handler):
        def emit(self, record: logging.LogRecord) -> None:
            records.append(record)

    handler = _ListHandler(level=level)
    logger.addHandler(handler)
    try:
        yield records
    finally:
        logger.removeHandler(handler)


@pytest.fixture(scope="session")
def qa_corpus() -> list[dict[str, str]]:
    corpus_path = Path(__file__).parent / "data" / "qa_corpus.json"
    with open(corpus_path) as f:
        return json.load(f)


@pytest.fixture
def temp_db_path(tmp_path):
    """Create a temporary database path for testing.

    Note: Tests that need a database should use HaikuRAG with create=True.
    """
    return tmp_path / "test.lancedb"


@pytest.fixture
def temp_yaml_config(tmp_path, monkeypatch):
    """Create a temporary YAML config file for testing.

    This fixture creates a config file in a temp directory and sets
    the environment variable so config.py will load it.
    """
    config_file = tmp_path / "test-config.yaml"
    config_data = {
        "environment": "production",
        "storage": {
            "data_dir": "",
            "monitor_directories": [],
            "vacuum_retention_seconds": 60,
        },
        "embeddings": {
            "model": {
                "provider": "ollama",
                "name": "qwen3-embedding:4b",
                "vector_dim": 2560,
            }
        },
        "qa": {"model": {"provider": "ollama", "name": "qwen3.8"}},
    }

    with open(config_file, "w") as f:
        yaml.dump(config_data, f)

    # Set env var so config loader will find it
    monkeypatch.setenv("HAIKU_RAG_CONFIG_PATH", str(config_file))

    yield config_file


@pytest.fixture
def allow_model_requests():
    with pydantic_ai.models.override_allow_model_requests(True):
        yield


@pytest.fixture(autouse=True)
def set_mock_api_keys(monkeypatch):
    """Set mock API keys for providers that require them during initialization."""
    if not os.getenv("OPENAI_API_KEY"):
        monkeypatch.setenv("OPENAI_API_KEY", "sk-mock-key-for-vcr-playback")
    if not os.getenv("ANTHROPIC_API_KEY"):
        monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-mock-key-for-vcr-playback")
    if not os.getenv("CO_API_KEY"):
        monkeypatch.setenv("CO_API_KEY", "mock-cohere-key-for-vcr-playback")
    if not os.getenv("ZEROENTROPY_API_KEY"):
        monkeypatch.setenv("ZEROENTROPY_API_KEY", "mock-ze-key-for-vcr-playback")
    if not os.getenv("VOYAGE_API_KEY"):
        monkeypatch.setenv("VOYAGE_API_KEY", "mock-voyage-key-for-vcr-playback")
    if not os.getenv("GROQ_API_KEY"):
        monkeypatch.setenv("GROQ_API_KEY", "mock-groq-key-for-vcr-playback")
    if not os.getenv("GOOGLE_API_KEY"):
        monkeypatch.setenv("GOOGLE_API_KEY", "mock-google-key-for-vcr-playback")
    if not os.getenv("AWS_DEFAULT_REGION"):
        monkeypatch.setenv("AWS_DEFAULT_REGION", "us-east-1")


def pytest_recording_configure(config: Any, vcr: "VCR"):
    from . import json_body_serializer

    vcr.register_serializer("yaml", json_body_serializer)


@pytest.fixture(scope="module")
def vcr_config():
    return {
        "ignore_localhost": False,
        "ignore_hosts": ["huggingface.co"],
        "filter_headers": ["authorization", "x-api-key"],
        "decode_compressed_response": True,
    }


@pytest.fixture(scope="session")
def doclaynet_first_page_pdf(tmp_path_factory) -> Path:
    """One-page extract of ``tests/data/doclaynet.pdf`` (the full DocLayNet
    arXiv paper). Most existing tests only need a small PDF with at least
    one picture; this avoids running docling over all nine pages of the
    paper just to assert ``pictures != []``. The full paper is used
    directly by the split-and-merge integration test."""
    import pypdfium2 as pdfium

    src_path = Path(__file__).parent / "data" / "doclaynet.pdf"
    out_dir = tmp_path_factory.mktemp("doclaynet")
    out_path = out_dir / "page0.pdf"

    src = pdfium.PdfDocument(str(src_path))
    try:
        dst = pdfium.PdfDocument.new()
        try:
            dst.import_pages(src, [0])
            with open(out_path, "wb") as f:
                dst.save(f)
        finally:
            dst.close()
    finally:
        src.close()
    return out_path


# --- external services for integration tests ---
#
# Integration tests (marked `integration`, excluded in CI via `-m "not
# integration"`) need external services. Bring them up with
#   docker compose -f tests/docker/docker-compose.yml up -d
# These fixtures hand the test a connection URL when the service is reachable
# (or when the matching env var points at an external instance), and skip the
# test otherwise.

_COMPOSE_HINT = (
    "start it with `docker compose -f tests/docker/docker-compose.yml up -d`"
)


@pytest.fixture(scope="session")
def postgres_dburi() -> str:
    """A reachable Postgres queue URL. Uses HAIKU_RAG_TEST_PG_DBURI when set,
    otherwise the docker-compose `postgres` service. Skips when neither is up."""
    override = os.environ.get("HAIKU_RAG_TEST_PG_DBURI")
    if override:
        return override
    if not reachable("localhost", 55432):
        pytest.skip(f"Postgres not reachable on localhost:55432 — {_COMPOSE_HINT}")
    return "postgresql+asyncpg://haiku:secret@localhost:55432/haiku_rag_test"


@pytest.fixture(scope="session")
def docling_serve_url() -> str:
    """A reachable docling-serve base URL. Uses HAIKU_RAG_TEST_DOCLING_SERVE_URL
    when set, otherwise the docker-compose `docling-serve` service. Skips when
    neither is up."""
    override = os.environ.get("HAIKU_RAG_TEST_DOCLING_SERVE_URL")
    if override:
        return override
    if not reachable("localhost", 5001):
        pytest.skip(f"docling-serve not reachable on localhost:5001 — {_COMPOSE_HINT}")
    return "http://localhost:5001"


def writing(client: "HaikuRAG") -> "SingleDatabaseSession":
    """The database a write implementation works on, from a client holding one.

    Write implementations take a session, never a client, so a set can
    never reach them. Tests that call one directly go through here."""
    from haiku.rag.client.session import SingleDatabaseSession

    assert isinstance(client._session, SingleDatabaseSession)
    return client._session


def for_path(
    db_path: "Path | str | None" = None, config: "AppConfig | None" = None
) -> "DatabaseScope":
    """A scope covering one database at `db_path`.

    The application layer takes the databases it works on, already resolved.
    Tests that hold a path and need a scope go through here.
    """
    from haiku.rag.client.scope import DatabaseScope
    from haiku.rag.config import get_config

    return DatabaseScope.resolve(
        config if config is not None else get_config(), database_path=db_path
    )


@contextmanager
def _covering_returns(stub, client):
    """Make a patched `HaikuRAG` hand back `client` however it is constructed.

    The TUIs build their client through `HaikuRAG._covering`, so patching the
    constructor alone leaves `_covering` answering with a fresh Mock.
    """
    stub.return_value = client
    stub._covering.return_value = client
    yield stub