cnbugs
49 lines
1.8 KiB
Python
49 lines
1.8 KiB
Python
import asyncio
|
|
import json
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import backend.main as main
|
|
|
|
|
|
def parse_sse(event):
|
|
return json.loads(event[6:].strip())
|
|
|
|
|
|
class AgentApiTest(unittest.TestCase):
|
|
def test_agent_stream_forwards_agent_events(self):
|
|
async def fake_agent(report, question="", max_steps=6, execution_mode="ai"):
|
|
self.assertEqual(execution_mode, "manual")
|
|
yield {"type": "thinking", "step": 1, "message": "规划"}
|
|
yield {"type": "final", "content": "完成"}
|
|
|
|
async def run():
|
|
main._last_diagnosis = {"score": 60}
|
|
main._last_report = "report"
|
|
with patch.object(main, "run_diagnostic_agent", fake_agent):
|
|
response = await main.agent_diagnose_stream(question="检查 Pod", execution_mode="manual")
|
|
return [parse_sse(event) async for event in response.body_iterator]
|
|
|
|
events = asyncio.run(run())
|
|
self.assertEqual(events[0]["type"], "start")
|
|
self.assertEqual(events[-1]["type"], "done")
|
|
self.assertEqual(events[-2]["type"], "final")
|
|
|
|
def test_execute_approved_command_once(self):
|
|
approval = {"command": "kubectl scale deployment/api --replicas=2", "reason": "恢复副本"}
|
|
|
|
async def fake_execute(command):
|
|
return {"command": command, "mode": "write", "returncode": 0, "stdout": "scaled", "stderr": ""}
|
|
|
|
async def run():
|
|
with patch.object(main, "take_pending_approval", return_value=approval), patch.object(main, "execute_kubectl", fake_execute):
|
|
return await main.execute_approved_agent_command(main.AgentApprovalRequest(approval_id="abc", approved=True))
|
|
|
|
result = asyncio.run(run())
|
|
self.assertEqual(result["returncode"], 0)
|
|
self.assertEqual(result["stdout"], "scaled")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|