96 lines
3.6 KiB
Python
96 lines
3.6 KiB
Python
import pytest
|
|
import pytest_asyncio
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from deerflow.persistence.base import Base
|
|
from deerflow.persistence.skill_knowledge import SkillKnowledgeRepository
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def repository():
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(Base.metadata.create_all)
|
|
yield SkillKnowledgeRepository(async_sessionmaker(engine, expire_on_commit=False))
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_matrix_job_creates_one_item_and_binding_per_pair(repository: SkillKnowledgeRepository) -> None:
|
|
job = await repository.create_job(
|
|
skill_names=["a", "b", "c"],
|
|
targets=[
|
|
{"target_type": "weknora", "target_mode": "wiki", "target_id": "kb-1", "target_name": "一"},
|
|
{"target_type": "assistant", "target_mode": "assistant", "target_id": "base", "target_name": "知识梳理"},
|
|
],
|
|
created_by="admin",
|
|
)
|
|
detail = await repository.get_job(job["id"])
|
|
bindings = await repository.list_bindings()
|
|
|
|
assert detail is not None
|
|
assert detail["total_count"] == 6
|
|
assert len(detail["items"]) == 6
|
|
assert len(bindings) == 6
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_detach_delete_only_removes_binding_contributions(repository: SkillKnowledgeRepository) -> None:
|
|
job = await repository.create_job(
|
|
skill_names=["a", "b"],
|
|
targets=[{"target_type": "assistant", "target_mode": "assistant", "target_id": "base", "target_name": "知识梳理"}],
|
|
created_by="admin",
|
|
review_mode="off",
|
|
)
|
|
detail = await repository.get_job(job["id"])
|
|
assert detail is not None
|
|
binding = (await repository.list_bindings(skill_names=["a"]))[0]
|
|
detached = await repository.detach_binding(binding["id"], delete_contributions=True)
|
|
assert detached and detached["status"] == "detached_deleted"
|
|
assert len(await repository.list_bindings()) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_item_claim_is_atomic_and_records_worker_lease(repository: SkillKnowledgeRepository) -> None:
|
|
job = await repository.create_job(
|
|
skill_names=["a"],
|
|
targets=[
|
|
{
|
|
"target_type": "assistant",
|
|
"target_mode": "assistant",
|
|
"target_id": "base",
|
|
"target_name": "知识梳理",
|
|
}
|
|
],
|
|
created_by="admin",
|
|
)
|
|
detail = await repository.get_job(job["id"])
|
|
item_id = detail["items"][0]["id"]
|
|
|
|
first = await repository.claim_item(item_id, owner="worker-a")
|
|
second = await repository.claim_item(item_id, owner="worker-b")
|
|
|
|
assert first and first["lease_owner"] == "worker-a"
|
|
assert first["attempts"] == 1
|
|
assert second is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_binding_allows_only_one_active_item(repository: SkillKnowledgeRepository) -> None:
|
|
target = {
|
|
"target_type": "assistant",
|
|
"target_mode": "assistant",
|
|
"target_id": "base",
|
|
"target_name": "知识梳理",
|
|
}
|
|
first = await repository.create_job(skill_names=["a"], targets=[target], created_by="admin")
|
|
duplicate = await repository.create_job(skill_names=["a"], targets=[target], created_by="admin")
|
|
assert first["total_count"] == 1
|
|
assert duplicate["total_count"] == 0
|
|
assert duplicate["status"] == "completed"
|
|
|
|
detail = await repository.get_job(first["id"])
|
|
await repository.update_item(detail["items"][0]["id"], status="failed")
|
|
retry = await repository.create_job(skill_names=["a"], targets=[target], created_by="admin")
|
|
assert retry["total_count"] == 1
|