d367b1eb5a
The CLI decides only what it knows — that --db and --database are the same thing said twice, and whether a command reads more than one — and hands the resolved scope down. Nothing rewrites the configuration, so a named database keeps the name results and citations carry, and a remote one opens the URI it was configured with rather than the local path standing in for it. HaikuRAGApp, ChatApp and InspectorApp take that scope and nothing else. Selection reaches the client through a private constructor, so the public signature still takes a path or names.
262 lines
8.7 KiB
Python
262 lines
8.7 KiB
Python
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": "gpt-oss"}},
|
|
}
|
|
|
|
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 rather than 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 rather than 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
|