Coverage for agentos/tests/test_task_scheduler.py: 0%

348 statements  

« prev     ^ index     » next       coverage.py v7.14.3, created at 2026-07-08 13:14 +0800

1"""Tests for agentos.tools.task_scheduler — Task, PriorityTaskQueue, TaskScheduler, WorkerPool.""" 

2 

3import pytest 

4import threading 

5import time 

6from agentos.tools.task_scheduler import ( 

7 PriorityTaskQueue, 

8 Task, 

9 TaskScheduler, 

10 TaskStatus, 

11 WorkerPool, 

12) 

13 

14 

15# ============================================================================ 

16# TaskStatus 

17# ============================================================================ 

18 

19class TestTaskStatus: 

20 def test_enum_members(self): 

21 assert TaskStatus.PENDING is not None 

22 assert TaskStatus.RUNNING is not None 

23 assert TaskStatus.COMPLETED is not None 

24 assert TaskStatus.FAILED is not None 

25 assert TaskStatus.EXPIRED is not None 

26 assert TaskStatus.CANCELLED is not None 

27 

28 def test_unique_values(self): 

29 values = [s.value for s in TaskStatus] 

30 assert len(values) == len(set(values)) 

31 

32 

33# ============================================================================ 

34# Task 

35# ============================================================================ 

36 

37class TestTask: 

38 def test_creation_defaults(self): 

39 t = Task(task_id="t1", func=lambda: 1) 

40 assert t.task_id == "t1" 

41 assert t.args == () 

42 assert t.kwargs == {} 

43 assert t.priority == 0 

44 assert t.deadline is None 

45 assert t.status == TaskStatus.PENDING 

46 assert t.result is None 

47 assert t.error is None 

48 assert t.started_at is None 

49 assert t.completed_at is None 

50 

51 def test_creation_with_args(self): 

52 t = Task(task_id="t2", func=lambda x: x, args=(5,), kwargs={"k": "v"}, 

53 priority=3, deadline=100.0) 

54 assert t.args == (5,) 

55 assert t.kwargs == {"k": "v"} 

56 assert t.priority == 3 

57 assert t.deadline == 100.0 

58 

59 def test_run_success(self): 

60 t = Task(task_id="t3", func=lambda a, b: a + b, args=(1, 2)) 

61 t.run() 

62 assert t.status == TaskStatus.COMPLETED 

63 assert t.result == 3 

64 assert t.error is None 

65 assert t.started_at is not None 

66 assert t.completed_at is not None 

67 

68 def test_run_failure(self): 

69 def boom(): 

70 raise ValueError("kaboom") 

71 

72 t = Task(task_id="t4", func=boom) 

73 t.run() 

74 assert t.status == TaskStatus.FAILED 

75 assert t.result is None 

76 assert t.error == "kaboom" 

77 assert t.completed_at is not None 

78 

79 def test_elapsed_returns_float(self): 

80 t = Task(task_id="t5", func=lambda: 1) 

81 assert t.elapsed is None # not started 

82 t.run() 

83 assert isinstance(t.elapsed, float) 

84 assert t.elapsed >= 0 

85 

86 def test_elapsed_before_start(self): 

87 t = Task(task_id="t6", func=lambda: 1) 

88 assert t.elapsed is None 

89 

90 

91# ============================================================================ 

92# PriorityTaskQueue — Basic Operations 

93# ============================================================================ 

94 

95class TestPriorityTaskQueueBasic: 

96 def test_creation_defaults(self): 

97 q = PriorityTaskQueue() 

98 assert q.size == 0 

99 assert q._max_size == 1000 

100 

101 def test_creation_custom_size(self): 

102 q = PriorityTaskQueue(max_size=10) 

103 assert q._max_size == 10 

104 

105 def test_enqueue_dequeue(self): 

106 q = PriorityTaskQueue() 

107 t = Task(task_id="t1", func=lambda: 42) 

108 assert q.enqueue(t) is True 

109 assert q.size == 1 

110 

111 result = q.dequeue() 

112 assert result is t 

113 assert q.size == 0 

114 

115 def test_dequeue_empty(self): 

116 q = PriorityTaskQueue() 

117 assert q.dequeue() is None 

118 

119 def test_enqueue_full(self): 

120 q = PriorityTaskQueue(max_size=2) 

121 q.enqueue(Task(task_id="a", func=lambda: 1)) 

