import asyncio import unittest from unittest.mock import patch from backend.ai_agent import run_diagnostic_agent async def collect_events(stream): return [event async for event in stream] class DiagnosticAgentTest(unittest.TestCase): def test_executes_read_only_tool_and_returns_final_answer(self): responses = iter([ '{"action":"command","command":"kubectl get pods -A","reason":"查看异常 Pod"}', '{"action":"final","analysis":"发现一个异常 Pod","verified":true,"remaining_issues":0}', ]) async def fake_completion(messages): return next(responses) async def fake_execute(command): return {"command": command, "mode": "read", "returncode": 0, "stdout": "pod-a Running", "stderr": ""} async def run(): with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_command", fake_execute): return await collect_events(run_diagnostic_agent("report", max_steps=3)) events = asyncio.run(run()) self.assertEqual( [event["type"] for event in events], ["command", "result", "thinking", "command", "result", "thinking", "final"], ) self.assertEqual(events[3]["mode"], "read") self.assertEqual(events[-1]["content"], "发现一个异常 Pod") def test_pauses_mutating_command_for_human_approval(self): async def fake_completion(messages): return '{"action":"command","command":"kubectl rollout restart deployment/api -n prod","reason":"重启故障工作负载"}' async def run(): with patch("backend.ai_agent._complete_chat", fake_completion): return await collect_events(run_diagnostic_agent("report", max_steps=3)) events = asyncio.run(run()) self.assertEqual(events[-1]["type"], "approval_required") self.assertEqual(events[-1]["command"], "kubectl rollout restart deployment/api -n prod") self.assertTrue(events[-1]["approval_id"]) def test_manual_mode_never_executes_command_and_waits_for_result(self): async def fake_completion(messages): return '{"action":"command","command":"kubectl get pods -A","reason":"查看 Pod"}' async def run(): with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_command") as execute: events = await collect_events(run_diagnostic_agent("report", max_steps=3, execution_mode="manual")) return events, execute events, execute = asyncio.run(run()) execute.assert_not_called() self.assertEqual(events[-1]["type"], "manual_command") self.assertEqual(events[-1]["command"], "kubectl get pods -A") self.assertTrue(events[-1]["command_id"]) def test_ai_mode_keeps_write_command_behind_approval(self): async def fake_completion(messages): return '{"action":"command","command":"kubectl delete pod api-0 -n prod","reason":"删除故障 Pod"}' executed = [] async def run(): async def fake_execute(command): executed.append(command) return {"command": command, "mode": "read", "returncode": 0, "stdout": "nodes ready", "stderr": ""} with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_command", fake_execute): return await collect_events(run_diagnostic_agent("report", max_steps=3, execution_mode="ai")) events = asyncio.run(run()) self.assertEqual(executed, ["kubectl get nodes -o wide"]) self.assertEqual(events[-1]["type"], "approval_required") def test_rejects_unknown_execution_mode(self): async def run(): return await collect_events(run_diagnostic_agent("report", execution_mode="invalid")) events = asyncio.run(run()) self.assertEqual(events[-1]["type"], "error") self.assertIn("执行模式", events[-1]["message"]) def test_ai_mode_forces_initial_read_only_probe_before_model_planning(self): responses = iter([ '{"action":"final","analysis":"根据采集结果完成分析","verified":true,"remaining_issues":0}', ]) async def fake_completion(messages): return next(responses) executed = [] async def fake_execute(command): executed.append(command) return {"command": command, "mode": "read", "returncode": 0, "stdout": "node-a Ready", "stderr": ""} async def run(): with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_command", fake_execute): return await collect_events(run_diagnostic_agent("report", max_steps=3, execution_mode="ai")) events = asyncio.run(run()) self.assertEqual(executed, ["kubectl get nodes -o wide"]) self.assertEqual([event["type"] for event in events[:2]], ["command", "result"]) self.assertEqual(events[-1]["type"], "final") def test_parse_agent_action_accepts_markdown_json_fence(self): from backend.ai_agent import _parse_agent_action action = _parse_agent_action( '我将先检查 Pod。\n```json\n{"action":"command","command":"kubectl get pods -A","reason":"查看 Pod"}\n```' ) self.assertEqual(action["action"], "command") self.assertEqual(action["command"], "kubectl get pods -A") def test_ai_mode_continues_beyond_six_commands_until_model_finishes(self): command_count = 8 responses = iter([ *[ ( '{"action":"command","command":"kubectl get pods -A ' f'--field-selector metadata.name=pod-{index}","reason":"继续收集证据"}}' ) for index in range(command_count) ], '{"action":"final","analysis":"已完成根因分析","verified":true,"remaining_issues":0}', ]) executed = [] async def fake_completion(messages): return next(responses) async def fake_execute(command): executed.append(command) return {"command": command, "mode": "read", "returncode": 0, "stdout": "evidence", "stderr": ""} async def run(): with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_command", fake_execute): return await collect_events(run_diagnostic_agent("report", execution_mode="ai", max_steps=20)) events = asyncio.run(run()) self.assertEqual(len(executed), command_count + 1) self.assertEqual(events[-1]["type"], "final") self.assertEqual(events[-1]["content"], "已完成根因分析") self.assertNotIn("最大轮次", events[-1]["content"]) def test_safety_limit_reports_agent_error_instead_of_sending_user_to_manual_diagnosis(self): async def fake_completion(messages): return '{"action":"command","command":"kubectl get pods -A","reason":"继续调查"}' async def fake_execute(command): return {"command": command, "mode": "read", "returncode": 0, "stdout": "same result", "stderr": ""} async def run(): with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_command", fake_execute): return await collect_events(run_diagnostic_agent("report", execution_mode="ai", max_steps=2)) events = asyncio.run(run()) self.assertEqual(events[-1]["type"], "error") self.assertIn("安全上限", events[-1]["message"]) self.assertNotIn("人工排查", events[-1]["message"]) def test_agent_rejects_final_until_verification_reports_no_remaining_issues(self): responses = iter([ '{"action":"final","analysis":"已经修复"}', '{"action":"command","command":"kubectl get pods -A","reason":"执行全量复检"}', '{"action":"final","analysis":"所有工作负载恢复正常","verified":true,"remaining_issues":0}', ]) executed = [] async def fake_completion(messages): return next(responses) async def fake_execute(command): executed.append(command) stdout = "node-a Ready" if command == "kubectl get nodes -o wide" else "all pods Running" return {"command": command, "mode": "read", "returncode": 0, "stdout": stdout, "stderr": ""} async def run(): with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_command", fake_execute): return await collect_events(run_diagnostic_agent("report", execution_mode="ai", max_steps=5)) events = asyncio.run(run()) self.assertEqual(executed, ["kubectl get nodes -o wide", "kubectl get pods -A"]) self.assertEqual(events[-1]["type"], "final") self.assertTrue(events[-1]["verified"]) self.assertEqual(events[-1]["remaining_issues"], 0) def test_parse_final_requires_explicit_verification_fields(self): from backend.ai_agent import _parse_agent_action with self.assertRaises(ValueError): _parse_agent_action('{"action":"final","analysis":"看起来正常"}') action = _parse_agent_action( '{"action":"final","analysis":"全部正常","verified":true,"remaining_issues":0}' ) self.assertTrue(action["verified"]) self.assertEqual(action["remaining_issues"], 0) def test_invalid_json_response_is_retried_instead_of_stopping_agent(self): responses = iter([ "我需要继续检查，但没有按格式输出", '{"action":"command","command":"kubectl get pods -A","reason":"继续检查"}', '{"action":"final","analysis":"检查完成","verified":true,"remaining_issues":0}', ]) executed = [] async def fake_completion(messages): return next(responses) async def fake_execute(command): executed.append(command) return {"command": command, "mode": "read", "returncode": 0, "stdout": "healthy", "stderr": ""} async def run(): with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_command", fake_execute): return await collect_events(run_diagnostic_agent("report", max_steps=5)) events = asyncio.run(run()) self.assertTrue(any(event["type"] == "response_rejected" for event in events)) self.assertEqual(events[-1]["type"], "final") self.assertEqual(executed, ["kubectl get nodes -o wide", "kubectl get pods -A"]) def test_stops_after_three_consecutive_invalid_json_responses(self): async def fake_completion(messages): return "still invalid" async def fake_execute(command): return {"command": command, "mode": "read", "returncode": 0, "stdout": "node Ready", "stderr": ""} async def run(): with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_command", fake_execute): return await collect_events(run_diagnostic_agent("report", max_steps=10)) events = asyncio.run(run()) self.assertEqual(events[-1]["type"], "error") self.assertIn("连续 3 次", events[-1]["message"]) def test_stops_when_model_returns_invalid_action(self): async def fake_completion(messages): return "not-json" async def run(): with patch("backend.ai_agent._complete_chat", fake_completion): return await collect_events(run_diagnostic_agent("report", max_steps=3)) events = asyncio.run(run()) self.assertEqual(events[-1]["type"], "error") self.assertIn("JSON", events[-1]["message"]) def test_agent_can_mix_kubectl_and_shell_commands(self): """Agent 应能在同一轮诊断中交替使用 kubectl 和系统诊断 shell 命令。""" responses = iter([ '{"action":"command","command":"kubectl get nodes -o wide","reason":"查看节点状态"}', '{"action":"command","command":"systemctl status kubelet","reason":"检查 kubelet 服务"}', '{"action":"command","command":"journalctl -u kubelet -n 50 --no-pager","reason":"查看 kubelet 日志"}', '{"action":"final","analysis":"kubelet 服务正常，节点就绪","verified":true,"remaining_issues":0}', ]) executed = [] async def fake_completion(messages): return next(responses) async def fake_execute(command): executed.append(command) if command.startswith("kubectl"): stdout = "node-a Ready" elif command.startswith("systemctl"): stdout = "active (running)" else: stdout = "no errors" return {"command": command, "mode": "read", "returncode": 0, "stdout": stdout, "stderr": ""} async def run(): with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_command", fake_execute): return await collect_events(run_diagnostic_agent("report", execution_mode="ai", max_steps=10)) events = asyncio.run(run()) # 初始探测 + 3 条 Agent 命令 = 4 次执行 self.assertEqual( executed, ["kubectl get nodes -o wide", "kubectl get nodes -o wide", "systemctl status kubelet", "journalctl -u kubelet -n 50 --no-pager"], ) # 所有命令都是只读模式 (自动执行) command_events = [e for e in events if e["type"] == "command"] self.assertTrue(all(e["mode"] == "read" for e in command_events)) self.assertEqual(events[-1]["type"], "final") def test_agent_rejects_non_whitelisted_shell_command(self): """Agent 应拒绝不在白名单内的 shell 命令 (如 rm)。""" responses = iter([ '{"action":"command","command":"rm -rf /tmp/test","reason":"清理测试文件"}', '{"action":"final","analysis":"已完成","verified":true,"remaining_issues":0}', ]) executed = [] async def fake_completion(messages): return next(responses) async def fake_execute(command): executed.append(command) return {"command": command, "mode": "read", "returncode": 0, "stdout": "", "stderr": ""} async def run(): with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_command", fake_execute): return await collect_events(run_diagnostic_agent("report", execution_mode="ai", max_steps=5)) events = asyncio.run(run()) # rm 命令应被拒绝 rejected = [e for e in events if e["type"] == "command_rejected"] self.assertTrue(rejected) self.assertIn("不允许执行该命令", rejected[0]["reason"]) # rm 不应被执行 self.assertNotIn("rm -rf /tmp/test", executed) self.assertEqual(events[-1]["type"], "final") if __name__ == "__main__": unittest.main()