65 lines
1.9 KiB
Python
65 lines
1.9 KiB
Python
import os
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from claude_agent_sdk.testing import run_session_store_conformance
|
|
from sqlalchemy import delete
|
|
|
|
from beaver_gateway.storage import Database, PostgresSessionStore, TranscriptEntry
|
|
|
|
POSTGRES_URL = os.environ.get("BEAVER_TEST_DATABASE_URL")
|
|
|
|
|
|
async def test_sqlite_conformance() -> None:
|
|
root = Path(tempfile.mkdtemp(prefix="beaver-store-"))
|
|
counter = 0
|
|
|
|
async def make_store() -> PostgresSessionStore:
|
|
nonlocal counter
|
|
counter += 1
|
|
db = Database(f"sqlite:///{root / f'{counter}.db'}")
|
|
await db.create_all()
|
|
return PostgresSessionStore(db)
|
|
|
|
await run_session_store_conformance(make_store)
|
|
|
|
|
|
@pytest.mark.skipif(POSTGRES_URL is None, reason="BEAVER_TEST_DATABASE_URL not set")
|
|
async def test_postgres_conformance() -> None:
|
|
assert POSTGRES_URL is not None
|
|
db = Database(POSTGRES_URL)
|
|
await db.create_all()
|
|
|
|
async def make_store() -> PostgresSessionStore:
|
|
async with db.session() as session:
|
|
await session.execute(delete(TranscriptEntry))
|
|
await session.commit()
|
|
return PostgresSessionStore(db)
|
|
|
|
try:
|
|
await run_session_store_conformance(make_store)
|
|
finally:
|
|
await db.dispose()
|
|
|
|
|
|
async def test_append_dedups_by_uuid_and_keeps_floats() -> None:
|
|
db = Database(f"sqlite:///{tempfile.mkdtemp(prefix='beaver-store-')}/d.db")
|
|
await db.create_all()
|
|
store = PostgresSessionStore(db)
|
|
key = {"project_key": "p", "session_id": "s"}
|
|
batch = [
|
|
{"type": "x", "uuid": "a", "n": 1.5},
|
|
{"type": "y", "n": 2},
|
|
{"type": "x", "uuid": "a", "n": 999},
|
|
]
|
|
await store.append(key, batch)
|
|
await store.append(key, batch)
|
|
loaded = await store.load(key)
|
|
assert loaded == [
|
|
{"type": "x", "uuid": "a", "n": 1.5},
|
|
{"type": "y", "n": 2},
|
|
{"type": "y", "n": 2},
|
|
]
|
|
assert isinstance(loaded[0]["n"], float)
|