"""ACEManager 策略热力图功能测试."""
import pytest
from unittest.mock import MagicMock, patch
from src.ace_manager import ACEManager
from src.config_manager import ConfigManager
class TestACEManagerStrategyHeatmap:
"""验证 get_strategy_heatmap 的聚合逻辑."""
@pytest.fixture
def config_manager(self, temp_config_dir):
"""创建独立配置管理器."""
return ConfigManager(config_path=temp_config_dir / ".graphitiace" / "config.json")
def _create_manager(self, config_manager, graphiti_client):
"""创建禁用初始化的 ACEManager 以便注入依赖."""
with patch.object(ACEManager, "_initialize_ace", lambda self: None):
manager = ACEManager(config_manager, graphiti_client)
manager.enabled = True
return manager
def test_get_strategy_heatmap_success(self, config_manager):
"""在连接的情况下应返回聚合数据."""
graphiti_client = MagicMock()
graphiti_client.is_connected.return_value = True
session = MagicMock()
session.run.return_value = [
{
"tool_name": "search_entities",
"bucket": "70-100%",
"strategy_count": 2,
"total_usage": 10,
"avg_success_rate": 0.8,
"last_updated": "2025-12-01T00:00:00Z",
},
{
"tool_name": "add_episode",
"bucket": "40-70%",
"strategy_count": 1,
"total_usage": 4,
"avg_success_rate": 0.6,
"last_updated": "2025-12-01T00:00:00Z",
},
]
graphiti_client.driver.session.return_value.__enter__.return_value = session
manager = self._create_manager(config_manager, graphiti_client)
result = manager.get_strategy_heatmap(limit=5, group_by="tool")
assert result["success"] is True
assert len(result["entries"]) == 2
assert result["summary"]["total_usage"] == 14
assert result["summary"]["group_by"] == "tool"
assert result["summary"]["groups"][0]["name"] == "search_entities"
def test_get_strategy_heatmap_not_connected(self, config_manager):
"""未连接 Neo4j 时应返回错误信息."""
graphiti_client = MagicMock()
graphiti_client.is_connected.return_value = False
manager = self._create_manager(config_manager, graphiti_client)
result = manager.get_strategy_heatmap()
assert result["success"] is False
assert "无法生成策略热力图" in result["message"]
