Coverage for agentos/tests/test_subagent_parent_child.py: 0%
250 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-09 10:19 +0800
1"""测试 SubAgent 父子通信 — 状态共享、心跳、生命周期管理。"""
3import asyncio
5import pytest
7pytestmark = pytest.mark.asyncio
8from agentos.subagent import ( # noqa: E402
9 ChildContext,
10 ChildHandle,
11 ChildHeartbeat,
12 ChildStatus,
13 SharedState,
14 SubAgentManager,
15 SubAgentSpec,
16)
19class TestSharedState:
20 async def test_set_get(self):
21 ss = SharedState()
22 await ss.set("key1", "val1")
23 assert await ss.get("key1") == "val1"
24 assert await ss.get("missing", "def") == "def"
26 async def test_update_snapshot(self):
27 ss = SharedState()
28 await ss.update({"a": 1, "b": 2})
29 snap = await ss.snapshot()
30 assert snap == {"a": 1, "b": 2}
32 async def test_sync_ops(self):
33 ss = SharedState()
34 ss.set_sync("x", 42)
35 assert ss.get_sync("x") == 42
37 async def test_concurrent_writes(self):
38 ss = SharedState()
40 async def writer(key: str, n: int):
41 for i in range(n):
42 await ss.set(key, i)
43 await asyncio.sleep(0)
45 await asyncio.gather(writer("a", 20), writer("b", 20))
46 assert await ss.get("a") == 19
47 assert await ss.get("b") == 19
50class TestChildContext:
51 async def test_progress_report(self):
52 hbs = []
54 async def hb_cb(hb: ChildHeartbeat):
55 hbs.append(hb)
57 ctx = ChildContext("test-1", heartbeat_callback=hb_cb)
58 await ctx.report_progress(0.5, "step1", "half done")
59 assert ctx.progress == 0.5
60 assert len(hbs) == 1
61 assert hbs[0].progress == 0.5
62 assert hbs[0].current_step == "step1"
64 async def test_step_and_heartbeat(self):
65 hbs = []
67 async def hb_cb(hb: ChildHeartbeat):
68 hbs.append(hb)
70 ctx = ChildContext("test-2", heartbeat_callback=hb_cb)
71 await ctx.step(1, "init")
72 await ctx.send_heartbeat("alive")
73 assert len(hbs) == 1
74 assert hbs[0].iteration == 1
76 async def test_done(self):
77 hbs = []
79 async def hb_cb(hb: ChildHeartbeat):
80 hbs.append(hb)
82 ctx = ChildContext("test-3", heartbeat_callback=hb_cb)
83 await ctx.done("all good")
84 assert len(hbs) == 1
85 assert hbs[0].status == ChildStatus.COMPLETED
86 assert hbs[0].progress == 1.0
87 assert hbs[0].message == "all good"
89 async def test_fail(self):
90 hbs = []
92 async def hb_cb(hb: ChildHeartbeat):
93 hbs.append(hb)
95 ctx = ChildContext("test-4", heartbeat_callback=hb_cb)
96 await ctx.fail("something broke")
97 assert len(hbs) == 1
98 assert hbs[0].status == ChildStatus.FAILED
99 assert hbs[0].message == "something broke"
101 async def test_cancel_detection(self):
102 cancelled = [False]
104 def on_cancel():
105 return cancelled[0]
107 ctx = ChildContext("test-5", on_cancel=on_cancel)
108 assert not ctx.cancelled
109 status = await ctx.check_control()
110 assert status == ChildStatus.RUNNING
112 cancelled[0] = True
113 status = await ctx.check_control()
114 assert status == ChildStatus.CANCELLED
115 assert ctx.cancelled
117 async def test_pause_resume(self):
118 paused = [True]
119 resume_triggered = [False]
121 async def on_pause():
122 resume_triggered[0] = True
123 paused[0] = False
125 ctx = ChildContext("test-6", on_pause=on_pause)
126 ctx._paused = paused[0]
127 status = await ctx.check_control()
128 assert status == ChildStatus.PAUSED
129 assert resume_triggered[0]
132class TestChildHandle:
133 async def test_create_context(self):
134 handle = ChildHandle("h1", "do stuff", "fork")
135 ctx = handle.create_context()
136 assert ctx.agent_id == "h1"
137 assert handle.context is ctx
138 assert handle.shared_state is ctx.shared_state
140 async def test_pause_resume(self):
141 handle = ChildHandle("h2", "task", "fork")
142 handle.create_context()
143 assert handle.status == ChildStatus.IDLE
145 await handle.pause()
146 assert handle.status == ChildStatus.PAUSED
148 await handle.resume()
149 assert handle.status == ChildStatus.RUNNING
151 async def test_cancel(self):
152 handle = ChildHandle("h3", "task", "fork")
153 handle.create_context()
154 await handle.cancel()
155 assert handle.status == ChildStatus.CANCELLED
157 async def test_get_status(self):
158 handle = ChildHandle("h4", "analyze", "fork")
159 handle.create_context()
160 handle.info.progress = 0.7
161 handle.info.current_step = "parsing"
162 handle.info.iterations = 12
164 status = handle.get_status()
165 assert status["agent_id"] == "h4"
166 assert status["progress"] == 0.7
167 assert status["current_step"] == "parsing"
168 assert status["iterations"] == 12
169 assert "elapsed" in status
171 async def test_timeout_detection(self):
172 handle = ChildHandle("h5", "task", "fork", timeout=0.1)
173 await asyncio.sleep(0.15)
174 assert handle.check_timeout()
176 async def test_no_timeout_when_unset(self):
177 handle = ChildHandle("h6", "task", "fork", timeout=None)
178 assert not handle.check_timeout()
180 async def test_heartbeat_timeout(self):
181 handle = ChildHandle("h7", "task", "fork", heartbeat_interval=0.1)
182 await asyncio.sleep(0.35)
183 assert handle.check_heartbeat_timeout()
185 async def test_heartbeat_updates_info(self):
186 handle = ChildHandle("h8", "task", "fork")
187 handle.create_context()
188 await handle._receive_heartbeat(
189 ChildHeartbeat(
190 agent_id="h8",
191 progress=0.5,
192 current_step="s1",
193 message="working",
194 iteration=5,
195 )
196 )
197 assert handle.info.progress == 0.5
198 assert handle.info.current_step == "s1"
199 assert handle.info.iterations == 5
201 async def test_shared_state_parent_child(self):
202 handle = ChildHandle("h9", "task", "fork")
203 ctx = handle.create_context()
205 await ctx.shared_state.set("data", [1, 2, 3])
206 val = await handle.shared_state.get("data")
207 assert val == [1, 2, 3]
209 await handle.shared_state.set("status", "ok")
210 assert await ctx.shared_state.get("status") == "ok"
213class TestSubAgentManager:
214 async def test_spawn_fork_with_child_context(self):
216 async def run_func(spec: SubAgentSpec, ctx: ChildContext):
217 await ctx.report_progress(0.3, "init")
218 await ctx.step(1, "load")
219 await ctx.report_progress(0.7, "process")
220 await ctx.done("success")
221 return ("success", 2)
223 mgr = SubAgentManager()
224 result = await mgr.spawn_fork("test task", run_func=run_func)
225 assert result.output == "success"
226 assert result.iterations == 2
227 assert result.handle is not None
228 assert result.handle.status == ChildStatus.COMPLETED
230 async def test_spawn_fork_failure(self):
231 async def run_func(spec, ctx):
232 await ctx.report_progress(0.1, "start")
233 raise ValueError("boom")
235 mgr = SubAgentManager()
236 result = await mgr.spawn_fork("bad task", run_func=run_func)
237 assert result.error == "boom"
238 assert result.handle.status == ChildStatus.FAILED
239 assert result.handle.info.error == "boom"
241 async def test_spawn_fork_pause_resume_flow(self):
242 """模拟父子协作：父暂停→子暂停→父恢复→子继续→完成。"""
243 state = {"phase": "init"}
245 async def run_func(spec, ctx: ChildContext):
246 state["phase"] = "running"
247 await ctx.report_progress(0.2, "step1")
249 # 检查控制信号
250 status = await ctx.check_control()
251 if status == ChildStatus.PAUSED:
252 state["phase"] = "paused"
254 # 再次检查（模拟恢复后继续）
255 status = await ctx.check_control()
256 if status == ChildStatus.RUNNING:
257 state["phase"] = "resumed"
259 await ctx.done("ok")
260 return ("ok", 3)
262 mgr = SubAgentManager()
264 # 启动
265 task = asyncio.create_task(mgr.spawn_fork("pause test", run_func=run_func))
267 await asyncio.sleep(0.05) # 让子Agent跑到 step
268 mgr.get_handle(task.result().handle.agent_id) if hasattr(task, "result") else None
270 # 等task完成
271 result = await task
272 assert result.error is None or result.error == ""
273 assert state["phase"] in ("running", "paused", "resumed")
275 async def test_swarm_parallel(self):
276 results_log = []
278 async def run_func(spec: SubAgentSpec, ctx: ChildContext):
279 await ctx.report_progress(0.5, spec.task)
280 await asyncio.sleep(0.01)
281 await ctx.done(f"done_{spec.task}")
282 results_log.append(spec.task)
283 return (f"done_{spec.task}", 1)
285 mgr = SubAgentManager()
286 results = await mgr.spawn_swarm(["A", "B", "C"], run_func=run_func)
287 assert len(results) == 3
288 assert len(results_log) == 3
289 for r in results:
290 assert r.handle is not None
291 assert r.handle.status == ChildStatus.COMPLETED
293 async def test_cancel_all(self):
294 async def run_func(spec, ctx: ChildContext):
295 await ctx.report_progress(0.1, "init")
296 for i in range(50):
297 await ctx.step(i, f"step_{i}")
298 status = await ctx.check_control()
299 if status == ChildStatus.CANCELLED:
300 return ("cancelled", i)
301 await asyncio.sleep(0.01)
302 return ("done", 50)
304 mgr = SubAgentManager()
305 t1 = asyncio.create_task(mgr.spawn_fork("long task 1", run_func=run_func))
306 t2 = asyncio.create_task(mgr.spawn_fork("long task 2", run_func=run_func))
308 await asyncio.sleep(0.05)
309 await mgr.cancel_all()
311 r1, r2 = await asyncio.gather(t1, t2)
312 assert r1.handle.status == ChildStatus.CANCELLED
313 assert r2.handle.status == ChildStatus.CANCELLED
315 async def test_list_children(self):
316 mgr = SubAgentManager()
317 r = await mgr.spawn_fork("task1")
318 children = mgr.list_children()
319 assert len(children) == 1
320 assert children[0]["agent_id"] == r.agent_id
322 async def test_cleanup(self):
323 mgr = SubAgentManager()
324 await mgr.spawn_fork("cleanup test")
325 assert len(mgr._agents) == 1
327 cleaned = await mgr.cleanup(max_age_seconds=-1.0)
328 assert cleaned == 1
329 assert len(mgr._agents) == 0
331 async def test_heartbeat_monitoring(self):
332 mgr = SubAgentManager()
333 handle = ChildHandle("hb-test", "task", "fork", timeout=0.05)
334 mgr._agents["hb-test"] = handle
335 handle.info.status = ChildStatus.RUNNING
337 monitor = asyncio.create_task(mgr.monitor_heartbeats(interval=0.02))
338 await asyncio.sleep(0.1)
339 monitor.cancel()
340 try:
341 await monitor
342 except asyncio.CancelledError:
343 pass
345 assert handle.status in (ChildStatus.TIMEOUT, ChildStatus.RUNNING)