63 lines
1.9 KiB
Python
63 lines
1.9 KiB
Python
import pytest
|
|
|
|
from deerflow.skill_knowledge.graph_writer import (
|
|
EDGE_QUERY,
|
|
NODE_QUERY,
|
|
execute_plan,
|
|
parameterized_statements,
|
|
plan_from_ir,
|
|
)
|
|
|
|
|
|
class RecordingTransaction:
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple[str, dict]] = []
|
|
|
|
async def run(self, query: str, parameters: dict) -> None:
|
|
self.calls.append((query, parameters))
|
|
|
|
|
|
def _plan():
|
|
return plan_from_ir(
|
|
{
|
|
"entities": [
|
|
{"local_id": "a", "type": "person", "name": "张三", "confidence": 1},
|
|
{"local_id": "b", "type": "system", "name": "知识库", "confidence": 1},
|
|
],
|
|
"relations": [
|
|
{
|
|
"source_local_id": "a",
|
|
"target_local_id": "b",
|
|
"predicate": "uses",
|
|
"confidence": 0.9,
|
|
"evidence": [{"file": "SKILL.md", "line": 2}],
|
|
}
|
|
],
|
|
},
|
|
target_type="weknora",
|
|
target_id="kb-1",
|
|
binding_id="binding-1",
|
|
snapshot_id="snapshot-1",
|
|
job_id="job-1",
|
|
)
|
|
|
|
|
|
def test_direct_graph_plan_uses_only_fixed_parameterized_templates() -> None:
|
|
statements = parameterized_statements(_plan())
|
|
assert len(statements) == 3
|
|
assert {row["statement"] for row in statements} == {NODE_QUERY, EDGE_QUERY}
|
|
assert all("parameters" in row for row in statements)
|
|
assert all("张三" not in row["statement"] for row in statements)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_direct_graph_writer_obeys_explicit_enable_switch() -> None:
|
|
transaction = RecordingTransaction()
|
|
disabled = await execute_plan(transaction, _plan(), enabled=False)
|
|
assert disabled == {"nodes": 0, "edges": 0}
|
|
assert transaction.calls == []
|
|
|
|
enabled = await execute_plan(transaction, _plan(), enabled=True)
|
|
assert enabled == {"nodes": 2, "edges": 1}
|
|
assert len(transaction.calls) == 3
|