cnbugs
98 lines
4.2 KiB
Python
98 lines
4.2 KiB
Python
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"}',
|
|
])
|
|
|
|
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_kubectl", 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], ["thinking", "command", "result", "thinking", "final"])
|
|
self.assertEqual(events[1]["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_kubectl") 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"}'
|
|
|
|
async def run():
|
|
with patch("backend.ai_agent._complete_chat", fake_completion), patch("backend.ai_agent.execute_kubectl") as execute:
|
|
events = await collect_events(run_diagnostic_agent("report", max_steps=3, execution_mode="ai"))
|
|
return events, execute
|
|
|
|
events, execute = asyncio.run(run())
|
|
execute.assert_not_called()
|
|
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_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"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|