diff --git a/tests/test_s3_integration.py b/tests/test_s3_integration.py index 7276101b..05185499 100644 --- a/tests/test_s3_integration.py +++ b/tests/test_s3_integration.py @@ -5,11 +5,13 @@ from uuid import uuid4 +import obstore 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.s3 import make_s3_store from haiku.rag.store.engine import Store from tests.services import reachable @@ -43,11 +45,28 @@ def _make_config() -> AppConfig: ) +@pytest.fixture +def config(): + """A config pointing at a unique S3 prefix, cleaned up after the test. + + SeaweedFS reclaims volume space lazily, so leaving each run's data behind + eventually fills the volume server and seals its volumes read-only, which + makes every subsequent write block forever. + """ + config = _make_config() + yield config + + bucket, _, prefix = config.lancedb.uri.removeprefix("s3://").partition("/") + store = make_s3_store(bucket, S3_STORAGE_OPTIONS) + paths = [obj["path"] for batch in store.list(prefix=f"{prefix}/") for obj in batch] + if paths: + obstore.delete(store, paths) + + @pytest.mark.asyncio -async def test_store_connect_and_create(tmp_path): +async def test_store_connect_and_create(tmp_path, config): from haiku.rag.store.engine import get_database_stats - config = _make_config() async with Store(tmp_path / "unused", config=config, create=True) as store: stats = await get_database_stats(store.db) assert stats["documents"]["exists"] @@ -55,17 +74,15 @@ async def test_store_connect_and_create(tmp_path): @pytest.mark.asyncio -async def test_store_vacuum(tmp_path): - config = _make_config() +async def test_store_vacuum(tmp_path, config): async with Store(tmp_path / "unused", config=config, create=True) as store: await store.vacuum() @pytest.mark.asyncio -async def test_store_add_document(tmp_path): +async def test_store_add_document(tmp_path, config): from haiku.rag.store.engine import DocumentRecord, get_database_stats - config = _make_config() async with Store(tmp_path / "unused", config=config, create=True) as store: doc = DocumentRecord(content="The quick brown fox jumps over the lazy dog.") await store.documents_table.add([doc]) @@ -75,8 +92,7 @@ async def test_store_add_document(tmp_path): @pytest.mark.asyncio -async def test_client_create_document(tmp_path): - config = _make_config() +async def test_client_create_document(tmp_path, config): async with HaikuRAG(tmp_path / "unused", config=config, create=True) as rag: doc = await rag.create_document( "Python is a programming language.", uri="test://python" @@ -86,8 +102,7 @@ async def test_client_create_document(tmp_path): @pytest.mark.asyncio -async def test_client_list_documents(tmp_path): - config = _make_config() +async def test_client_list_documents(tmp_path, config): async with HaikuRAG(tmp_path / "unused", config=config, create=True) as rag: await rag.create_document("First document.", uri="test://first") await rag.create_document("Second document.", uri="test://second") @@ -97,8 +112,7 @@ async def test_client_list_documents(tmp_path): @pytest.mark.asyncio -async def test_client_search(tmp_path): - config = _make_config() +async def test_client_search(tmp_path, config): async with HaikuRAG(tmp_path / "unused", config=config, create=True) as rag: await rag.create_document( "The Eiffel Tower is located in Paris, France.", uri="test://eiffel" @@ -109,8 +123,7 @@ async def test_client_search(tmp_path): @pytest.mark.asyncio -async def test_client_delete_document(tmp_path): - config = _make_config() +async def test_client_delete_document(tmp_path, config): async with HaikuRAG(tmp_path / "unused", config=config, create=True) as rag: doc = await rag.create_document("Temporary document.", uri="test://temp") await rag.delete_document(doc.id) @@ -119,8 +132,7 @@ async def test_client_delete_document(tmp_path): @pytest.mark.asyncio -async def test_app_info(tmp_path, capsys): - config = _make_config() +async def test_app_info(tmp_path, capsys, config): async with HaikuRAG(tmp_path / "unused", config=config, create=True) as rag: await rag.create_document("Info test document.", uri="test://info") @@ -134,8 +146,7 @@ async def test_app_info(tmp_path, capsys): @pytest.mark.asyncio -async def test_app_info_empty_db(tmp_path, capsys): - config = _make_config() +async def test_app_info_empty_db(tmp_path, capsys, config): app = HaikuRAGApp(db_path=tmp_path / "unused", config=config) await app.info()