122 q.enqueue(Task(task_id="b", func=lambda: 2)) 

123 assert q.enqueue(Task(task_id="c", func=lambda: 3)) is False 

124 assert q.size == 2 

125 

126 def test_priority_ordering(self): 

127 q = PriorityTaskQueue() 

128 t1 = Task(task_id="low", func=lambda: 1, priority=10) 

129 t2 = Task(task_id="high", func=lambda: 2, priority=1) 

130 q.enqueue(t1) 

131 q.enqueue(t2) 

132 

133 first = q.dequeue() 

134 second = q.dequeue() 

135 assert first.task_id == "high" 

136 assert second.task_id == "low" 

137 

138 def test_same_priority_fifo(self): 

139 q = PriorityTaskQueue() 

140 t1 = Task(task_id="first", func=lambda: 1, priority=5) 

141 t2 = Task(task_id="second", func=lambda: 2, priority=5) 

142 q.enqueue(t1) 

143 q.enqueue(t2) 

144 

145 assert q.dequeue().task_id == "first" 

146 assert q.dequeue().task_id == "second" 

147 

148 

149# ============================================================================ 

150# PriorityTaskQueue — Cancellation 

151# ============================================================================ 

152 

153class TestPriorityTaskQueueCancel: 

154 def test_cancel_existing(self): 

155 q = PriorityTaskQueue() 

156 t = Task(task_id="x", func=lambda: 1) 

157 q.enqueue(t) 

158 assert q.cancel("x") is True 

159 assert q.size == 1 # still on heap (cleaned on dequeue) 

160 assert t.status == TaskStatus.CANCELLED 

161 

162 def test_cancel_nonexistent(self): 

163 q = PriorityTaskQueue() 

164 assert q.cancel("no-such") is False 

165 

166 def test_cancel_removed_on_dequeue(self): 

167 q = PriorityTaskQueue() 

168 t1 = Task(task_id="c", func=lambda: 1) 

169 t2 = Task(task_id="d", func=lambda: 2) 

170 q.enqueue(t1) 

171 q.enqueue(t2) 

172 q.cancel("c") 

173 result = q.dequeue() 

174 assert result is t2 # cancelled t1 skipped, t2 returned 

175 assert q.dequeue() is None 

176 

177 

178# ============================================================================ 

179# PriorityTaskQueue — Expiration 

180# ============================================================================ 

181 

182class TestPriorityTaskQueueExpire: 

183 def test_deadline_expired(self): 

184 q = PriorityTaskQueue() 

185 t = Task(task_id="e", func=lambda: 1, deadline=time.monotonic() - 10) 

186 q.enqueue(t) 

187 result = q.dequeue() 

188 assert result is None 

189 assert t.status == TaskStatus.EXPIRED 

190 

191 def test_deadline_future(self): 

192 q = PriorityTaskQueue() 

193 t = Task(task_id="f", func=lambda: 1, deadline=time.monotonic() + 3600) 

194 q.enqueue(t) 

195 result = q.dequeue() 

196 assert result is t 

197 assert t.status == TaskStatus.PENDING # not expired 

198 

199 def test_cancelled_before_expired(self): 

200 q = PriorityTaskQueue() 

201 t = Task(task_id="g", func=lambda: 1, deadline=time.monotonic() - 10) 

202 q.enqueue(t) 

203 q.cancel("g") 

204 result = q.dequeue() 

205 assert result is None # cancelled and removed 

206 assert t.status == TaskStatus.CANCELLED 

207 

208 def test_expired_stats(self): 

209 q = PriorityTaskQueue() 

210 t = Task(task_id="h", func=lambda: 1, deadline=time.monotonic() - 10) 

211 q.enqueue(t) 

212 q.dequeue() 

213 assert q.stats["total_expired"] == 1 

214 

215 

216# ============================================================================ 

217# PriorityTaskQueue — Stats 

218# ============================================================================ 

219 

220class TestPriorityTaskQueueStats: 

221 def test_default_stats(self): 

222 q = PriorityTaskQueue(max_size=500) 

223 s = q.stats 

224 assert s["size"] == 0 

225 assert s["max_size"] == 500 

226 assert s["total_enqueued"] == 0 

227 assert s["total_dequeued"] == 0 

228 assert s["total_expired"] == 0 

229 

