87 lines
2.8 KiB
Python
87 lines
2.8 KiB
Python
"""``IndividualDatabase`` —— plugin_owned slot 路径单元测试。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
from kilostar.core.postgres_database.module.individual import IndividualDatabase
|
|
from kilostar.utils.error import BusinessError
|
|
|
|
|
|
def _make_db():
|
|
session = AsyncMock()
|
|
session.__aenter__ = AsyncMock(return_value=session)
|
|
session.__aexit__ = AsyncMock(return_value=False)
|
|
session_maker = MagicMock(return_value=session)
|
|
return IndividualDatabase(session_maker), session
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delete_blocks_plugin_owned_row():
|
|
db, session = _make_db()
|
|
fake = MagicMock()
|
|
fake.plugin_owned = "data_analytics"
|
|
execute_result = MagicMock()
|
|
execute_result.scalar_one_or_none.return_value = fake
|
|
session.execute = AsyncMock(return_value=execute_result)
|
|
|
|
with pytest.raises(BusinessError, match="不可删除"):
|
|
await db.delete_worker_individual("agent-x")
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delete_allows_user_owned_row():
|
|
db, session = _make_db()
|
|
fake = MagicMock()
|
|
fake.plugin_owned = None
|
|
execute_result = MagicMock()
|
|
execute_result.scalar_one_or_none.return_value = fake
|
|
session.execute = AsyncMock(return_value=execute_result)
|
|
session.delete = AsyncMock()
|
|
session.commit = AsyncMock()
|
|
|
|
ok = await db.delete_worker_individual("agent-x")
|
|
assert ok is True
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_upsert_plugin_slot_inserts_when_missing():
|
|
db, session = _make_db()
|
|
execute_result = MagicMock()
|
|
execute_result.scalar_one_or_none.return_value = None # 不存在
|
|
session.execute = AsyncMock(return_value=execute_result)
|
|
session.add = MagicMock()
|
|
session.commit = AsyncMock()
|
|
session.refresh = AsyncMock()
|
|
|
|
row = await db.upsert_plugin_slot(
|
|
plugin_name="data_analytics",
|
|
slot_name="analyst",
|
|
description="数据分析师",
|
|
)
|
|
session.add.assert_called_once()
|
|
added = session.add.call_args[0][0]
|
|
assert added.plugin_owned == "data_analytics"
|
|
assert added.agent_name == "analyst"
|
|
assert added.provider_title == "" # 等用户装配
|
|
assert row is added
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_upsert_plugin_slot_refreshes_when_exists():
|
|
db, session = _make_db()
|
|
fake = MagicMock()
|
|
fake.description = "old"
|
|
fake.node_affinity = "cpu"
|
|
execute_result = MagicMock()
|
|
execute_result.scalar_one_or_none.return_value = fake
|
|
session.execute = AsyncMock(return_value=execute_result)
|
|
session.add = MagicMock()
|
|
session.commit = AsyncMock()
|
|
session.refresh = AsyncMock()
|
|
|
|
await db.upsert_plugin_slot("data_analytics", "analyst", "新描述", node_affinity="gpu")
|
|
assert fake.description == "新描述"
|
|
assert fake.node_affinity == "gpu"
|