Files
awoooi/apps/api/tests/test_mcp_tool_registry.py
Your Name 7f3722c7f7
Some checks failed
CD Pipeline / tests (push) Successful in 1m22s
Code Review / ai-code-review (push) Successful in 13s
CD Pipeline / build-and-deploy (push) Successful in 4m10s
CD Pipeline / post-deploy-checks (push) Has been cancelled
fix(ai): improve docker repair verification signals
2026-06-01 19:27:36 +08:00

430 lines
18 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""
MCPToolRegistry 測試
====================
ADR-081: Phase 1 動態工具登記冊
測試項目:
- SensorDimension 枚舉完整性
- RegisteredTool dataclass
- register_provider() — 工具登記 / 重複登記防護 / 停用 Provider
- register_tool_manually() — 測試用手動注入
- suggest_tools() — alertname 過濾 / priority 排序 / max_tools 限制
- _classify_tool() — 工具名稱 → 感官維度自動推斷
2026-04-15 Claude Sonnet 4.6 + ogt: Phase 1 初始建立
"""
from __future__ import annotations
import pytest
from src.plugins.mcp.interfaces import MCPTool, MCPToolProvider, MCPToolResult
from src.services.mcp_tool_registry import (
MCPToolRegistry,
SensorDimension,
_classify_tool,
)
# ─────────────────────────────────────────────────────────────────────────────
# Stubs
# ─────────────────────────────────────────────────────────────────────────────
def _make_tool(name: str, server: str = "test") -> MCPTool:
return MCPTool(name=name, description="test tool", input_schema={}, server_name=server)
class _StubProvider(MCPToolProvider):
"""最小可用 Provider stub無外部依賴"""
def __init__(self, name: str, tools: list[str], enabled: bool = True) -> None:
self._name = name
self._tools = [_make_tool(t, name) for t in tools]
self._enabled = enabled
@property
def name(self) -> str:
return self._name
@property
def enabled(self) -> bool:
return self._enabled
async def list_tools(self) -> list[MCPTool]:
return self._tools
async def execute(self, tool_name: str, parameters: dict) -> MCPToolResult:
return MCPToolResult(success=True, execution_id="stub", output={})
# ─────────────────────────────────────────────────────────────────────────────
# SensorDimension
# ─────────────────────────────────────────────────────────────────────────────
class TestSensorDimension:
"""8D 感官維度枚舉必須完整"""
def test_all_8_dimensions_exist(self):
dims = {d.value for d in SensorDimension}
expected = {
"d1_k8s_state", "d2_logs", "d3_metrics", "d4_changes",
"d5_business", "d6_history", "d7_peers", "d8_topology",
}
assert dims == expected
def test_is_str_enum(self):
assert isinstance(SensorDimension.D1_K8S_STATE, str)
assert SensorDimension.D1_K8S_STATE == "d1_k8s_state"
# ─────────────────────────────────────────────────────────────────────────────
# _classify_tool — 自動維度推斷
# ─────────────────────────────────────────────────────────────────────────────
class TestClassifyTool:
"""工具名稱 → 感官維度推斷規則"""
def _provider(self) -> _StubProvider:
return _StubProvider("p", [])
def test_pod_describe_is_d1(self):
reg = _classify_tool(_make_tool("kubectl_describe_pod"), self._provider())
assert SensorDimension.D1_K8S_STATE in reg.dimensions
assert reg.priority == 2
def test_event_is_d1(self):
reg = _classify_tool(_make_tool("kubectl_get_events"), self._provider())
assert SensorDimension.D1_K8S_STATE in reg.dimensions
def test_logs_is_d2(self):
reg = _classify_tool(_make_tool("kubectl_logs"), self._provider())
assert SensorDimension.D2_LOGS in reg.dimensions
assert reg.priority == 2
def test_metrics_is_d3(self):
reg = _classify_tool(_make_tool("prometheus_query"), self._provider())
assert SensorDimension.D3_METRICS in reg.dimensions
assert reg.priority == 3
def test_deploy_diff_is_d4(self):
reg = _classify_tool(_make_tool("git_diff"), self._provider())
assert SensorDimension.D4_CHANGES in reg.dimensions
def test_sli_is_d5(self):
reg = _classify_tool(_make_tool("grafana_sli"), self._provider())
assert SensorDimension.D5_BUSINESS in reg.dimensions
def test_knowledge_is_d6(self):
reg = _classify_tool(_make_tool("rag_knowledge_search"), self._provider())
assert SensorDimension.D6_HISTORY in reg.dimensions
def test_peer_is_d7(self):
reg = _classify_tool(_make_tool("peer_health_check"), self._provider())
assert SensorDimension.D7_PEERS in reg.dimensions
def test_topology_is_d8(self):
reg = _classify_tool(_make_tool("istio_topology"), self._provider())
assert SensorDimension.D8_TOPOLOGY in reg.dimensions
def test_ssh_covers_d1_d2_d3(self):
reg = _classify_tool(_make_tool("ssh_exec"), self._provider())
assert SensorDimension.D1_K8S_STATE in reg.dimensions
assert SensorDimension.D2_LOGS in reg.dimensions
assert SensorDimension.D3_METRICS in reg.dimensions
def test_unknown_tool_defaults_to_d1(self):
reg = _classify_tool(_make_tool("some_unknown_tool"), self._provider())
assert SensorDimension.D1_K8S_STATE in reg.dimensions
def test_kube_type_hints_for_d1(self):
reg = _classify_tool(_make_tool("kubectl_describe"), self._provider())
assert "Kube" in reg.incident_type_hints or any(
h in reg.incident_type_hints for h in ["Pod", "Deploy", "Node"]
)
def test_rollout_tool_is_high_priority_d1(self):
reg = _classify_tool(_make_tool("k8s_watch_rollout"), self._provider())
assert SensorDimension.D1_K8S_STATE in reg.dimensions
assert reg.priority == 1
assert "Deploy" in reg.incident_type_hints
def test_ssh_type_hints(self):
reg = _classify_tool(_make_tool("ssh_run"), self._provider())
assert any(h in reg.incident_type_hints for h in ["Host", "Docker"])
# ─────────────────────────────────────────────────────────────────────────────
# MCPToolRegistry — register_provider
# ─────────────────────────────────────────────────────────────────────────────
class TestRegisterProvider:
"""Provider 登記行為"""
@pytest.mark.asyncio
async def test_register_adds_tools(self):
registry = MCPToolRegistry()
provider = _StubProvider("k8s", ["kubectl_describe", "kubectl_logs"])
count = await registry.register_provider(provider)
assert count == 2
assert registry.tool_count == 2
assert registry.provider_count == 1
@pytest.mark.asyncio
async def test_duplicate_provider_skipped(self):
registry = MCPToolRegistry()
provider = _StubProvider("k8s", ["kubectl_get"])
await registry.register_provider(provider)
# 再次登記同一 Provider
count2 = await registry.register_provider(provider)
assert count2 == 0
assert registry.tool_count == 1 # 不重複
@pytest.mark.asyncio
async def test_disabled_provider_skipped(self):
registry = MCPToolRegistry()
disabled = _StubProvider("ssh", ["ssh_exec"], enabled=False)
count = await registry.register_provider(disabled)
assert count == 0
assert registry.tool_count == 0
@pytest.mark.asyncio
async def test_multiple_providers_accumulate(self):
registry = MCPToolRegistry()
p1 = _StubProvider("k8s", ["kubectl_describe"])
p2 = _StubProvider("ssh", ["ssh_exec"])
await registry.register_provider(p1)
await registry.register_provider(p2)
assert registry.provider_count == 2
assert registry.tool_count == 2
# ─────────────────────────────────────────────────────────────────────────────
# MCPToolRegistry — register_tool_manually
# ─────────────────────────────────────────────────────────────────────────────
class TestRegisterToolManually:
"""手動工具注入(測試用)"""
def test_manual_register_adds_tool(self):
registry = MCPToolRegistry()
provider = _StubProvider("test", [])
tool = _make_tool("custom_tool")
registry.register_tool_manually(
tool=tool,
provider=provider,
dimensions=[SensorDimension.D3_METRICS],
priority=3,
)
assert registry.tool_count == 1
all_tools = registry.get_all_tools()
assert all_tools[0].tool.name == "custom_tool"
assert SensorDimension.D3_METRICS in all_tools[0].dimensions
def test_manual_register_default_priority(self):
registry = MCPToolRegistry()
provider = _StubProvider("test", [])
registry.register_tool_manually(
tool=_make_tool("t"),
provider=provider,
dimensions=[SensorDimension.D1_K8S_STATE],
)
assert registry.get_all_tools()[0].priority == 5
# ─────────────────────────────────────────────────────────────────────────────
# MCPToolRegistry — suggest_tools
# ─────────────────────────────────────────────────────────────────────────────
class TestSuggestTools:
"""動態工具推薦邏輯"""
def _registry_with_tools(self) -> MCPToolRegistry:
registry = MCPToolRegistry()
provider = _StubProvider("test", [])
# K8s 工具(僅適用 Kube* 告警)
k8s_tool = _make_tool("kubectl_describe")
registry.register_tool_manually(
tool=k8s_tool, provider=provider,
dimensions=[SensorDimension.D1_K8S_STATE],
incident_type_hints=["Kube", "Pod"],
priority=2,
)
# 通用工具(所有告警適用)
generic_tool = _make_tool("prometheus_query")
registry.register_tool_manually(
tool=generic_tool, provider=provider,
dimensions=[SensorDimension.D3_METRICS],
incident_type_hints=[],
priority=3,
)
return registry
def test_kube_alertname_gets_k8s_tool(self):
registry = self._registry_with_tools()
tools = registry.suggest_tools(alertname="KubePodCrashLooping")
names = [t.tool.name for t in tools]
assert "kubectl_describe" in names
def test_non_kube_alertname_skips_k8s_tool(self):
registry = self._registry_with_tools()
tools = registry.suggest_tools(alertname="HostDiskUsageHigh")
names = [t.tool.name for t in tools]
assert "kubectl_describe" not in names
assert "prometheus_query" in names
def test_custom_alert_with_deployment_label_gets_k8s_tool(self):
registry = self._registry_with_tools()
tools = registry.suggest_tools(
alertname="AwoooPT16Canary",
incident_labels={"deployment": "awoooi-auto-repair-canary", "namespace": "awoooi-prod"},
)
names = [t.tool.name for t in tools]
assert "kubectl_describe" in names
assert "prometheus_query" in names
def test_namespace_only_non_kube_alert_does_not_get_k8s_tool(self):
registry = self._registry_with_tools()
tools = registry.suggest_tools(
alertname="HostErrorLogFlood",
incident_labels={"namespace": "infra"},
)
names = [t.tool.name for t in tools]
assert "kubectl_describe" not in names
assert "prometheus_query" in names
def test_empty_alertname_gets_generic_only(self):
registry = self._registry_with_tools()
tools = registry.suggest_tools(alertname="")
names = [t.tool.name for t in tools]
assert "prometheus_query" in names
assert "kubectl_describe" not in names
def test_max_tools_limit(self):
registry = MCPToolRegistry()
provider = _StubProvider("test", [])
for i in range(10):
registry.register_tool_manually(
tool=_make_tool(f"tool_{i}"),
provider=provider,
dimensions=[SensorDimension.D3_METRICS],
)
tools = registry.suggest_tools(max_tools=3)
assert len(tools) == 3
def test_priority_ordering(self):
registry = MCPToolRegistry()
provider = _StubProvider("test", [])
registry.register_tool_manually(
tool=_make_tool("low_priority"), provider=provider,
dimensions=[SensorDimension.D3_METRICS], priority=8,
)
registry.register_tool_manually(
tool=_make_tool("high_priority"), provider=provider,
dimensions=[SensorDimension.D1_K8S_STATE], priority=1,
)
tools = registry.suggest_tools()
assert tools[0].tool.name == "high_priority"
assert tools[1].tool.name == "low_priority"
def test_disabled_provider_excluded(self):
registry = MCPToolRegistry()
enabled_p = _StubProvider("enabled", [], enabled=True)
disabled_p = _StubProvider("disabled", [], enabled=False)
registry.register_tool_manually(
tool=_make_tool("disabled_tool"), provider=disabled_p,
dimensions=[SensorDimension.D1_K8S_STATE],
)
registry.register_tool_manually(
tool=_make_tool("enabled_tool"), provider=enabled_p,
dimensions=[SensorDimension.D1_K8S_STATE],
)
tools = registry.suggest_tools()
names = [t.tool.name for t in tools]
assert "enabled_tool" in names
assert "disabled_tool" not in names
def test_empty_registry_returns_empty(self):
registry = MCPToolRegistry()
assert registry.suggest_tools() == []
@pytest.mark.asyncio
async def test_host_error_log_flood_gets_host_observability_tools(self):
registry = MCPToolRegistry()
ssh_provider = _StubProvider(
"ssh_host",
[
"ssh_diagnose",
"ssh_get_top_processes",
"ssh_get_container_logs",
"ssh_get_container_status",
"ssh_get_service_status",
"ssh_check_port",
],
)
k8s_provider = _StubProvider("kubernetes", ["kubectl_describe"])
signoz_provider = _StubProvider("signoz", ["query_logs"])
prometheus_provider = _StubProvider("prometheus", ["prometheus_query"])
await registry.register_provider(k8s_provider)
await registry.register_provider(ssh_provider)
await registry.register_provider(signoz_provider)
await registry.register_provider(prometheus_provider)
tools = registry.suggest_tools(
alertname="HostErrorLogFlood",
incident_labels={
"target": "ollama",
"sensor_host": "ollama",
"sensor_ip": "192.168.0.188",
"host": "192.168.0.188",
},
max_tools=8,
)
names = [reg.tool.name for reg in tools]
assert "ssh_diagnose" in names
assert "ssh_get_top_processes" in names
assert "query_logs" in names
assert "prometheus_query" in names
assert "kubectl_describe" not in names
@pytest.mark.asyncio
async def test_docker_container_alert_does_not_treat_container_as_pod_locator(self):
registry = MCPToolRegistry()
ssh_provider = _StubProvider("ssh_host", ["ssh_diagnose", "ssh_get_container_status"])
k8s_provider = _StubProvider("kubernetes", ["kubectl_describe", "k8s_get_pod_logs"])
prometheus_provider = _StubProvider("prometheus", ["prometheus_query"])
await registry.register_provider(k8s_provider)
await registry.register_provider(ssh_provider)
await registry.register_provider(prometheus_provider)
tools = registry.suggest_tools(
alertname="DockerContainerMemoryLimitPressure",
incident_labels={
"namespace": "default",
"container": "momo-pro-system",
"container_name": "momo-pro-system",
"host": "188",
},
max_tools=8,
)
names = [reg.tool.name for reg in tools]
assert "ssh_diagnose" in names
assert "ssh_get_container_status" in names
assert "prometheus_query" in names
assert "kubectl_describe" not in names
assert "k8s_get_pod_logs" not in names
def test_get_all_tools_returns_all(self):
registry = MCPToolRegistry()
provider = _StubProvider("test", [])
registry.register_tool_manually(
tool=_make_tool("t1"), provider=provider,
dimensions=[SensorDimension.D1_K8S_STATE],
)
registry.register_tool_manually(
tool=_make_tool("t2"), provider=provider,
dimensions=[SensorDimension.D3_METRICS],
)
assert len(registry.get_all_tools()) == 2