add SeaweedFS integration tests for S3Watcher

This commit is contained in:
Yiorgis Gozadinos 2026-04-29 14:19:55 +03:00
parent 37099988a6
commit f87d894b01
No known key found for this signature in database
2 changed files with 190 additions and 1 deletions

View file

@ -0,0 +1,24 @@
#!/bin/bash
# Run integration tests against a local SeaweedFS instance.
#
# Brings up SeaweedFS via docker compose, runs every test marked `integration`,
# and tears down on exit. Extra args are forwarded to pytest:
#
# ./scripts/run-integration-tests.sh
# ./scripts/run-integration-tests.sh tests/test_s3_integration.py::test_s3_watcher_initial_sweep
set -euo pipefail
REPO_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
COMPOSE_FILE="$REPO_ROOT/tests/docker/docker-compose.s3.yml"
cleanup() {
echo "==> Tearing down SeaweedFS"
docker compose -f "$COMPOSE_FILE" down -v >/dev/null 2>&1 || true
}
trap cleanup EXIT
echo "==> Bringing up SeaweedFS"
docker compose -f "$COMPOSE_FILE" up -d --wait
echo "==> Running integration tests"
uv run pytest -m integration "$@"

View file

@ -3,6 +3,7 @@
# Stop after: # Stop after:
# docker compose -f tests/docker/docker-compose.s3.yml down -v # docker compose -f tests/docker/docker-compose.s3.yml down -v
import importlib.util
import socket import socket
from uuid import uuid4 from uuid import uuid4
@ -10,9 +11,11 @@ import pytest
from haiku.rag.app import HaikuRAGApp from haiku.rag.app import HaikuRAGApp
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig, LanceDBConfig from haiku.rag.config.models import AppConfig, LanceDBConfig, S3MonitorEntry
from haiku.rag.store.engine import Store from haiku.rag.store.engine import Store
HAS_AIOBOTO3 = importlib.util.find_spec("aioboto3") is not None
S3_ENDPOINT = "http://localhost:8333" S3_ENDPOINT = "http://localhost:8333"
S3_BUCKET = "test-bucket" S3_BUCKET = "test-bucket"
S3_STORAGE_OPTIONS = { S3_STORAGE_OPTIONS = {
@ -149,3 +152,165 @@ async def test_app_info_empty_db(tmp_path, capsys):
out = capsys.readouterr().out out = capsys.readouterr().out
assert "Database is empty" in out assert "Database is empty" in out
# ----------------------- S3 watcher integration tests ----------------------- #
# These exercise S3Watcher against the live SeaweedFS instance. Documents are
# uploaded as raw S3 objects under a unique per-test prefix; the watcher's
# refresh() is invoked directly so tests stay deterministic. LanceDB stays
# local — these tests verify the watcher path, not LanceDB-on-S3.
_aioboto3_required = pytest.mark.skipif(
not HAS_AIOBOTO3,
reason="aioboto3 not installed (uv sync --extra s3)",
)
async def _put_object(prefix: str, key: str, body: bytes) -> None:
import aioboto3 # type: ignore[import-not-found] # ty: ignore[unresolved-import]
session = aioboto3.Session(
aws_access_key_id=S3_STORAGE_OPTIONS["aws_access_key_id"],
aws_secret_access_key=S3_STORAGE_OPTIONS["aws_secret_access_key"],
region_name=S3_STORAGE_OPTIONS["region"],
)
async with session.client(
"s3", endpoint_url=S3_STORAGE_OPTIONS["endpoint"], use_ssl=False
) as s3:
try:
await s3.create_bucket(Bucket=S3_BUCKET)
except Exception:
pass # bucket already exists
await s3.put_object(Bucket=S3_BUCKET, Key=f"{prefix}/{key}", Body=body)
async def _delete_object(prefix: str, key: str) -> None:
import aioboto3 # type: ignore[import-not-found] # ty: ignore[unresolved-import]
session = aioboto3.Session(
aws_access_key_id=S3_STORAGE_OPTIONS["aws_access_key_id"],
aws_secret_access_key=S3_STORAGE_OPTIONS["aws_secret_access_key"],
region_name=S3_STORAGE_OPTIONS["region"],
)
async with session.client(
"s3", endpoint_url=S3_STORAGE_OPTIONS["endpoint"], use_ssl=False
) as s3:
await s3.delete_object(Bucket=S3_BUCKET, Key=f"{prefix}/{key}")
def _watcher_entry(prefix: str, **overrides) -> S3MonitorEntry:
return S3MonitorEntry(
uri=overrides.pop("uri", f"s3://{S3_BUCKET}/{prefix}/"),
storage_options=overrides.pop("storage_options", S3_STORAGE_OPTIONS),
include_patterns=overrides.pop("include_patterns", ["*.txt"]),
delete_orphans=overrides.pop("delete_orphans", False),
poll_interval=overrides.pop("poll_interval", 60),
**overrides,
)
@_aioboto3_required
@pytest.mark.asyncio
async def test_s3_watcher_initial_sweep(tmp_path):
from haiku.rag.monitor import S3Watcher
prefix = f"watcher-init-{uuid4().hex[:8]}"
await _put_object(prefix, "alpha.txt", b"alpha content")
await _put_object(prefix, "beta.txt", b"beta content")
async with HaikuRAG(tmp_path / "db.lancedb", create=True) as rag:
watcher = S3Watcher(
client=rag,
entry=_watcher_entry(prefix),
supported_extensions=[".txt", ".md", ".pdf"],
)
await watcher.refresh()
docs = await rag.list_documents()
uris = sorted(d.uri or "" for d in docs)
assert uris == [
f"s3://{S3_BUCKET}/{prefix}/alpha.txt",
f"s3://{S3_BUCKET}/{prefix}/beta.txt",
]
@_aioboto3_required
@pytest.mark.asyncio
async def test_s3_watcher_detects_new_object(tmp_path):
from haiku.rag.monitor import S3Watcher
prefix = f"watcher-new-{uuid4().hex[:8]}"
await _put_object(prefix, "first.txt", b"first content")
async with HaikuRAG(tmp_path / "db.lancedb", create=True) as rag:
watcher = S3Watcher(
client=rag,
entry=_watcher_entry(prefix),
supported_extensions=[".txt"],
)
await watcher.refresh()
assert await rag.count_documents() == 1
await _put_object(prefix, "second.txt", b"second content")
await watcher.refresh()
assert await rag.count_documents() == 2
@_aioboto3_required
@pytest.mark.asyncio
async def test_s3_watcher_detects_modified_object(tmp_path):
from haiku.rag.monitor import S3Watcher
prefix = f"watcher-mod-{uuid4().hex[:8]}"
uri = f"s3://{S3_BUCKET}/{prefix}/file.txt"
await _put_object(prefix, "file.txt", b"original content")
async with HaikuRAG(tmp_path / "db.lancedb", create=True) as rag:
watcher = S3Watcher(
client=rag,
entry=_watcher_entry(prefix),
supported_extensions=[".txt"],
)
await watcher.refresh()
first = await rag.get_document_by_uri(uri)
assert first is not None
first_md5 = first.metadata["md5"]
await _put_object(prefix, "file.txt", b"new content body")
await watcher.refresh()
second = await rag.get_document_by_uri(uri)
assert second is not None
assert second.id == first.id
assert second.metadata["md5"] != first_md5
assert "new content body" in second.content
@_aioboto3_required
@pytest.mark.asyncio
async def test_s3_watcher_orphan_deletion(tmp_path):
from haiku.rag.monitor import S3Watcher
prefix = f"watcher-orphan-{uuid4().hex[:8]}"
await _put_object(prefix, "kept.txt", b"keep me")
await _put_object(prefix, "doomed.txt", b"will be deleted")
async with HaikuRAG(tmp_path / "db.lancedb", create=True) as rag:
watcher = S3Watcher(
client=rag,
entry=_watcher_entry(prefix, delete_orphans=True),
supported_extensions=[".txt"],
)
await watcher.refresh()
assert await rag.count_documents() == 2
await _delete_object(prefix, "doomed.txt")
await watcher.refresh()
docs = await rag.list_documents()
assert len(docs) == 1
assert docs[0].uri == f"s3://{S3_BUCKET}/{prefix}/kept.txt"