k8smanager-cli/tests/test_diagnostic_agent.py

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],
            ["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_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"}'

        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_kubectl", 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":"根据采集结果完成分析"}',
        ])

        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_kubectl", 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_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()