add SeaweedFS integration tests for S3Watcher
This commit is contained in:
parent
37099988a6
commit
f87d894b01
2 changed files with 190 additions and 1 deletions
24
scripts/run-integration-tests.sh
Executable file
24
scripts/run-integration-tests.sh
Executable 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 "$@"
|
||||||
|
|
@ -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"
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue