b005fb1ef6
Some checks are pending
release / release (push) Waiting to run
Package layout (zubble/), vendored benchlib.stream_chat, LocalEnv, minimal CLI, 38 pure-logic tests, and a tag-triggered Forgejo Actions release pipeline that builds a single-file PyInstaller executable.
247 lines
9.5 KiB
Python
247 lines
9.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Tests for zubble agent loop orchestration.
|
|
|
|
These tests mock only I/O (stream_chat HTTP calls, env.exec/verify),
|
|
never the orchestration logic that wires detectors, protocol, ledger,
|
|
and abort handling together.
|
|
|
|
python3 -m unittest discover -s tests -v
|
|
"""
|
|
import os
|
|
import sys
|
|
import json
|
|
import time
|
|
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from zubble import benchlib # noqa: E402
|
|
from zubble.agent import run_agent_loop # noqa: E402
|
|
from zubble import protocol # noqa: E402
|
|
|
|
|
|
class _MockEnv:
|
|
"""Minimal env stub matching scenariobench's interface."""
|
|
intro = "You have a root shell on a host."
|
|
report_path = "/root/report.md"
|
|
|
|
def __init__(self, exec_results=None, verify_result=(1, "unhealthy")):
|
|
self._exec_results = exec_results or []
|
|
self._exec_idx = 0
|
|
self._verify_result = verify_result
|
|
|
|
def exec(self, cmd, timeout=90):
|
|
if self._exec_idx < len(self._exec_results):
|
|
r = self._exec_results[self._exec_idx]
|
|
self._exec_idx += 1
|
|
return r
|
|
return (0, "")
|
|
|
|
def verify(self):
|
|
return self._verify_result
|
|
|
|
def setup(self):
|
|
return True, "ok"
|
|
|
|
def teardown(self):
|
|
pass
|
|
|
|
def motd(self):
|
|
return "welcome"
|
|
|
|
def report(self):
|
|
return ""
|
|
|
|
|
|
def _make_stream_response(content):
|
|
"""Build a minimal perf dict + content for a mock stream_chat return."""
|
|
perf = dict(ttft=0.01, t_answer=0.01, wall=0.01,
|
|
reasoning=0, visible=len(content), total=len(content),
|
|
decode=0.0, user_tps=0.0, prompt_tokens=10,
|
|
prefill=0.0, finish="stop")
|
|
return content, "", perf
|
|
|
|
|
|
class AgentLoopActionKillTest(unittest.TestCase):
|
|
"""Verify that repeated identical actions trigger kill → fresh context
|
|
→ reset cap → doom_loop."""
|
|
|
|
def test_action_repetition_leads_to_doom_loop(self):
|
|
repeated_action = json.dumps({
|
|
"action": "run",
|
|
"command": "kubectl get pods",
|
|
"plan_step_id": 1,
|
|
"reasoning": "checking pods",
|
|
})
|
|
|
|
def mock_stream_chat(*args, **kwargs):
|
|
return _make_stream_response(repeated_action)
|
|
|
|
with patch.object(benchlib, "stream_chat", side_effect=mock_stream_chat):
|
|
env = _MockEnv(exec_results=[(0, "NAME READY STATUS\npod-a Running")] * 100)
|
|
deadline = time.monotonic() + 10
|
|
result = run_agent_loop(
|
|
"http://x", env, "Fix the cluster", "Connected.",
|
|
deadline, "off", None, max_turns=200, maxtok=4000, name="test")
|
|
|
|
self.assertEqual(result["outcome"], "doom_loop")
|
|
self.assertTrue(result["turns"] > 0)
|
|
|
|
def test_plan_then_run_then_submit_passes(self):
|
|
responses = [
|
|
json.dumps({"action": "plan", "steps": ["diagnose", "fix", "verify"],
|
|
"reasoning": "3-step plan"}),
|
|
json.dumps({"action": "run", "command": "kubectl get pods",
|
|
"plan_step_id": 1, "reasoning": "step 1"}),
|
|
json.dumps({"action": "run", "command": "kubectl apply -f fix.yaml",
|
|
"plan_step_id": 2, "reasoning": "step 2"}),
|
|
json.dumps({"action": "submit", "report": "Fixed the cert.",
|
|
"reasoning": "done"}),
|
|
]
|
|
idx = [0]
|
|
|
|
def mock_stream_chat(*args, **kwargs):
|
|
r = _make_stream_response(responses[idx[0]])
|
|
idx[0] += 1
|
|
return r
|
|
|
|
with patch.object(benchlib, "stream_chat", side_effect=mock_stream_chat):
|
|
env = _MockEnv(
|
|
exec_results=[(0, "pods running"), (0, "applied")],
|
|
verify_result=(0, "cluster healthy"),
|
|
)
|
|
deadline = time.monotonic() + 10
|
|
result = run_agent_loop(
|
|
"http://x", env, "Fix the cluster", "Connected.",
|
|
deadline, "off", None, max_turns=200, maxtok=4000, name="test")
|
|
|
|
self.assertTrue(result["passed"])
|
|
self.assertEqual(result["outcome"], "passed")
|
|
self.assertEqual(result["submissions"], 1)
|
|
self.assertEqual(result["commands"], 2)
|
|
|
|
def test_agent_requests_constrained_json_schema(self):
|
|
response = json.dumps({"action": "submit", "report": "fixed",
|
|
"reasoning": "verified"})
|
|
|
|
with patch.object(benchlib, "stream_chat",
|
|
return_value=_make_stream_response(response)) as stream_chat:
|
|
env = _MockEnv(verify_result=(0, "cluster healthy"))
|
|
result = run_agent_loop(
|
|
"http://x", env, "Fix the cluster", "Connected.",
|
|
time.monotonic() + 10, "off", None,
|
|
max_turns=200, maxtok=4000, name="test")
|
|
|
|
self.assertTrue(result["passed"])
|
|
self.assertEqual(stream_chat.call_args.kwargs["response_format"],
|
|
protocol.RESPONSE_FORMAT)
|
|
|
|
def test_rejected_submit_continues_until_later_pass(self):
|
|
responses = [
|
|
json.dumps({"action": "plan", "steps": ["check"], "reasoning": "plan"}),
|
|
json.dumps({"action": "run", "command": "kubectl get pods",
|
|
"plan_step_id": 1, "reasoning": "check"}),
|
|
json.dumps({"action": "submit", "report": "first attempt",
|
|
"reasoning": "ready to verify"}),
|
|
json.dumps({"action": "run", "command": "kubectl apply -f fix.yaml",
|
|
"plan_step_id": 1, "reasoning": "apply correction"}),
|
|
json.dumps({"action": "submit", "report": "fix applied",
|
|
"reasoning": "ready to verify again"}),
|
|
]
|
|
idx = [0]
|
|
verify_calls = [0]
|
|
|
|
def mock_stream_chat(*args, **kwargs):
|
|
r = _make_stream_response(responses[idx[0]])
|
|
idx[0] += 1
|
|
return r
|
|
|
|
def verify():
|
|
verify_calls[0] += 1
|
|
return ((1, "still unhealthy") if verify_calls[0] == 1
|
|
else (0, "cluster healthy"))
|
|
|
|
with patch.object(benchlib, "stream_chat", side_effect=mock_stream_chat):
|
|
env = _MockEnv(exec_results=[(0, "pods running"), (0, "applied")])
|
|
env.verify = verify
|
|
deadline = time.monotonic() + 10
|
|
result = run_agent_loop(
|
|
"http://x", env, "Fix the cluster", "Connected.",
|
|
deadline, "off", None, max_turns=200, maxtok=4000, name="test")
|
|
|
|
self.assertTrue(result["passed"])
|
|
self.assertEqual(result["outcome"], "passed")
|
|
self.assertEqual(result["submissions"], 2)
|
|
self.assertEqual(result["rejected"], 1)
|
|
self.assertEqual(result["commands"], 2)
|
|
|
|
def test_malformed_response_then_valid_repair_executes_action(self):
|
|
responses = [
|
|
'{"action":"finished","reasoning":"done"}',
|
|
json.dumps({"action": "submit", "report": "fixed",
|
|
"reasoning": "ready to verify"}),
|
|
]
|
|
idx = [0]
|
|
|
|
def mock_stream_chat(*args, **kwargs):
|
|
r = _make_stream_response(responses[idx[0]])
|
|
idx[0] += 1
|
|
return r
|
|
|
|
with patch.object(benchlib, "stream_chat", side_effect=mock_stream_chat):
|
|
env = _MockEnv(verify_result=(0, "cluster healthy"))
|
|
result = run_agent_loop(
|
|
"http://x", env, "Fix the cluster", "Connected.",
|
|
time.monotonic() + 10, "off", None,
|
|
max_turns=200, maxtok=4000, name="test")
|
|
|
|
self.assertTrue(result["passed"])
|
|
self.assertEqual(result["submissions"], 1)
|
|
self.assertEqual(result["parse_failures"], 1)
|
|
|
|
def test_malformed_response_and_repair_end_as_parse_error(self):
|
|
def mock_stream_chat(*args, **kwargs):
|
|
return _make_stream_response('{"action":"submit","report":"bad\nreport"}')
|
|
|
|
with patch.object(benchlib, "stream_chat", side_effect=mock_stream_chat):
|
|
env = _MockEnv()
|
|
result = run_agent_loop(
|
|
"http://x", env, "Fix the cluster", "Connected.",
|
|
time.monotonic() + 10, "off", None,
|
|
max_turns=200, maxtok=4000, name="test")
|
|
|
|
self.assertFalse(result["passed"])
|
|
self.assertEqual(result["outcome"], "parse_error")
|
|
self.assertEqual(result["turns"], 1)
|
|
self.assertEqual(result["parse_failures"], 2)
|
|
self.assertEqual(result["submissions"], 0)
|
|
|
|
|
|
class ProtocolTest(unittest.TestCase):
|
|
|
|
def test_final_and_finished_are_rejected_with_submit_guidance(self):
|
|
for action in ("final", "finished"):
|
|
with self.subTest(action=action):
|
|
with self.assertRaisesRegex(protocol.ProtocolError,
|
|
"use 'submit'"):
|
|
protocol.parse_response(json.dumps({
|
|
"action": action, "report": "done", "reasoning": "done",
|
|
}))
|
|
|
|
def test_trailing_content_is_rejected(self):
|
|
with self.assertRaisesRegex(protocol.ProtocolError, "trailing content"):
|
|
protocol.parse_response(
|
|
'{"action":"submit","report":"done","reasoning":"done"} trailing')
|
|
|
|
def test_long_escaped_report_is_valid(self):
|
|
response = json.dumps({
|
|
"action": "submit",
|
|
"report": "line one\nline two with {braces} and \\\"quotes\\\"",
|
|
"reasoning": "verified",
|
|
})
|
|
self.assertEqual(protocol.parse_response(response)["action"], "submit")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|