Clean up S3 test data per-test to keep SeaweedFS writable

This commit is contained in:
Yiorgis Gozadinos 2026-06-26 16:41:20 +03:00
parent 98f2a088d8
commit ce27a7aae5
No known key found for this signature in database

View file

@ -5,11 +5,13 @@
from uuid import uuid4 from uuid import uuid4
import obstore
import pytest 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
from haiku.rag.s3 import make_s3_store
from haiku.rag.store.engine import Store from haiku.rag.store.engine import Store
from tests.services import reachable 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 @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 from haiku.rag.store.engine import get_database_stats
config = _make_config()
async with Store(tmp_path / "unused", config=config, create=True) as store: async with Store(tmp_path / "unused", config=config, create=True) as store:
stats = await get_database_stats(store.db) stats = await get_database_stats(store.db)
assert stats["documents"]["exists"] assert stats["documents"]["exists"]
@ -55,17 +74,15 @@ async def test_store_connect_and_create(tmp_path):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_store_vacuum(tmp_path): async def test_store_vacuum(tmp_path, config):
config = _make_config()
async with Store(tmp_path / "unused", config=config, create=True) as store: async with Store(tmp_path / "unused", config=config, create=True) as store:
await store.vacuum() await store.vacuum()
@pytest.mark.asyncio @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 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: async with Store(tmp_path / "unused", config=config, create=True) as store:
doc = DocumentRecord(content="The quick brown fox jumps over the lazy dog.") doc = DocumentRecord(content="The quick brown fox jumps over the lazy dog.")
await store.documents_table.add([doc]) await store.documents_table.add([doc])
@ -75,8 +92,7 @@ async def test_store_add_document(tmp_path):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_create_document(tmp_path): async def test_client_create_document(tmp_path, config):
config = _make_config()
async with HaikuRAG(tmp_path / "unused", config=config, create=True) as rag: async with HaikuRAG(tmp_path / "unused", config=config, create=True) as rag:
doc = await rag.create_document( doc = await rag.create_document(
"Python is a programming language.", uri="test://python" "Python is a programming language.", uri="test://python"
@ -86,8 +102,7 @@ async def test_client_create_document(tmp_path):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_list_documents(tmp_path): async def test_client_list_documents(tmp_path, config):
config = _make_config()
async with HaikuRAG(tmp_path / "unused", config=config, create=True) as rag: 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("First document.", uri="test://first")
await rag.create_document("Second document.", uri="test://second") await rag.create_document("Second document.", uri="test://second")
@ -97,8 +112,7 @@ async def test_client_list_documents(tmp_path):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_search(tmp_path): async def test_client_search(tmp_path, config):
config = _make_config()
async with HaikuRAG(tmp_path / "unused", config=config, create=True) as rag: async with HaikuRAG(tmp_path / "unused", config=config, create=True) as rag:
await rag.create_document( await rag.create_document(
"The Eiffel Tower is located in Paris, France.", uri="test://eiffel" "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 @pytest.mark.asyncio
async def test_client_delete_document(tmp_path): async def test_client_delete_document(tmp_path, config):
config = _make_config()
async with HaikuRAG(tmp_path / "unused", config=config, create=True) as rag: async with HaikuRAG(tmp_path / "unused", config=config, create=True) as rag:
doc = await rag.create_document("Temporary document.", uri="test://temp") doc = await rag.create_document("Temporary document.", uri="test://temp")
await rag.delete_document(doc.id) await rag.delete_document(doc.id)
@ -119,8 +132,7 @@ async def test_client_delete_document(tmp_path):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_app_info(tmp_path, capsys): async def test_app_info(tmp_path, capsys, config):
config = _make_config()
async with HaikuRAG(tmp_path / "unused", config=config, create=True) as rag: async with HaikuRAG(tmp_path / "unused", config=config, create=True) as rag:
await rag.create_document("Info test document.", uri="test://info") 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 @pytest.mark.asyncio
async def test_app_info_empty_db(tmp_path, capsys): async def test_app_info_empty_db(tmp_path, capsys, config):
config = _make_config()
app = HaikuRAGApp(db_path=tmp_path / "unused", config=config) app = HaikuRAGApp(db_path=tmp_path / "unused", config=config)
await app.info() await app.info()