Clean up S3 test data per-test to keep SeaweedFS writable
This commit is contained in:
parent
98f2a088d8
commit
ce27a7aae5
1 changed files with 29 additions and 18 deletions
|
|
@ -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()
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue