From f87d894b017fbaa4c0856f7054867a0a8dbb2fe3 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 29 Apr 2026 14:19:55 +0300 Subject: [PATCH] add SeaweedFS integration tests for S3Watcher --- scripts/run-integration-tests.sh | 24 +++++ tests/test_s3_integration.py | 167 ++++++++++++++++++++++++++++++- 2 files changed, 190 insertions(+), 1 deletion(-) create mode 100755 scripts/run-integration-tests.sh diff --git a/scripts/run-integration-tests.sh b/scripts/run-integration-tests.sh new file mode 100755 index 00000000..86c168c6 --- /dev/null +++ b/scripts/run-integration-tests.sh @@ -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 "$@" diff --git a/tests/test_s3_integration.py b/tests/test_s3_integration.py index c13cf029..9dfe90be 100644 --- a/tests/test_s3_integration.py +++ b/tests/test_s3_integration.py @@ -3,6 +3,7 @@ # Stop after: # docker compose -f tests/docker/docker-compose.s3.yml down -v +import importlib.util import socket from uuid import uuid4 @@ -10,9 +11,11 @@ import pytest from haiku.rag.app import HaikuRAGApp 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 +HAS_AIOBOTO3 = importlib.util.find_spec("aioboto3") is not None + S3_ENDPOINT = "http://localhost:8333" S3_BUCKET = "test-bucket" S3_STORAGE_OPTIONS = { @@ -149,3 +152,165 @@ async def test_app_info_empty_db(tmp_path, capsys): out = capsys.readouterr().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"