test_work_graph_execute.py
9.75 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
"""U-006/U-012 tests: c_work_graph_execute's dynamic-workflow wrapper (U-006)
plus CLASSIFIER_BLOCKED / USAGE_LIMIT_REACHED detection and the
p_usage_limit_wait retry loop (U-012), mocked at the subprocess boundary
(CLAUDE.md Testing constraint). Real end-to-end coverage is in
src/test_work_graph_execute_e2e.py (RUN_REAL_AGENT_E2E=1), reusing the
exact spike script the A-1 assumption was validated against.
Run: python3 -m unittest src.test_work_graph_execute -v
"""
import json
import os
import subprocess
import tempfile
import unittest
from unittest.mock import patch
from src import work_graph_execute
def _completed_process(returncode=0, stdout_obj=None, stderr=""):
stdout = json.dumps(stdout_obj) if stdout_obj is not None else ""
return subprocess.CompletedProcess(args=["claude"], returncode=returncode, stdout=stdout, stderr=stderr)
class WriteWorkflowScriptTest(unittest.TestCase):
def test_writes_into_target_claude_workflows_not_graphlemon_repo(self):
with tempfile.TemporaryDirectory() as tmp:
path = work_graph_execute.write_workflow_script(tmp, "run-123", "export const meta = {}")
self.assertEqual(path, os.path.join(tmp, ".claude", "workflows", "run-123.js"))
with open(path, encoding="utf-8") as f:
self.assertEqual(f.read(), "export const meta = {}")
class ExecuteWorkGraphTest(unittest.TestCase):
SCHEMA = {"type": "object", "required": ["sum"], "properties": {"sum": {"type": "number"}}}
def test_happy_path_returns_structured_output(self):
with patch("src.work_graph_execute.subprocess.run") as run:
run.return_value = _completed_process(
stdout_obj={"is_error": False, "structured_output": {"sum": 29}}
)
output = work_graph_execute.execute_work_graph(
"/tmp/some-project", "run-123", {"numbers": [2, 3, 4]}, self.SCHEMA
)
self.assertEqual(output, {"sum": 29})
def test_invokes_claude_with_expected_args_and_cwd(self):
with patch("src.work_graph_execute.subprocess.run") as run:
run.return_value = _completed_process(stdout_obj={"is_error": False, "structured_output": {"sum": 1}})
work_graph_execute.execute_work_graph("/tmp/some-project", "run-123", {"n": 1}, self.SCHEMA)
args, kwargs = run.call_args
argv = args[0]
self.assertEqual(argv[0], "claude")
self.assertIn("-p", argv)
self.assertIn("Run /run-123 with args {\"n\": 1}", argv)
self.assertIn("--json-schema", argv)
self.assertIn(json.dumps(self.SCHEMA), argv)
self.assertIn("--permission-mode", argv)
self.assertIn("auto", argv)
self.assertNotIn("--bare", argv)
self.assertNotIn("--dangerously-skip-permissions", argv)
self.assertNotIn("bypassPermissions", argv)
self.assertEqual(kwargs["cwd"], "/tmp/some-project")
self.assertEqual(kwargs["stdin"], subprocess.DEVNULL)
def test_nonzero_exit_raises(self):
with patch("src.work_graph_execute.subprocess.run") as run:
run.return_value = _completed_process(returncode=1, stdout_obj=None, stderr="boom")
with self.assertRaises(work_graph_execute.WorkGraphExecuteError):
work_graph_execute.execute_work_graph("/tmp/some-project", "run-123", {}, self.SCHEMA)
def test_is_error_true_raises(self):
with patch("src.work_graph_execute.subprocess.run") as run:
run.return_value = _completed_process(stdout_obj={"is_error": True})
with self.assertRaises(work_graph_execute.WorkGraphExecuteError):
work_graph_execute.execute_work_graph("/tmp/some-project", "run-123", {}, self.SCHEMA)
def test_missing_structured_output_raises(self):
with patch("src.work_graph_execute.subprocess.run") as run:
run.return_value = _completed_process(stdout_obj={"is_error": False})
with self.assertRaises(work_graph_execute.WorkGraphExecuteError):
work_graph_execute.execute_work_graph("/tmp/some-project", "run-123", {}, self.SCHEMA)
class ClassifierBlockedDetectionTest(unittest.TestCase):
SCHEMA = {"type": "object"}
def test_three_or_more_permission_denials_raises_classifier_blocked(self):
with patch("src.work_graph_execute.subprocess.run") as run:
run.return_value = _completed_process(
stdout_obj={"is_error": True, "permission_denials": ["a", "b", "c"]}
)
with self.assertRaises(work_graph_execute.ClassifierBlockedError):
work_graph_execute.execute_work_graph("/tmp/p", "run-1", {}, self.SCHEMA)
def test_fewer_than_three_denials_is_generic_error_not_classifier_blocked(self):
with patch("src.work_graph_execute.subprocess.run") as run:
run.return_value = _completed_process(
stdout_obj={"is_error": True, "permission_denials": ["a"]}
)
with self.assertRaises(work_graph_execute.WorkGraphExecuteError) as ctx:
work_graph_execute.execute_work_graph("/tmp/p", "run-1", {}, self.SCHEMA)
self.assertNotIsInstance(ctx.exception, work_graph_execute.ClassifierBlockedError)
class UsageLimitDetectionTest(unittest.TestCase):
SCHEMA = {"type": "object"}
def test_session_limit_phrase_in_stderr_on_nonzero_exit(self):
with patch("src.work_graph_execute.subprocess.run") as run:
run.return_value = _completed_process(
returncode=1, stdout_obj=None, stderr="Error: You've hit your session limit"
)
with self.assertRaises(work_graph_execute.UsageLimitReachedError) as ctx:
work_graph_execute.execute_work_graph("/tmp/p", "run-1", {}, self.SCHEMA)
self.assertEqual(ctx.exception.quota_limit_type, "session")
def test_weekly_limit_phrase_in_error_result(self):
with patch("src.work_graph_execute.subprocess.run") as run:
run.return_value = _completed_process(
stdout_obj={"is_error": True, "result": "You've hit your weekly limit"}
)
with self.assertRaises(work_graph_execute.UsageLimitReachedError) as ctx:
work_graph_execute.execute_work_graph("/tmp/p", "run-1", {}, self.SCHEMA)
self.assertEqual(ctx.exception.quota_limit_type, "weekly")
class RetryWrapperTest(unittest.TestCase):
SCHEMA = {"type": "object"}
def setUp(self):
self.sleeps = []
self.clock = [0.0]
def _sleep(self, seconds):
self.sleeps.append(seconds)
self.clock[0] += seconds
def _now(self):
return self.clock[0]
def test_recovers_after_backoff_no_timeout_hit(self):
responses = [
_completed_process(returncode=1, stdout_obj=None, stderr="hit your session limit"),
_completed_process(stdout_obj={"is_error": False, "structured_output": {"ok": True}}),
]
with patch("src.work_graph_execute.subprocess.run", side_effect=responses):
result = work_graph_execute.execute_work_graph_with_usage_limit_retry(
"/tmp/p", "run-1", {}, self.SCHEMA, sleep=self._sleep, now=self._now
)
self.assertEqual(result, {"ok": True})
self.assertEqual(self.sleeps, [5 * 60]) # first backoff step
def test_backoff_schedule_escalates_then_holds_at_30(self):
# 5 failures before success -> backoff steps 5,10,20,30,30 minutes
responses = [
_completed_process(returncode=1, stdout_obj=None, stderr="hit your session limit")
] * 5 + [_completed_process(stdout_obj={"is_error": False, "structured_output": {"ok": True}})]
with patch("src.work_graph_execute.subprocess.run", side_effect=responses):
work_graph_execute.execute_work_graph_with_usage_limit_retry(
"/tmp/p", "run-1", {}, self.SCHEMA, sleep=self._sleep, now=self._now
)
self.assertEqual(self.sleeps, [5 * 60, 10 * 60, 20 * 60, 30 * 60, 30 * 60])
def test_session_timeout_after_8_hours_raises_timeout_error(self):
def always_fails(*args, **kwargs):
return _completed_process(returncode=1, stdout_obj=None, stderr="hit your session limit")
with patch("src.work_graph_execute.subprocess.run", side_effect=always_fails):
with self.assertRaises(work_graph_execute.UsageLimitWaitTimeoutError):
work_graph_execute.execute_work_graph_with_usage_limit_retry(
"/tmp/p", "run-1", {}, self.SCHEMA, sleep=self._sleep, now=self._now
)
self.assertGreaterEqual(self.clock[0], 8 * 3600)
def test_opus_quota_raises_timeout_error_immediately_no_backoff(self):
with patch("src.work_graph_execute.subprocess.run") as run:
run.return_value = _completed_process(
returncode=1, stdout_obj=None, stderr="opus model quota exhausted"
)
with self.assertRaises(work_graph_execute.UsageLimitWaitTimeoutError):
work_graph_execute.execute_work_graph_with_usage_limit_retry(
"/tmp/p", "run-1", {}, self.SCHEMA, sleep=self._sleep, now=self._now
)
self.assertEqual(self.sleeps, [])
def test_classifier_blocked_is_not_retried(self):
with patch("src.work_graph_execute.subprocess.run") as run:
run.return_value = _completed_process(
stdout_obj={"is_error": True, "permission_denials": ["a", "b", "c"]}
)
with self.assertRaises(work_graph_execute.ClassifierBlockedError):
work_graph_execute.execute_work_graph_with_usage_limit_retry(
"/tmp/p", "run-1", {}, self.SCHEMA, sleep=self._sleep, now=self._now
)
self.assertEqual(self.sleeps, [])
self.assertEqual(run.call_count, 1)
if __name__ == "__main__":
unittest.main()