121 lines
4.1 KiB
Python
121 lines
4.1 KiB
Python
"""DbConfirmationStore 单元测试。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import AsyncGenerator
|
|
|
|
import pytest_asyncio
|
|
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
|
|
from sqlalchemy.orm import sessionmaker
|
|
|
|
from app.db.base import Base
|
|
from app.modules.safety.confirmation import (
|
|
DbConfirmationStore,
|
|
PendingConfirmation,
|
|
new_confirmation_id,
|
|
)
|
|
|
|
_TEST_DATABASE_URL = "sqlite+aiosqlite://"
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def db_session() -> AsyncGenerator[AsyncSession, None]:
|
|
engine = create_async_engine(_TEST_DATABASE_URL, echo=False)
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
async with async_session() as session:
|
|
yield session
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def store(db_session: AsyncSession) -> DbConfirmationStore:
|
|
"""创建绑定 db_session 的 DbConfirmationStore。"""
|
|
session_factory = sessionmaker(db_session.bind, class_=AsyncSession, expire_on_commit=False)
|
|
return DbConfirmationStore(session_factory)
|
|
|
|
|
|
class TestDbConfirmationStore:
|
|
"""数据库确认存储 CRUD 测试。"""
|
|
|
|
async def test_create_and_get(self, store: DbConfirmationStore) -> None:
|
|
pending = PendingConfirmation(
|
|
confirmation_id=new_confirmation_id(),
|
|
tool="fetch_url",
|
|
input={"q": "test"},
|
|
risk_level="L2",
|
|
reason="测试确认",
|
|
)
|
|
created = await store.create(pending)
|
|
assert created.confirmation_id == pending.confirmation_id
|
|
|
|
fetched = await store.get(pending.confirmation_id)
|
|
assert fetched is not None
|
|
assert fetched.tool == "fetch_url"
|
|
assert fetched.input == {"q": "test"}
|
|
|
|
async def test_get_non_existent(self, store: DbConfirmationStore) -> None:
|
|
result = await store.get("confirm_nonexistent")
|
|
assert result is None
|
|
|
|
async def test_pop_removes_and_returns(self, store: DbConfirmationStore) -> None:
|
|
pending = PendingConfirmation(
|
|
confirmation_id=new_confirmation_id(),
|
|
tool="shell",
|
|
input={"cmd": "ls"},
|
|
risk_level="L3",
|
|
reason="高风险操作",
|
|
)
|
|
await store.create(pending)
|
|
|
|
popped = await store.pop(pending.confirmation_id)
|
|
assert popped is not None
|
|
assert popped.tool == "shell"
|
|
|
|
# pop 后不应再可取到
|
|
fetched = await store.get(pending.confirmation_id)
|
|
assert fetched is None
|
|
|
|
async def test_pop_non_existent(self, store: DbConfirmationStore) -> None:
|
|
result = await store.pop("confirm_nonexistent")
|
|
assert result is None
|
|
|
|
async def test_list_pending(self, store: DbConfirmationStore) -> None:
|
|
# 插入两条
|
|
id1 = new_confirmation_id()
|
|
id2 = new_confirmation_id()
|
|
await store.create(
|
|
PendingConfirmation(
|
|
confirmation_id=id1, tool="fetch_url", input={}, risk_level="L1", reason="a"
|
|
)
|
|
)
|
|
await store.create(
|
|
PendingConfirmation(
|
|
confirmation_id=id2, tool="shell", input={}, risk_level="L3", reason="b"
|
|
)
|
|
)
|
|
items = await store.list_pending()
|
|
assert len(items) == 2
|
|
# 按创建时间升序
|
|
assert items[0].confirmation_id == id1
|
|
assert items[1].confirmation_id == id2
|
|
|
|
async def test_list_pending_empty(self, store: DbConfirmationStore) -> None:
|
|
items = await store.list_pending()
|
|
assert items == []
|
|
|
|
async def test_popped_not_in_list_pending(self, store: DbConfirmationStore) -> None:
|
|
pending = PendingConfirmation(
|
|
confirmation_id=new_confirmation_id(),
|
|
tool="fetch_url",
|
|
input={},
|
|
risk_level="L2",
|
|
reason="测试",
|
|
)
|
|
await store.create(pending)
|
|
await store.pop(pending.confirmation_id)
|
|
items = await store.list_pending()
|
|
ids = [i.confirmation_id for i in items]
|
|
assert pending.confirmation_id not in ids
|