230 def test_stats_after_ops(self): 

231 q = PriorityTaskQueue() 

232 q.enqueue(Task(task_id="a", func=lambda: 1)) 

233 q.enqueue(Task(task_id="b", func=lambda: 2)) 

234 q.dequeue() 

235 s = q.stats 

236 assert s["total_enqueued"] == 2 

237 assert s["total_dequeued"] == 1 

238 assert s["size"] == 1 

239 

240 def test_stats_includes_expired(self): 

241 q = PriorityTaskQueue() 

242 t = Task(task_id="x", func=lambda: 1, deadline=time.monotonic() - 1) 

243 q.enqueue(t) 

244 q.dequeue() 

245 s = q.stats 

246 assert s["total_expired"] == 1 

247 

248 

249# ============================================================================ 

250# TaskScheduler — Submit & Immediate Execution 

251# ============================================================================ 

252 

253class TestTaskSchedulerSubmit: 

254 def test_submit_and_run_once(self): 

255 sched = TaskScheduler() 

256 t = Task(task_id="t1", func=lambda: 99) 

257 assert sched.submit(t) is True 

258 result = sched.run_once() 

259 assert result is t 

260 assert t.status == TaskStatus.COMPLETED 

261 assert t.result == 99 

262 

263 def test_submit_queue_full(self): 

264 sched = TaskScheduler() 

265 sched._queue = PriorityTaskQueue(max_size=1) 

266 t1 = Task(task_id="a", func=lambda: 1) 

267 t2 = Task(task_id="b", func=lambda: 2) 

268 assert sched.submit(t1) is True 

269 assert sched.submit(t2) is False 

270 

271 def test_run_once_empty(self): 

272 sched = TaskScheduler() 

273 assert sched.run_once() is None 

274 

275 def test_pending_count(self): 

276 sched = TaskScheduler() 

277 assert sched.pending == 0 

278 sched.submit(Task(task_id="x", func=lambda: 1)) 

279 assert sched.pending == 1 

280 sched.run_once() 

281 assert sched.pending == 0 

282 

283 def test_run_loop(self): 

284 sched = TaskScheduler() 

285 sched.submit(Task(task_id="a", func=lambda: 1)) 

286 sched.submit(Task(task_id="b", func=lambda: 2)) 

287 sched.submit(Task(task_id="c", func=lambda: 3)) 

288 

289 count = sched.run_loop() 

290 assert count == 3 

291 assert sched.pending == 0 

292 

293 def test_run_loop_max_tasks(self): 

294 sched = TaskScheduler() 

295 for i in range(5): 

296 sched.submit(Task(task_id=f"t{i}", func=lambda: i)) 

297 

298 count = sched.run_loop(max_tasks=2) 

299 assert count == 2 

300 assert sched.pending == 3 

301 

302 

303# ============================================================================ 

304# TaskScheduler — Schedule After 

305# ============================================================================ 

306 

307class TestTaskSchedulerScheduleAfter: 

308 def test_schedule_after_fires(self): 

309 sched = TaskScheduler() 

310 sched.start() 

311 t = sched.schedule_after(lambda: 42, 0.05) 

312 time.sleep(0.15) 

313 sched.stop() 

314 # Note: schedule_after sets deadline; task must be dequeued manually or via run_loop 

315 assert t.status == TaskStatus.PENDING # not auto-executed, just scheduled 

316 

317 def test_schedule_after_task_attributes(self): 

318 sched = TaskScheduler() 

319 t = sched.schedule_after(lambda x: x, 1.5, 10) 

320 assert t.priority == 1500 # int(delay * 1000) 

321 assert t.deadline is not None 

322 assert t.args == (10,) 

323 

324 

325# ============================================================================ 

326# TaskScheduler — Start / Stop / Stats 

327# ============================================================================ 

328 

329class TestTaskSchedulerLifecycle: 

330 def test_start_stop(self): 

331 sched = TaskScheduler() 

332 assert sched._running is False 

333 sched.start() 

334 assert sched._running is True 

335 sched.stop() 

336 assert sched._running is False 

337 

338 def test_stats(self): 

339 sched = TaskScheduler() 

340 sched.submit(Task(task_id="x", func=lambda: 1)) 

341 sched.run_once() 

342 

343 s = sched.stats 

344 assert s["scheduled_count"] == 1 

345 assert s["executed_count"] == 1 

