"""工具添加 Episode 实体抽取测试."""
import pytest
from unittest.mock import Mock, MagicMock, patch
from src.tools import handle_tool_call
from src.config_manager import ConfigManager
class TestToolsAddEpisodeEntities:
"""工具添加 Episode 实体抽取测试类."""
@pytest.fixture
def config_manager(self, temp_config_dir):
"""创建配置管理器."""
return ConfigManager(config_path=temp_config_dir / ".graphitiace" / "config.json")
@pytest.mark.asyncio
async def test_add_episode_with_entities_extracted(self, config_manager):
"""测试添加 Episode 并成功抽取实体."""
from src.graphiti_client import GraphitiClient
from unittest.mock import AsyncMock
mock_client = Mock(spec=GraphitiClient)
mock_client.is_connected.return_value = True
mock_client.add_episode = AsyncMock(return_value={
"success": True,
"message": "Episode 添加成功",
"episode_id": 1,
"entities_extracted": True,
"entities": [{"name": "Entity1", "type": "Technology"}]
})
result = await handle_tool_call(
tool_name="add_episode",
arguments={
"content": "我喜欢使用 TypeScript 和 React",
"metadata": {"type": "preference"}
},
config_manager=config_manager,
graphiti_client=mock_client
)
assert result is not None
assert isinstance(result, list)
result_text = "".join([str(item.text) if hasattr(item, 'text') else str(item) for item in result])
assert "成功" in result_text or "success" in result_text.lower()
@pytest.mark.asyncio
async def test_add_episode_without_entities_extracted(self, config_manager):
"""测试添加 Episode 但未抽取实体（需要 Cursor AI 手动完成）."""
from src.graphiti_client import GraphitiClient
from unittest.mock import AsyncMock
mock_client = Mock(spec=GraphitiClient)
mock_client.is_connected.return_value = True
mock_client.add_episode = AsyncMock(return_value={
"success": True,
"message": "Episode 添加成功",
"episode_id": 1,
"entities_extracted": False,
"suggestion": "可以使用 Cursor AI 手动提取实体"
})
result = await handle_tool_call(
tool_name="add_episode",
arguments={
"content": "测试内容",
"metadata": {"type": "test"}
},
config_manager=config_manager,
graphiti_client=mock_client
)
assert result is not None
assert isinstance(result, list)
result_text = "".join([str(item.text) if hasattr(item, 'text') else str(item) for item in result])
assert "成功" in result_text or "success" in result_text.lower()
# 验证包含建议信息
assert "Cursor" in result_text or "建议" in result_text or "suggestion" in result_text.lower() or "💡" in result_text
