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