329 lines
11 KiB
Python
329 lines
11 KiB
Python
"""
|
||
Test cases for Step-KB Binding feature.
|
||
[Step-KB-Binding] 步骤关联知识库功能的测试用例
|
||
|
||
测试覆盖：
|
||
1. 步骤配置的增删改查与参数校验
|
||
2. 配置步骤KB范围后，检索仅在范围内发生
|
||
3. 未配置时回退原逻辑
|
||
4. 多知识库同名内容场景下，步骤约束生效
|
||
5. trace 字段完整性校验
|
||
"""
|
||
|
||
import pytest
|
||
from unittest.mock import AsyncMock, MagicMock, patch
|
||
from dataclasses import dataclass
|
||
from typing import Any
|
||
|
||
|
||
class TestStepKbBindingModel:
|
||
"""测试步骤KB绑定数据模型"""
|
||
|
||
def test_flow_step_with_kb_binding_fields(self):
|
||
"""测试 FlowStep 包含 KB 绑定字段"""
|
||
from app.models.entities import FlowStep
|
||
|
||
step = FlowStep(
|
||
step_no=1,
|
||
content="测试步骤",
|
||
allowed_kb_ids=["kb-1", "kb-2"],
|
||
preferred_kb_ids=["kb-1"],
|
||
kb_query_hint="查找产品相关信息",
|
||
max_kb_calls_per_step=2,
|
||
)
|
||
|
||
assert step.allowed_kb_ids == ["kb-1", "kb-2"]
|
||
assert step.preferred_kb_ids == ["kb-1"]
|
||
assert step.kb_query_hint == "查找产品相关信息"
|
||
assert step.max_kb_calls_per_step == 2
|
||
|
||
def test_flow_step_without_kb_binding(self):
|
||
"""测试 FlowStep 不配置 KB 绑定时的默认值"""
|
||
from app.models.entities import FlowStep
|
||
|
||
step = FlowStep(
|
||
step_no=1,
|
||
content="测试步骤",
|
||
)
|
||
|
||
assert step.allowed_kb_ids is None
|
||
assert step.preferred_kb_ids is None
|
||
assert step.kb_query_hint is None
|
||
assert step.max_kb_calls_per_step is None
|
||
|
||
def test_max_kb_calls_validation(self):
|
||
"""测试 max_kb_calls_per_step 的范围校验"""
|
||
from app.models.entities import FlowStep
|
||
from pydantic import ValidationError
|
||
|
||
# 有效范围 1-5
|
||
step = FlowStep(step_no=1, content="test", max_kb_calls_per_step=3)
|
||
assert step.max_kb_calls_per_step == 3
|
||
|
||
# 超出上限
|
||
with pytest.raises(Exception): # ValidationError
|
||
FlowStep(step_no=1, content="test", max_kb_calls_per_step=10)
|
||
|
||
|
||
class TestStepKbConfig:
|
||
"""测试 StepKbConfig 数据类"""
|
||
|
||
def test_step_kb_config_creation(self):
|
||
"""测试 StepKbConfig 创建"""
|
||
from app.services.mid.kb_search_dynamic_tool import StepKbConfig
|
||
|
||
config = StepKbConfig(
|
||
allowed_kb_ids=["kb-1", "kb-2"],
|
||
preferred_kb_ids=["kb-1"],
|
||
kb_query_hint="查找产品信息",
|
||
max_kb_calls=2,
|
||
step_id="flow-1_step_1",
|
||
)
|
||
|
||
assert config.allowed_kb_ids == ["kb-1", "kb-2"]
|
||
assert config.preferred_kb_ids == ["kb-1"]
|
||
assert config.kb_query_hint == "查找产品信息"
|
||
assert config.max_kb_calls == 2
|
||
assert config.step_id == "flow-1_step_1"
|
||
|
||
def test_step_kb_config_defaults(self):
|
||
"""测试 StepKbConfig 默认值"""
|
||
from app.services.mid.kb_search_dynamic_tool import StepKbConfig
|
||
|
||
config = StepKbConfig()
|
||
|
||
assert config.allowed_kb_ids is None
|
||
assert config.preferred_kb_ids is None
|
||
assert config.kb_query_hint is None
|
||
assert config.max_kb_calls == 1
|
||
assert config.step_id is None
|
||
|
||
|
||
class TestKbSearchDynamicToolWithStepConfig:
|
||
"""测试 KbSearchDynamicTool 与步骤配置的集成"""
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_kb_search_with_allowed_kb_ids(self):
|
||
"""测试配置 allowed_kb_ids 后检索范围受限"""
|
||
from app.services.mid.kb_search_dynamic_tool import (
|
||
KbSearchDynamicTool,
|
||
KbSearchDynamicConfig,
|
||
StepKbConfig,
|
||
)
|
||
|
||
mock_session = MagicMock()
|
||
mock_timeout_governor = MagicMock()
|
||
|
||
tool = KbSearchDynamicTool(
|
||
session=mock_session,
|
||
timeout_governor=mock_timeout_governor,
|
||
config=KbSearchDynamicConfig(enabled=True),
|
||
)
|
||
|
||
step_config = StepKbConfig(
|
||
allowed_kb_ids=["kb-allowed-1", "kb-allowed-2"],
|
||
step_id="test_step",
|
||
)
|
||
|
||
with patch.object(tool, '_do_retrieve', new_callable=AsyncMock) as mock_retrieve:
|
||
mock_retrieve.return_value = [
|
||
{"id": "1", "content": "test", "score": 0.8, "metadata": {"kb_id": "kb-allowed-1"}}
|
||
]
|
||
|
||
result = await tool.execute(
|
||
query="测试查询",
|
||
tenant_id="tenant-1",
|
||
step_kb_config=step_config,
|
||
)
|
||
|
||
# 验证检索调用时传入了正确的 kb_ids
|
||
call_args = mock_retrieve.call_args
|
||
assert call_args[1]['step_kb_config'] == step_config
|
||
|
||
# 验证返回结果包含 step_kb_binding 信息
|
||
assert result.step_kb_binding is not None
|
||
assert result.step_kb_binding['allowed_kb_ids'] == ["kb-allowed-1", "kb-allowed-2"]
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_kb_search_without_step_config(self):
|
||
"""测试未配置步骤KB时的回退行为"""
|
||
from app.services.mid.kb_search_dynamic_tool import (
|
||
KbSearchDynamicTool,
|
||
KbSearchDynamicConfig,
|
||
)
|
||
|
||
mock_session = MagicMock()
|
||
mock_timeout_governor = MagicMock()
|
||
|
||
tool = KbSearchDynamicTool(
|
||
session=mock_session,
|
||
timeout_governor=mock_timeout_governor,
|
||
config=KbSearchDynamicConfig(enabled=True),
|
||
)
|
||
|
||
with patch.object(tool, '_do_retrieve', new_callable=AsyncMock) as mock_retrieve:
|
||
mock_retrieve.return_value = [
|
||
{"id": "1", "content": "test", "score": 0.8, "metadata": {}}
|
||
]
|
||
|
||
result = await tool.execute(
|
||
query="测试查询",
|
||
tenant_id="tenant-1",
|
||
)
|
||
|
||
# 验证检索调用时未传入 step_kb_config
|
||
call_args = mock_retrieve.call_args
|
||
assert call_args[1]['step_kb_config'] is None
|
||
|
||
# 验证返回结果不包含 step_kb_binding
|
||
assert result.step_kb_binding is None
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_kb_search_result_includes_used_kb_ids(self):
|
||
"""测试检索结果包含实际使用的知识库ID"""
|
||
from app.services.mid.kb_search_dynamic_tool import (
|
||
KbSearchDynamicTool,
|
||
KbSearchDynamicConfig,
|
||
StepKbConfig,
|
||
)
|
||
|
||
mock_session = MagicMock()
|
||
mock_timeout_governor = MagicMock()
|
||
|
||
tool = KbSearchDynamicTool(
|
||
session=mock_session,
|
||
timeout_governor=mock_timeout_governor,
|
||
config=KbSearchDynamicConfig(enabled=True),
|
||
)
|
||
|
||
step_config = StepKbConfig(
|
||
allowed_kb_ids=["kb-1", "kb-2"],
|
||
step_id="test_step",
|
||
)
|
||
|
||
with patch.object(tool, '_do_retrieve', new_callable=AsyncMock) as mock_retrieve:
|
||
mock_retrieve.return_value = [
|
||
{"id": "1", "content": "test1", "score": 0.9, "metadata": {"kb_id": "kb-1"}},
|
||
{"id": "2", "content": "test2", "score": 0.8, "metadata": {"kb_id": "kb-1"}},
|
||
{"id": "3", "content": "test3", "score": 0.7, "metadata": {"kb_id": "kb-2"}},
|
||
]
|
||
|
||
result = await tool.execute(
|
||
query="测试查询",
|
||
tenant_id="tenant-1",
|
||
step_kb_config=step_config,
|
||
)
|
||
|
||
# 验证 used_kb_ids 包含所有命中的知识库
|
||
assert result.step_kb_binding is not None
|
||
assert set(result.step_kb_binding['used_kb_ids']) == {"kb-1", "kb-2"}
|
||
assert result.step_kb_binding['kb_hit'] is True
|
||
|
||
|
||
class TestTraceInfoStepKbBinding:
|
||
"""测试 TraceInfo 中的 step_kb_binding 字段"""
|
||
|
||
def test_trace_info_with_step_kb_binding(self):
|
||
"""测试 TraceInfo 包含 step_kb_binding 字段"""
|
||
from app.models.mid.schemas import TraceInfo, ExecutionMode
|
||
|
||
trace = TraceInfo(
|
||
mode=ExecutionMode.AGENT,
|
||
step_kb_binding={
|
||
"step_id": "flow-1_step_2",
|
||
"allowed_kb_ids": ["kb-1", "kb-2"],
|
||
"used_kb_ids": ["kb-1"],
|
||
"kb_hit": True,
|
||
},
|
||
)
|
||
|
||
assert trace.step_kb_binding is not None
|
||
assert trace.step_kb_binding['step_id'] == "flow-1_step_2"
|
||
assert trace.step_kb_binding['allowed_kb_ids'] == ["kb-1", "kb-2"]
|
||
assert trace.step_kb_binding['used_kb_ids'] == ["kb-1"]
|
||
|
||
def test_trace_info_without_step_kb_binding(self):
|
||
"""测试 TraceInfo 默认不包含 step_kb_binding"""
|
||
from app.models.mid.schemas import TraceInfo, ExecutionMode
|
||
|
||
trace = TraceInfo(mode=ExecutionMode.AGENT)
|
||
|
||
assert trace.step_kb_binding is None
|
||
|
||
|
||
class TestFlowStepKbBindingIntegration:
|
||
"""测试流程步骤与KB绑定的集成"""
|
||
|
||
def test_script_flow_steps_with_kb_binding(self):
|
||
"""测试 ScriptFlow 的 steps 包含 KB 绑定配置"""
|
||
from app.models.entities import ScriptFlowCreate
|
||
|
||
flow_create = ScriptFlowCreate(
|
||
name="测试流程",
|
||
steps=[
|
||
{
|
||
"step_no": 1,
|
||
"content": "步骤1",
|
||
"allowed_kb_ids": ["kb-1"],
|
||
"preferred_kb_ids": None,
|
||
"kb_query_hint": "查找产品信息",
|
||
"max_kb_calls_per_step": 2,
|
||
},
|
||
{
|
||
"step_no": 2,
|
||
"content": "步骤2",
|
||
# 不配置 KB 绑定
|
||
},
|
||
],
|
||
)
|
||
|
||
assert flow_create.steps[0]['allowed_kb_ids'] == ["kb-1"]
|
||
assert flow_create.steps[0]['kb_query_hint'] == "查找产品信息"
|
||
assert flow_create.steps[1].get('allowed_kb_ids') is None
|
||
|
||
|
||
class TestKbBindingLogging:
|
||
"""测试 KB 绑定的日志记录"""
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_step_kb_config_logging(self, caplog):
|
||
"""测试步骤KB配置的日志记录"""
|
||
import logging
|
||
from app.services.mid.kb_search_dynamic_tool import (
|
||
KbSearchDynamicTool,
|
||
KbSearchDynamicConfig,
|
||
StepKbConfig,
|
||
)
|
||
|
||
mock_session = MagicMock()
|
||
mock_timeout_governor = MagicMock()
|
||
|
||
tool = KbSearchDynamicTool(
|
||
session=mock_session,
|
||
timeout_governor=mock_timeout_governor,
|
||
config=KbSearchDynamicConfig(enabled=True),
|
||
)
|
||
|
||
step_config = StepKbConfig(
|
||
allowed_kb_ids=["kb-1"],
|
||
step_id="flow-1_step_1",
|
||
)
|
||
|
||
with patch.object(tool, '_do_retrieve', new_callable=AsyncMock) as mock_retrieve:
|
||
mock_retrieve.return_value = []
|
||
|
||
with caplog.at_level(logging.INFO):
|
||
await tool.execute(
|
||
query="测试",
|
||
tenant_id="tenant-1",
|
||
step_kb_config=step_config,
|
||
)
|
||
|
||
# 验证日志包含 Step-KB-Binding 标记
|
||
assert any("Step-KB-Binding" in record.message for record in caplog.records)
|
||
|
||
|
||
# 运行测试的入口
|
||
if __name__ == "__main__":
|
||
pytest.main([__file__, "-v"])
|