392 lines
11 KiB
Python
392 lines
11 KiB
Python
import asyncio
|
|
import sys
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config import AppConfig, MonitorConfig, S3MonitorEntry
|
|
from haiku.rag.store.models.document import Document
|
|
|
|
|
|
@pytest.fixture
|
|
def fake_aioboto3(monkeypatch):
|
|
fake = MagicMock()
|
|
monkeypatch.setitem(sys.modules, "aioboto3", fake)
|
|
return fake
|
|
|
|
|
|
@pytest.fixture
|
|
def s3_paginator():
|
|
"""Return (paginator_mock, set_pages) — `set_pages([page, ...])` rewires the iterator."""
|
|
pages: list[dict] = []
|
|
|
|
async def _paginate(**_kwargs):
|
|
for page in pages:
|
|
yield page
|
|
|
|
paginator = MagicMock()
|
|
paginator.paginate.side_effect = lambda **kw: _paginate(**kw)
|
|
|
|
def set_pages(new_pages):
|
|
pages.clear()
|
|
pages.extend(new_pages)
|
|
|
|
return paginator, set_pages
|
|
|
|
|
|
@pytest.fixture
|
|
def fake_s3_client(fake_aioboto3, s3_paginator):
|
|
paginator, set_pages = s3_paginator
|
|
s3_client = MagicMock()
|
|
s3_client.get_paginator.return_value = paginator
|
|
|
|
client_ctx = AsyncMock()
|
|
client_ctx.__aenter__.return_value = s3_client
|
|
client_ctx.__aexit__.return_value = None
|
|
|
|
session = MagicMock()
|
|
session.client.return_value = client_ctx
|
|
fake_aioboto3.Session.return_value = session
|
|
|
|
return s3_client, set_pages
|
|
|
|
|
|
def _entry(**kwargs):
|
|
base = {
|
|
"uri": "s3://my-bucket/incoming/",
|
|
"poll_interval": 60,
|
|
"delete_orphans": False,
|
|
"ignore_patterns": [],
|
|
"include_patterns": [],
|
|
"storage_options": {},
|
|
}
|
|
base.update(kwargs)
|
|
return S3MonitorEntry(**base)
|
|
|
|
|
|
def _doc(uri: str, etag: str, doc_id: str | None = None) -> Document:
|
|
return Document(
|
|
id=doc_id or uri,
|
|
content="...",
|
|
uri=uri,
|
|
metadata={"etag": etag, "md5": "deadbeef"},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_s3_watcher_refresh_upserts_new_objects(fake_s3_client):
|
|
s3, set_pages = fake_s3_client
|
|
set_pages(
|
|
[
|
|
{
|
|
"Contents": [
|
|
{"Key": "incoming/a.txt", "ETag": '"abc"'},
|
|
{"Key": "incoming/b.txt", "ETag": '"def"'},
|
|
]
|
|
}
|
|
]
|
|
)
|
|
|
|
from haiku.rag.monitor import S3Watcher
|
|
|
|
rag = AsyncMock(spec=HaikuRAG)
|
|
rag.list_documents.return_value = []
|
|
rag.create_document_from_source.return_value = Document(
|
|
id="x", content="...", uri="s3://my-bucket/incoming/a.txt"
|
|
)
|
|
|
|
watcher = S3Watcher(
|
|
client=rag, entry=_entry(), supported_extensions=[".txt", ".md", ".pdf"]
|
|
)
|
|
await watcher.refresh()
|
|
|
|
assert rag.create_document_from_source.await_count == 2
|
|
called_uris = {c.args[0] for c in rag.create_document_from_source.await_args_list}
|
|
assert called_uris == {
|
|
"s3://my-bucket/incoming/a.txt",
|
|
"s3://my-bucket/incoming/b.txt",
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_s3_watcher_skips_unchanged_etag(fake_s3_client):
|
|
s3, set_pages = fake_s3_client
|
|
set_pages([{"Contents": [{"Key": "incoming/a.txt", "ETag": '"abc"'}]}])
|
|
|
|
from haiku.rag.monitor import S3Watcher
|
|
|
|
rag = AsyncMock(spec=HaikuRAG)
|
|
rag.list_documents.return_value = [_doc("s3://my-bucket/incoming/a.txt", "abc")]
|
|
|
|
watcher = S3Watcher(client=rag, entry=_entry(), supported_extensions=[".txt"])
|
|
await watcher.refresh()
|
|
|
|
rag.create_document_from_source.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_s3_watcher_upserts_when_etag_differs(fake_s3_client):
|
|
s3, set_pages = fake_s3_client
|
|
set_pages([{"Contents": [{"Key": "incoming/a.txt", "ETag": '"new"'}]}])
|
|
|
|
from haiku.rag.monitor import S3Watcher
|
|
|
|
rag = AsyncMock(spec=HaikuRAG)
|
|
rag.list_documents.return_value = [_doc("s3://my-bucket/incoming/a.txt", "old")]
|
|
rag.create_document_from_source.return_value = Document(
|
|
id="x", content="...", uri="s3://my-bucket/incoming/a.txt"
|
|
)
|
|
|
|
watcher = S3Watcher(client=rag, entry=_entry(), supported_extensions=[".txt"])
|
|
await watcher.refresh()
|
|
|
|
rag.create_document_from_source.assert_awaited_once_with(
|
|
"s3://my-bucket/incoming/a.txt", storage_options={}
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_s3_watcher_strips_etag_quotes(fake_s3_client):
|
|
s3, set_pages = fake_s3_client
|
|
set_pages([{"Contents": [{"Key": "incoming/a.txt", "ETag": '"abc"'}]}])
|
|
|
|
from haiku.rag.monitor import S3Watcher
|
|
|
|
rag = AsyncMock(spec=HaikuRAG)
|
|
rag.list_documents.return_value = [
|
|
_doc("s3://my-bucket/incoming/a.txt", "abc") # already stripped in storage
|
|
]
|
|
|
|
watcher = S3Watcher(client=rag, entry=_entry(), supported_extensions=[".txt"])
|
|
await watcher.refresh()
|
|
|
|
rag.create_document_from_source.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_s3_watcher_deletes_orphans_when_enabled(fake_s3_client):
|
|
s3, set_pages = fake_s3_client
|
|
set_pages([{"Contents": [{"Key": "incoming/a.txt", "ETag": '"abc"'}]}])
|
|
|
|
from haiku.rag.monitor import S3Watcher
|
|
|
|
a_doc = _doc("s3://my-bucket/incoming/a.txt", "abc", doc_id="a-id")
|
|
orphan = _doc("s3://my-bucket/incoming/old.txt", "stale", doc_id="orphan-id")
|
|
|
|
rag = AsyncMock(spec=HaikuRAG)
|
|
rag.list_documents.return_value = [a_doc, orphan]
|
|
rag.get_document_by_uri.return_value = orphan
|
|
|
|
watcher = S3Watcher(
|
|
client=rag,
|
|
entry=_entry(delete_orphans=True),
|
|
supported_extensions=[".txt"],
|
|
)
|
|
await watcher.refresh()
|
|
|
|
rag.delete_document.assert_awaited_once_with("orphan-id")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_s3_watcher_does_not_delete_orphans_when_disabled(fake_s3_client):
|
|
s3, set_pages = fake_s3_client
|
|
set_pages([{"Contents": []}])
|
|
|
|
from haiku.rag.monitor import S3Watcher
|
|
|
|
orphan = _doc("s3://my-bucket/incoming/old.txt", "stale", doc_id="orphan-id")
|
|
|
|
rag = AsyncMock(spec=HaikuRAG)
|
|
rag.list_documents.return_value = [orphan]
|
|
|
|
watcher = S3Watcher(
|
|
client=rag,
|
|
entry=_entry(delete_orphans=False),
|
|
supported_extensions=[".txt"],
|
|
)
|
|
await watcher.refresh()
|
|
|
|
rag.delete_document.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_s3_watcher_orphan_scope_is_per_entry(fake_s3_client):
|
|
"""A doc under a different bucket prefix must not be touched."""
|
|
s3, set_pages = fake_s3_client
|
|
set_pages([{"Contents": []}])
|
|
|
|
from haiku.rag.monitor import S3Watcher
|
|
|
|
rag = AsyncMock(spec=HaikuRAG)
|
|
rag.list_documents.return_value = [] # filter scopes to my-bucket
|
|
|
|
watcher = S3Watcher(
|
|
client=rag,
|
|
entry=_entry(delete_orphans=True),
|
|
supported_extensions=[".txt"],
|
|
)
|
|
await watcher.refresh()
|
|
|
|
rag.list_documents.assert_awaited_once()
|
|
filter_kwarg = rag.list_documents.await_args.kwargs["filter"]
|
|
assert filter_kwarg == "uri LIKE 's3://my-bucket/incoming/%'"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_s3_watcher_applies_include_and_ignore_patterns(fake_s3_client):
|
|
s3, set_pages = fake_s3_client
|
|
set_pages(
|
|
[
|
|
{
|
|
"Contents": [
|
|
{"Key": "incoming/keep.md", "ETag": '"1"'},
|
|
{"Key": "incoming/draft.md", "ETag": '"2"'},
|
|
{"Key": "incoming/skip.txt", "ETag": '"3"'},
|
|
]
|
|
}
|
|
]
|
|
)
|
|
|
|
from haiku.rag.monitor import S3Watcher
|
|
|
|
rag = AsyncMock(spec=HaikuRAG)
|
|
rag.list_documents.return_value = []
|
|
rag.create_document_from_source.return_value = Document(
|
|
id="x", content="...", uri="s3://my-bucket/incoming/keep.md"
|
|
)
|
|
|
|
watcher = S3Watcher(
|
|
client=rag,
|
|
entry=_entry(include_patterns=["*.md"], ignore_patterns=["draft*"]),
|
|
supported_extensions=[".md", ".txt"],
|
|
)
|
|
await watcher.refresh()
|
|
|
|
assert rag.create_document_from_source.await_count == 1
|
|
assert (
|
|
rag.create_document_from_source.await_args.args[0]
|
|
== "s3://my-bucket/incoming/keep.md"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_s3_watcher_observe_survives_transient_list_failure(fake_s3_client):
|
|
"""First refresh succeeds; second refresh raises; loop survives and recovers."""
|
|
s3, set_pages = fake_s3_client
|
|
|
|
pages_initial = [{"Contents": [{"Key": "incoming/a.txt", "ETag": '"abc"'}]}]
|
|
pages_after = [{"Contents": [{"Key": "incoming/a.txt", "ETag": '"abc"'}]}]
|
|
|
|
set_pages(pages_initial)
|
|
paginate_calls = {"n": 0}
|
|
|
|
async def paginate_side_effect(**_kwargs):
|
|
paginate_calls["n"] += 1
|
|
if paginate_calls["n"] == 2:
|
|
raise RuntimeError("transient list failure")
|
|
for page in pages_after if paginate_calls["n"] > 1 else pages_initial:
|
|
yield page
|
|
|
|
s3.get_paginator.return_value.paginate.side_effect = paginate_side_effect
|
|
|
|
from haiku.rag.monitor import S3Watcher
|
|
|
|
rag = AsyncMock(spec=HaikuRAG)
|
|
rag.list_documents.return_value = []
|
|
rag.create_document_from_source.return_value = Document(
|
|
id="x", content="...", uri="s3://my-bucket/incoming/a.txt"
|
|
)
|
|
|
|
watcher = S3Watcher(
|
|
client=rag,
|
|
entry=_entry(poll_interval=0),
|
|
supported_extensions=[".txt"],
|
|
)
|
|
task = asyncio.create_task(watcher.observe())
|
|
|
|
# Let three iterations run: initial refresh, transient failure, recovery.
|
|
for _ in range(20):
|
|
await asyncio.sleep(0)
|
|
if paginate_calls["n"] >= 3:
|
|
break
|
|
|
|
task.cancel()
|
|
try:
|
|
await task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
assert paginate_calls["n"] >= 3 # loop kept going past the failure
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_s3_watcher_invalid_uri_rejected():
|
|
from haiku.rag.monitor import S3Watcher
|
|
|
|
rag = AsyncMock(spec=HaikuRAG)
|
|
with pytest.raises(ValueError, match="Invalid S3 monitor URI"):
|
|
S3Watcher(
|
|
client=rag,
|
|
entry=S3MonitorEntry(uri="s3://"),
|
|
supported_extensions=[".txt"],
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_serve_starts_one_s3_task_per_entry(monkeypatch, fake_s3_client):
|
|
"""`serve` wires one S3Watcher task per MonitorConfig.s3 entry."""
|
|
from haiku.rag import app as app_module
|
|
|
|
captured_tasks: list = []
|
|
original_create_task = asyncio.create_task
|
|
|
|
def tracking_create_task(coro, *args, **kwargs):
|
|
captured_tasks.append(coro)
|
|
return original_create_task(coro, *args, **kwargs)
|
|
|
|
monkeypatch.setattr(app_module.asyncio, "create_task", tracking_create_task)
|
|
|
|
config = AppConfig(
|
|
monitor=MonitorConfig(
|
|
s3=[
|
|
S3MonitorEntry(uri="s3://bucket-a/x/"),
|
|
S3MonitorEntry(uri="s3://bucket-b/y/"),
|
|
]
|
|
)
|
|
)
|
|
|
|
fw_observe_calls = {"n": 0}
|
|
|
|
async def fake_fw_observe(self):
|
|
fw_observe_calls["n"] += 1
|
|
|
|
monkeypatch.setattr(app_module.FileWatcher, "observe", fake_fw_observe)
|
|
|
|
s3_observe_calls = {"n": 0}
|
|
|
|
async def fake_s3_observe(self):
|
|
s3_observe_calls["n"] += 1
|
|
|
|
monkeypatch.setattr(app_module.S3Watcher, "observe", fake_s3_observe)
|
|
|
|
# Provide a dummy supported_extensions to skip docling import.
|
|
class _Conv:
|
|
supported_extensions = [".txt"]
|
|
|
|
monkeypatch.setattr("haiku.rag.converters.get_converter", lambda cfg: _Conv())
|
|
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
db_path = Path(tmp) / "db.lancedb"
|
|
app = app_module.HaikuRAGApp(db_path=db_path, config=config)
|
|
async with HaikuRAG(db_path, config=config, create=True):
|
|
pass # create the database
|
|
|
|
await app.serve(enable_monitor=True, enable_mcp=False)
|
|
|
|
# Both S3 entries should have triggered observe(), plus the FileWatcher.
|
|
assert fw_observe_calls["n"] == 1
|
|
assert s3_observe_calls["n"] == 2
|