393 lines
15 KiB
Python
393 lines
15 KiB
Python
"""
|
||
B3: LLM 動態 Telegram 按鈕 — 單元測試
|
||
=====================================
|
||
2026-04-27 Claude Sonnet 4.6: ADR-082 B3
|
||
2026-04-27 Claude Sonnet 4.6: H3+M6 Fix — short_id Redis 映射 + critical 過濾
|
||
2026-04-27 Claude Sonnet 4.6: P0 Fix — async setex 在 return 前完成（消除 race）
|
||
|
||
測試範圍：
|
||
1. USE_LLM_DYNAMIC_BUTTONS=false → 走 YAML 路徑（現有行為）
|
||
2. USE_LLM_DYNAMIC_BUTTONS=true + recommended_actions 空 → 走 YAML 路徑
|
||
3. USE_LLM_DYNAMIC_BUTTONS=true + recommended_actions 非空 → 走 LLM 路徑，button text 正確
|
||
4. high risk action 前綴 ⚠️
|
||
5. callback_data 格式是 la:{16-hex-chars}，不是 JSON（H3 Fix，P0 後改 16-hex）
|
||
6. [批准][拒絕] 第一排永遠存在
|
||
7. 超過 2 個 action → 多排
|
||
8. critical action 不出按鈕（M6 Fix）
|
||
9. Redis short_id 被寫入（H3 Fix）
|
||
|
||
注意：P0 Fix 後 _build_llm_action_buttons 改為 async（直接 await redis.setex）。
|
||
所有 TestBuildLlmActionButtons 測試改為 @pytest.mark.asyncio。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass, field
|
||
from typing import Literal
|
||
from unittest.mock import AsyncMock, MagicMock, patch
|
||
|
||
import pytest
|
||
|
||
|
||
# =============================================================================
|
||
# 輕量 Stub：避免載入完整 telegram_gateway（需要 Redis / httpx / OTEL）
|
||
# =============================================================================
|
||
|
||
@dataclass
|
||
class _RecommendedAction:
|
||
name: str
|
||
label: str
|
||
emoji: str
|
||
mcp_provider: str
|
||
mcp_tool: str
|
||
params: dict
|
||
risk: Literal["low", "medium", "high", "critical"]
|
||
reasoning: str = ""
|
||
|
||
|
||
@dataclass
|
||
class _ActionPlan:
|
||
recommended_actions: list[_RecommendedAction] = field(default_factory=list)
|
||
|
||
|
||
# =============================================================================
|
||
# 直接測試 _build_llm_action_buttons 靜態方法
|
||
# =============================================================================
|
||
|
||
def _import_builder():
|
||
"""
|
||
延遲 import，確保環境變數先設定。
|
||
返回 TelegramGateway._build_llm_action_buttons 靜態方法。
|
||
"""
|
||
# 使用 importlib 動態取得，避免頂層 import 觸發 settings 初始化
|
||
import importlib
|
||
import sys
|
||
|
||
# 清除快取以確保 env 設定生效
|
||
for mod_name in list(sys.modules.keys()):
|
||
if "telegram_gateway" in mod_name:
|
||
del sys.modules[mod_name]
|
||
|
||
mod = importlib.import_module("src.services.telegram_gateway")
|
||
return mod.TelegramGateway._build_llm_action_buttons, mod
|
||
|
||
|
||
# =============================================================================
|
||
# Fixtures
|
||
# =============================================================================
|
||
|
||
@pytest.fixture
|
||
def low_action():
|
||
return _RecommendedAction(
|
||
name="check_pod_logs",
|
||
label="查 Pod 日誌",
|
||
emoji="📋",
|
||
mcp_provider="k8s",
|
||
mcp_tool="get_pod_logs",
|
||
params={"pod": "{labels.pod}"},
|
||
risk="low",
|
||
)
|
||
|
||
|
||
@pytest.fixture
|
||
def high_action():
|
||
return _RecommendedAction(
|
||
name="restart_pod",
|
||
label="重啟 Pod",
|
||
emoji="🔄",
|
||
mcp_provider="k8s",
|
||
mcp_tool="restart_pod",
|
||
params={"pod": "{labels.pod}"},
|
||
risk="high",
|
||
)
|
||
|
||
|
||
@pytest.fixture
|
||
def medium_action():
|
||
return _RecommendedAction(
|
||
name="scale_deployment",
|
||
label="擴容",
|
||
emoji="📈",
|
||
mcp_provider="k8s",
|
||
mcp_tool="scale_deployment",
|
||
params={"replicas": "3"},
|
||
risk="medium",
|
||
)
|
||
|
||
|
||
# =============================================================================
|
||
# Test 1: _build_llm_action_buttons — 基本 text 格式
|
||
# =============================================================================
|
||
|
||
class TestBuildLlmActionButtons:
|
||
"""
|
||
測試 _build_llm_action_buttons 實例方法。
|
||
|
||
H3 Fix 後 callback_data 格式改為 la:{8-hex-chars}，不再是 JSON。
|
||
Redis 寫入透過 asyncio.ensure_future 背景執行，測試中 mock 掉。
|
||
"""
|
||
|
||
def _make_gateway(self):
|
||
"""建立最小 gateway 實例（不初始化 Redis / httpx）。"""
|
||
from src.services.telegram_gateway import TelegramGateway
|
||
gw = object.__new__(TelegramGateway)
|
||
gw._security = MagicMock()
|
||
return gw
|
||
|
||
async def _call_builder(self, gw, actions):
|
||
"""呼叫 _build_llm_action_buttons（async），mock 掉 Redis 寫入。"""
|
||
mock_redis = MagicMock()
|
||
mock_redis.setex = AsyncMock()
|
||
with patch("src.services.telegram_gateway.get_redis", return_value=mock_redis):
|
||
return await gw._build_llm_action_buttons(actions)
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_low_risk_text_format(self, low_action):
|
||
gw = self._make_gateway()
|
||
rows = await self._call_builder(gw, [low_action])
|
||
assert len(rows) == 1
|
||
btn = rows[0][0]
|
||
assert btn["text"] == "📋 查 Pod 日誌"
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_high_risk_prefix(self, high_action):
|
||
"""Test 4: high risk 前綴 ⚠️"""
|
||
gw = self._make_gateway()
|
||
rows = await self._call_builder(gw, [high_action])
|
||
btn = rows[0][0]
|
||
assert btn["text"].startswith("⚠️")
|
||
assert "重啟 Pod" in btn["text"]
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_medium_risk_no_prefix(self, medium_action):
|
||
gw = self._make_gateway()
|
||
rows = await self._call_builder(gw, [medium_action])
|
||
btn = rows[0][0]
|
||
assert not btn["text"].startswith("⚠️")
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_callback_data_format_is_short_id(self, low_action):
|
||
"""Test 5 (H3 Fix + P0): callback_data 格式是 la:{16-hex-chars}，不是 JSON。
|
||
P0 Fix 後 short_id 改為 secrets.token_hex(8) = 16-hex-chars，≤19 bytes。
|
||
"""
|
||
import re
|
||
gw = self._make_gateway()
|
||
rows = await self._call_builder(gw, [low_action])
|
||
cb = rows[0][0]["callback_data"]
|
||
# 格式必須是 la: 開頭，後接 16 位 hex（P0 Fix 後從 8 → 16 hex chars）
|
||
assert re.fullmatch(r"la:[0-9a-f]{16}", cb), f"格式不符: {cb!r}"
|
||
# 確認不是 JSON
|
||
assert not cb.startswith("{")
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_callback_data_within_19_bytes(self, low_action):
|
||
"""H3 Fix + P0: la:{16-hex-chars} 固定 ≤19 bytes，絕不截斷。"""
|
||
gw = self._make_gateway()
|
||
rows = await self._call_builder(gw, [low_action])
|
||
cb = rows[0][0]["callback_data"]
|
||
assert len(cb.encode("utf-8")) <= 19, f"callback_data 超出 19 bytes: {cb!r} ({len(cb.encode())} bytes)"
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_critical_action_skipped(self):
|
||
"""Test 8 (M6 Fix): critical risk action 不生成按鈕。"""
|
||
critical_action = _RecommendedAction(
|
||
name="nuke_cluster",
|
||
label="核爆叢集",
|
||
emoji="💣",
|
||
mcp_provider="k8s",
|
||
mcp_tool="delete_all_pods",
|
||
params={},
|
||
risk="critical",
|
||
)
|
||
gw = self._make_gateway()
|
||
rows = await self._call_builder(gw, [critical_action])
|
||
assert rows == [], f"critical action 不應生成按鈕，但得到 {rows}"
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_critical_mixed_with_others(self, low_action, medium_action):
|
||
"""M6 Fix: critical 被過濾，其他 action 正常生成。"""
|
||
critical_action = _RecommendedAction(
|
||
name="nuke_cluster",
|
||
label="核爆",
|
||
emoji="💣",
|
||
mcp_provider="k8s",
|
||
mcp_tool="delete_all_pods",
|
||
params={},
|
||
risk="critical",
|
||
)
|
||
gw = self._make_gateway()
|
||
rows = await self._call_builder(gw, [low_action, critical_action, medium_action])
|
||
# 只剩 low + medium 兩個，同排
|
||
assert len(rows) == 1
|
||
assert len(rows[0]) == 2
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_redis_short_id_written(self, low_action):
|
||
"""Test 9 (H3 Fix + P0): Redis tg:la:{short_id} 被 await setex 寫入（TTL=3600）。
|
||
P0 Fix 後改為直接 await redis.setex，不再使用 ensure_future。
|
||
"""
|
||
mock_redis = MagicMock()
|
||
mock_redis.setex = AsyncMock()
|
||
|
||
gw = self._make_gateway()
|
||
|
||
with patch("src.services.telegram_gateway.get_redis", return_value=mock_redis):
|
||
rows = await gw._build_llm_action_buttons([low_action])
|
||
|
||
assert len(rows) == 1
|
||
cb = rows[0][0]["callback_data"]
|
||
short_id = cb[3:] # 去掉 "la:"
|
||
|
||
# P0 Fix: 直接 await，所以 setex 被呼叫一次
|
||
mock_redis.setex.assert_called_once()
|
||
call_args = mock_redis.setex.call_args
|
||
assert call_args[0][0] == f"tg:la:{short_id}"
|
||
assert call_args[0][1] == 3600
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_two_actions_same_row(self, low_action, medium_action):
|
||
"""Test 7 前置：2 個 action → 同一排"""
|
||
gw = self._make_gateway()
|
||
rows = await self._call_builder(gw, [low_action, medium_action])
|
||
assert len(rows) == 1
|
||
assert len(rows[0]) == 2
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_three_actions_two_rows(self, low_action, medium_action, high_action):
|
||
"""Test 7: 超過 2 個 action → 多排（每排最多 2 個）"""
|
||
gw = self._make_gateway()
|
||
rows = await self._call_builder(gw, [low_action, medium_action, high_action])
|
||
assert len(rows) == 2
|
||
assert len(rows[0]) == 2
|
||
assert len(rows[1]) == 1
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_empty_actions_returns_empty(self):
|
||
gw = self._make_gateway()
|
||
rows = await self._call_builder(gw, [])
|
||
assert rows == []
|
||
|
||
|
||
# =============================================================================
|
||
# Test 2-3, 6: _build_inline_keyboard 路徑切換
|
||
# =============================================================================
|
||
|
||
class TestBuildInlineKeyboardRouting:
|
||
"""
|
||
測試 _build_inline_keyboard() 的路徑選擇邏輯。
|
||
用 MagicMock 模擬 security，避免初始化完整 gateway。
|
||
"""
|
||
|
||
def _make_gateway_cls(self):
|
||
from src.services.telegram_gateway import TelegramGateway
|
||
return TelegramGateway
|
||
|
||
def _make_mock_gateway(self):
|
||
cls = self._make_gateway_cls()
|
||
gw = object.__new__(cls)
|
||
mock_security = MagicMock()
|
||
mock_security.generate_callback_nonce.side_effect = (
|
||
lambda approval_id, action: f"nonce:{approval_id}:{action}"
|
||
)
|
||
gw._security = mock_security
|
||
return gw
|
||
|
||
def _first_row_texts(self, keyboard: dict) -> list[str]:
|
||
return [btn["text"] for btn in keyboard["inline_keyboard"][0]]
|
||
|
||
def _callback_data_values(self, keyboard: dict, *, start_row: int = 0) -> list[str]:
|
||
return [
|
||
btn["callback_data"]
|
||
for row in keyboard["inline_keyboard"][start_row:]
|
||
for btn in row
|
||
if "callback_data" in btn
|
||
]
|
||
|
||
async def _build_kb(self, gw, **kwargs) -> dict:
|
||
"""await _build_inline_keyboard，mock 掉 Redis。"""
|
||
mock_redis = MagicMock()
|
||
mock_redis.setex = AsyncMock()
|
||
with patch("src.services.telegram_gateway.get_redis", return_value=mock_redis):
|
||
return await gw._build_inline_keyboard(**kwargs)
|
||
|
||
# Test 1: flag=false → YAML 路徑（第一排還是批准/拒絕）
|
||
@pytest.mark.asyncio
|
||
@patch("src.services.telegram_gateway.USE_LLM_DYNAMIC_BUTTONS", False)
|
||
@patch("src.services.callback_dispatcher.list_actions_for_category", return_value=[])
|
||
async def test_flag_false_uses_yaml_path(self, mock_list, low_action):
|
||
gw = self._make_mock_gateway()
|
||
action_plan = _ActionPlan(recommended_actions=[low_action])
|
||
kb = await self._build_kb(gw, approval_id="APR-001", incident_id="INC-001", action_plan=action_plan)
|
||
first_texts = self._first_row_texts(kb)
|
||
assert "✅ 批准" in first_texts
|
||
assert "❌ 拒絕" in first_texts
|
||
# LLM 按鈕不應出現（la: 前綴代表 llm_action）
|
||
all_cbs = self._callback_data_values(kb)
|
||
assert not any(cb.startswith("la:") for cb in all_cbs)
|
||
|
||
# Test 2: flag=true + actions 空 → YAML fallback
|
||
@pytest.mark.asyncio
|
||
@patch("src.services.telegram_gateway.USE_LLM_DYNAMIC_BUTTONS", True)
|
||
@patch("src.services.callback_dispatcher.list_actions_for_category", return_value=[])
|
||
async def test_flag_true_empty_actions_uses_yaml(self, mock_list):
|
||
gw = self._make_mock_gateway()
|
||
action_plan = _ActionPlan(recommended_actions=[])
|
||
kb = await self._build_kb(gw, approval_id="APR-002", incident_id="INC-002", action_plan=action_plan)
|
||
first_texts = self._first_row_texts(kb)
|
||
assert "✅ 批准" in first_texts
|
||
assert "❌ 拒絕" in first_texts
|
||
all_cbs = self._callback_data_values(kb)
|
||
assert not any(cb.startswith("la:") for cb in all_cbs)
|
||
|
||
# Test 2b: flag=true + action_plan=None → YAML fallback
|
||
@pytest.mark.asyncio
|
||
@patch("src.services.telegram_gateway.USE_LLM_DYNAMIC_BUTTONS", True)
|
||
@patch("src.services.callback_dispatcher.list_actions_for_category", return_value=[])
|
||
async def test_flag_true_no_action_plan_uses_yaml(self, mock_list):
|
||
gw = self._make_mock_gateway()
|
||
kb = await self._build_kb(gw, approval_id="APR-003", incident_id="INC-003")
|
||
first_texts = self._first_row_texts(kb)
|
||
assert "✅ 批准" in first_texts
|
||
assert "❌ 拒絕" in first_texts
|
||
|
||
# Test 3: flag=true + actions 非空 → LLM 路徑
|
||
@pytest.mark.asyncio
|
||
@patch("src.services.telegram_gateway.USE_LLM_DYNAMIC_BUTTONS", True)
|
||
async def test_flag_true_with_actions_uses_llm(self, low_action):
|
||
import re
|
||
gw = self._make_mock_gateway()
|
||
action_plan = _ActionPlan(recommended_actions=[low_action])
|
||
kb = await self._build_kb(gw, approval_id="APR-004", incident_id="INC-004", action_plan=action_plan)
|
||
# Test 6: 第一排永遠是 [批准][拒絕]
|
||
first_texts = self._first_row_texts(kb)
|
||
assert "✅ 批准" in first_texts
|
||
assert "❌ 拒絕" in first_texts
|
||
|
||
# LLM 按鈕在第二排以後（P0 Fix 後 16-hex-chars）
|
||
all_cbs = self._callback_data_values(kb, start_row=1)
|
||
assert any(re.fullmatch(r"la:[0-9a-f]{16}", cb) for cb in all_cbs), (
|
||
f"LLM 按鈕 callback_data 格式應為 la:{{16-hex}}，實際: {all_cbs}"
|
||
)
|
||
|
||
# button text 正確
|
||
all_texts = [btn["text"] for row in kb["inline_keyboard"][1:] for btn in row]
|
||
assert any("查 Pod 日誌" in t for t in all_texts)
|
||
|
||
# Test 6: 第一排永遠是 [批准][拒絕]（LLM 路徑）
|
||
@pytest.mark.asyncio
|
||
@patch("src.services.telegram_gateway.USE_LLM_DYNAMIC_BUTTONS", True)
|
||
async def test_approve_reject_always_first_row_in_llm_path(self, low_action, high_action, medium_action):
|
||
gw = self._make_mock_gateway()
|
||
action_plan = _ActionPlan(recommended_actions=[low_action, high_action, medium_action])
|
||
kb = await self._build_kb(gw, approval_id="APR-005", incident_id="INC-005", action_plan=action_plan)
|
||
first_texts = self._first_row_texts(kb)
|
||
assert first_texts[0] == "✅ 批准"
|
||
assert first_texts[1] == "❌ 拒絕"
|
||
# callback 操作排數 = 1 (approve/reject) + 2 (3 actions, 2 per row)；
|
||
# AwoooP evidence URL 另加一排，沒有 callback_data。
|
||
callback_rows = [
|
||
row for row in kb["inline_keyboard"]
|
||
if any("callback_data" in btn for btn in row)
|
||
]
|
||
assert len(callback_rows) == 3
|