haiku.rag/tests/ingester/test_source_plugins.py

import hashlib
from datetime import UTC, datetime

import pytest

from haiku.rag.config import PluginSourceConfig
from haiku.rag.ingester.pollers.factory import build_source
from haiku.rag.ingester.pollers.periodic import PeriodicPoller
from haiku.rag.ingester.queue.models import JobOp
from haiku.rag.ingester.sources import plugins as plugins_module
from haiku.rag.ingester.sources import resolve_configured_source
from haiku.rag.ingester.sources.base import (
    FetchResult,
    Source,
    SourceEvent,
    SourceEventKind,
)
from haiku.rag.ingester.sources.plugins import (
    ENTRY_POINT_GROUP,
    load_source_factories,
)


class _MemorySource:
    """A real Source over an in-memory {uri: text} dict. Doubles as its own
    factory: build_source calls it with the plugin kwargs."""

    def __init__(
        self,
        *,
        source_id: str,
        options: dict,
        supported_extensions: list[str] | None,
        max_file_size: int | None,
    ):
        self.source_id = source_id
        self.options = options
        self.supported_extensions = supported_extensions
        self.max_file_size = max_file_size
        self._docs: dict[str, str] = options["docs"]

    def supports(self, uri: str) -> bool:
        return uri in self._docs

    async def head(self, uri: str) -> str | None:
        return "v1"

    async def aclose(self) -> None:
        return None

    async def fetch(self, uri: str) -> FetchResult:
        body = self._docs[uri].encode()
        return FetchResult(
            uri=uri,
            body=body,
            content_type="text/markdown",
            content_hash=hashlib.md5(body).hexdigest(),
            revision="v1",
        )

    async def discover(self, since=None, *, known_uris=None):
        for uri in self._docs:
            yield SourceEvent(
                source_id=self.source_id,
                uri=uri,
                kind=SourceEventKind.UPSERT,
                revision="v1",
                discovered_at=datetime.now(UTC),
            )


class _FakeEntryPoint:
    def __init__(self, name, factory):
        self.name = name
        self._factory = factory
        self.loaded = False

    def load(self):
        self.loaded = True
        return self._factory


def _register(monkeypatch, *eps):
    captured: dict = {}

    def fake_entry_points(*, group):
        captured["group"] = group
        return list(eps)

    monkeypatch.setattr(plugins_module, "entry_points", fake_entry_points)
    return captured


def test_load_keys_entry_points_by_name_without_loading(monkeypatch):
    ep = _FakeEntryPoint("memory", _MemorySource)
    captured = _register(monkeypatch, ep)

    discovered = load_source_factories()

    assert captured["group"] == ENTRY_POINT_GROUP
    assert discovered == {"memory": ep}
    assert ep.loaded is False


def test_load_is_empty_when_none_registered(monkeypatch):
    _register(monkeypatch)
    assert load_source_factories() == {}


def _config(**overrides):
    return PluginSourceConfig(
        type="plugin",
        id="mem",
        plugin="memory",
        options={"docs": {"mem://a.md": "hello"}},
        **overrides,
    )


def test_build_source_loads_referenced_plugin_with_kwargs(monkeypatch):
    ep = _FakeEntryPoint("memory", _MemorySource)
    _register(monkeypatch, ep)

    source = build_source(_config(max_file_size=1024), supported_extensions=[".md"])

    assert ep.loaded is True
    assert isinstance(source, _MemorySource)
    assert isinstance(source, Source)
    assert source.source_id == "mem"
    assert source.options == {"docs": {"mem://a.md": "hello"}}
    assert source.supported_extensions == [".md"]
    assert source.max_file_size == 1024


def test_build_source_raises_on_unknown_plugin(monkeypatch):
    _register(monkeypatch, _FakeEntryPoint("memory", _MemorySource))

    cfg = _config().model_copy(update={"plugin": "missing"})
    with pytest.raises(ValueError, match="unknown source plugin 'missing'"):
        build_source(cfg)


def test_build_source_raises_when_plugin_returns_non_source(monkeypatch):
    _register(monkeypatch, _FakeEntryPoint("memory", lambda **kw: object()))

    with pytest.raises(TypeError, match="does not satisfy the Source protocol"):
        build_source(_config())


def test_build_source_does_not_load_unreferenced_plugins(monkeypatch):
    def _explode(**kw):
        raise ImportError("optional dependency missing")

    used = _FakeEntryPoint("memory", _MemorySource)
    unused = _FakeEntryPoint("broken", _explode)
    _register(monkeypatch, used, unused)

    build_source(_config())

    assert unused.loaded is False


@pytest.mark.asyncio
async def test_plugin_source_drives_poller_and_fetch(monkeypatch, jobs, sync):
    _register(monkeypatch, _FakeEntryPoint("memory", _MemorySource))
    cfg = _config()
    source = build_source(cfg)

    poller = PeriodicPoller(source=source, config=cfg, job_repo=jobs, sync_repo=sync)
    assert await poller._sweep_once() is True

    queued = await jobs.list_jobs(source_id="mem")
    assert len(queued) == 1
    assert queued[0].op is JobOp.UPSERT
    assert queued[0].uri == "mem://a.md"

    fetcher = resolve_configured_source("mem://a.md", "mem", [source])
    result = await fetcher.fetch("mem://a.md")
    assert result.body == b"hello"
    assert result.content_type == "text/markdown"