346 assert s["pending"] == 0 

347 

348 def test_schedule_at_interval(self): 

349 results = [] 

350 

351 def put(): 

352 results.append(1) 

353 

354 sched = TaskScheduler() 

355 sched.start() 

356 sched.schedule_at_interval(put, 0.05) 

357 time.sleep(0.15) 

358 sched.stop() 

359 

360 # The interval thread submits tasks to the queue, but they aren't 

361 # auto-executed without a run_loop. The tasks are queued as PENDING. 

362 assert sched._scheduled_count >= 2 # at least 2 submissions in 0.15s 

363 

364 

365# ============================================================================ 

366# WorkerPool 

367# ============================================================================ 

368 

369class TestWorkerPool: 

370 def test_creation(self): 

371 pool = WorkerPool(num_workers=2, max_queue_size=100) 

372 assert pool._num_workers == 2 

373 assert pool._running is False 

374 assert pool.pending == 0 

375 

376 def test_submit_before_start(self): 

377 pool = WorkerPool(num_workers=2) 

378 t = pool.submit(lambda: 42) 

379 assert t.status == TaskStatus.PENDING 

380 assert pool.pending == 1 

381 

382 def test_submit_full_queue(self): 

383 pool = WorkerPool(num_workers=1, max_queue_size=1) 

384 pool.submit(lambda: 1) 

385 with pytest.raises(RuntimeError, match="full"): 

386 pool.submit(lambda: 2) 

387 

388 def test_start_stop_workers(self): 

389 results = [] 

390 pool = WorkerPool(num_workers=2, max_queue_size=10) 

391 pool.submit(lambda: results.append("a")) 

392 pool.submit(lambda: results.append("b")) 

393 

394 pool.start() 

395 time.sleep(0.15) 

396 pool.stop() 

397 

398 assert len(results) == 2 

399 assert "a" in results 

400 assert "b" in results 

401 

402 def test_multiple_tasks_workers(self): 

403 results = [] 

404 lock = threading.Lock() 

405 

406 def add(n): 

407 time.sleep(0.02) 

408 with lock: 

409 results.append(n) 

410 

411 pool = WorkerPool(num_workers=3, max_queue_size=20) 

412 for i in range(9): 

413 pool.submit(lambda n=i: add(n)) 

414 

415 pool.start() 

416 time.sleep(0.3) 

417 pool.stop() 

418 

419 assert len(results) == 9 

420 

421 def test_stop_with_wait(self): 

422 pool = WorkerPool(num_workers=2, max_queue_size=10) 

423 pool.submit(lambda: time.sleep(0.05)) 

424 pool.start() 

425 time.sleep(0.02) 

426 pool.stop(wait=True) 

427 assert pool._running is False 

428 

429 def test_stop_no_wait(self): 

430 pool = WorkerPool(num_workers=2, max_queue_size=10) 

431 pool.submit(lambda: time.sleep(0.5)) 

432 pool.start() 

433 time.sleep(0.02) 

434 pool.stop(wait=False) 

435 assert pool._running is False 

436 

437 def test_stats(self): 

438 pool = WorkerPool(num_workers=3, max_queue_size=10) 

439 s = pool.stats 

440 assert s["workers"] == 3 

441 assert s["running"] is False 

442 assert s["size"] == 0 

443 

444 def test_stats_after_execution(self): 

445 pool = WorkerPool(num_workers=2, max_queue_size=20) 

446 for _ in range(5): 

447 pool.submit(lambda: None) 

448 pool.start() 

449 time.sleep(0.15) 

450 pool.stop() 

451 

452 s = pool.stats 

453 assert s["total_dequeued"] == 5 

454 

455 def test_concurrent_submit(self): 

456 pool = WorkerPool(num_workers=4, max_queue_size=100) 

457 pool.start() 

458 

459 errors = [] 

460 

461 def submitter(): 

462 try: 

463 for _ in range(10): 

464 pool.submit(lambda: time.sleep(0.01)) 

465 except Exception as e: 

466 errors.append(e) 

467 

468 threads = [threading.Thread(target=submitter) for _ in range(4)] 

469 for t in threads: 

470 t.start() 

471 for t in threads: 

472 t.join() 

473 

474 pool.stop() 

475 assert len(errors) == 0 

476 s = pool.stats 

477 assert s["total_enqueued"] == 40