62 lines
2.0 KiB
Python
62 lines
2.0 KiB
Python
import sqlite3
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
from sqlmodel import select
|
|
|
|
from beaver_gateway.storage import Database
|
|
from beaver_gateway.storage.models import Conversation
|
|
|
|
|
|
async def test_create_all_adds_missing_columns() -> None:
|
|
path = Path(tempfile.mkdtemp(prefix="beaver-migrate-")) / "old.db"
|
|
raw = sqlite3.connect(path)
|
|
raw.execute(
|
|
"CREATE TABLE conversations (id INTEGER PRIMARY KEY, frontend VARCHAR NOT NULL, "
|
|
"external_id VARCHAR NOT NULL, agent_name VARCHAR NOT NULL, "
|
|
"created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL)"
|
|
)
|
|
raw.execute(
|
|
"INSERT INTO conversations VALUES (1, 'markdown', 'x', 'a', '2026-01-01', '2026-01-01')"
|
|
)
|
|
raw.commit()
|
|
raw.close()
|
|
|
|
db = Database(f"sqlite:///{path}")
|
|
await db.create_all()
|
|
async with db.session() as session:
|
|
conv = (await session.exec(select(Conversation))).one()
|
|
assert conv.session_id is None
|
|
conv.session_id = "sid"
|
|
session.add(conv)
|
|
await session.commit()
|
|
await db.dispose()
|
|
|
|
|
|
async def test_create_all_backfills_defaults_for_old_rows() -> None:
|
|
path = Path(tempfile.mkdtemp(prefix="beaver-migrate-")) / "old.db"
|
|
raw = sqlite3.connect(path)
|
|
raw.execute(
|
|
"CREATE TABLE conversations (id INTEGER PRIMARY KEY, frontend VARCHAR NOT NULL, "
|
|
"external_id VARCHAR NOT NULL, agent_name VARCHAR NOT NULL, session_id VARCHAR, "
|
|
"created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL)"
|
|
)
|
|
raw.execute(
|
|
"INSERT INTO conversations VALUES (1, 'markdown', 'x', 'a', NULL, '2026-01-01', '2026-01-01')"
|
|
)
|
|
raw.commit()
|
|
raw.close()
|
|
|
|
db = Database(f"sqlite:///{path}")
|
|
await db.create_all()
|
|
async with db.session() as session:
|
|
conv = (await session.exec(select(Conversation))).one()
|
|
assert (conv.kind, conv.status, conv.pending_question, conv.flags) == (
|
|
"deep",
|
|
"open",
|
|
False,
|
|
{},
|
|
)
|
|
assert conv.running_turn is None
|
|
await db.dispose()
|