初始化提交

backend/tests/modules/safety/test_db_confirmation_store.py

@@ -0,0 +1,120 @@
+"""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