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

348 statements  

« prev     ^ index     » next       coverage.py v7.14.3, created at 2026-07-09 10:28 +0800

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

2 

3import threading 

4import time 

5 

6import pytest 

7 

8from agentos.tools.task_scheduler import ( 

9 PriorityTaskQueue, 

10 Task, 

11 TaskScheduler, 

12 TaskStatus, 

13 WorkerPool, 

14) 

15 

16# ============================================================================ 

17# TaskStatus 

18# ============================================================================ 

19 

20class TestTaskStatus: 

21 def test_enum_members(self): 

22 assert TaskStatus.PENDING is not None 

23 assert TaskStatus.RUNNING is not None 

24 assert TaskStatus.COMPLETED is not None 

25 assert TaskStatus.FAILED is not None 

26 assert TaskStatus.EXPIRED is not None 

27 assert TaskStatus.CANCELLED is not None 

28 

29 def test_unique_values(self): 

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

31 assert len(values) == len(set(values)) 

32 

33 

34# ============================================================================ 

35# Task 

36# ============================================================================ 

37 

38class TestTask: 

39 def test_creation_defaults(self): 

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

41 assert t.task_id == "t1" 

42 assert t.args == () 

43 assert t.kwargs == {} 

44 assert t.priority == 0 

45 assert t.deadline is None 

46 assert t.status == TaskStatus.PENDING 

47 assert t.result is None 

48 assert t.error is None 

49 assert t.started_at is None 

50 assert t.completed_at is None 

51 

52 def test_creation_with_args(self): 

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

54 priority=3, deadline=100.0) 

55 assert t.args == (5,) 

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

57 assert t.priority == 3 

58 assert t.deadline == 100.0 

59 

60 def test_run_success(self): 

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

62 t.run() 

63 assert t.status == TaskStatus.COMPLETED 

64 assert t.result == 3 

65 assert t.error is None 

66 assert t.started_at is not None 

67 assert t.completed_at is not None 

68 

69 def test_run_failure(self): 

70 def boom(): 

71 raise ValueError("kaboom") 

72 

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

74 t.run() 

75 assert t.status == TaskStatus.FAILED 

76 assert t.result is None 

77 assert t.error == "kaboom" 

78 assert t.completed_at is not None 

79 

80 def test_elapsed_returns_float(self): 

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

82 assert t.elapsed is None # not started 

83 t.run() 

84 assert isinstance(t.elapsed, float) 

85 assert t.elapsed >= 0 

86 

87 def test_elapsed_before_start(self): 

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

89 assert t.elapsed is None 

90 

91 

92# ============================================================================ 

93# PriorityTaskQueue — Basic Operations 

94# ============================================================================ 

95 

96class TestPriorityTaskQueueBasic: 

97 def test_creation_defaults(self): 

98 q = PriorityTaskQueue() 

99 assert q.size == 0 

100 assert q._max_size == 1000 

101 

102 def test_creation_custom_size(self): 

103 q = PriorityTaskQueue(max_size=10) 

104 assert q._max_size == 10 

105 

106 def test_enqueue_dequeue(self): 

107 q = PriorityTaskQueue() 

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

109 assert q.enqueue(t) is True 

110 assert q.size == 1 

111 

112 result = q.dequeue() 

113 assert result is t 

114 assert q.size == 0 

115 

116 def test_dequeue_empty(self): 

117 q = PriorityTaskQueue() 

118 assert q.dequeue() is None 

119 

120 def test_enqueue_full(self): 

121 q = PriorityTaskQueue(max_size=2) 

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

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

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

125 assert q.size == 2 

126 

127 def test_priority_ordering(self): 

128 q = PriorityTaskQueue() 

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

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

131 q.enqueue(t1) 

132 q.enqueue(t2) 

133 

134 first = q.dequeue() 

135 second = q.dequeue() 

136 assert first.task_id == "high" 

137 assert second.task_id == "low" 

138 

139 def test_same_priority_fifo(self): 

140 q = PriorityTaskQueue() 

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

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

143 q.enqueue(t1) 

144 q.enqueue(t2) 

145 

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

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

148 

149 

150# ============================================================================ 

151# PriorityTaskQueue — Cancellation 

152# ============================================================================ 

153 

154class TestPriorityTaskQueueCancel: 

155 def test_cancel_existing(self): 

156 q = PriorityTaskQueue() 

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

158 q.enqueue(t) 

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

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

161 assert t.status == TaskStatus.CANCELLED 

162 

163 def test_cancel_nonexistent(self): 

164 q = PriorityTaskQueue() 

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

166 

167 def test_cancel_removed_on_dequeue(self): 

168 q = PriorityTaskQueue() 

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

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

171 q.enqueue(t1) 

172 q.enqueue(t2) 

173 q.cancel("c") 

174 result = q.dequeue() 

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

176 assert q.dequeue() is None 

177 

178 

179# ============================================================================ 

180# PriorityTaskQueue — Expiration 

181# ============================================================================ 

182 

183class TestPriorityTaskQueueExpire: 

184 def test_deadline_expired(self): 

185 q = PriorityTaskQueue() 

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

187 q.enqueue(t) 

188 result = q.dequeue() 

189 assert result is None 

190 assert t.status == TaskStatus.EXPIRED 

191 

192 def test_deadline_future(self): 

193 q = PriorityTaskQueue() 

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

195 q.enqueue(t) 

196 result = q.dequeue() 

197 assert result is t 

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

199 

200 def test_cancelled_before_expired(self): 

201 q = PriorityTaskQueue() 

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

203 q.enqueue(t) 

204 q.cancel("g") 

205 result = q.dequeue() 

206 assert result is None # cancelled and removed 

207 assert t.status == TaskStatus.CANCELLED 

208 

209 def test_expired_stats(self): 

210 q = PriorityTaskQueue() 

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

212 q.enqueue(t) 

213 q.dequeue() 

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

215 

216 

217# ============================================================================ 

218# PriorityTaskQueue — Stats 

219# ============================================================================ 

220 

221class TestPriorityTaskQueueStats: 

222 def test_default_stats(self): 

223 q = PriorityTaskQueue(max_size=500) 

224 s = q.stats 

225 assert s["size"] == 0 

226 assert s["max_size"] == 500 

227 assert s["total_enqueued"] == 0 

228 assert s["total_dequeued"] == 0 

229 assert s["total_expired"] == 0 

230 

231 def test_stats_after_ops(self): 

232 q = PriorityTaskQueue() 

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

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

235 q.dequeue() 

236 s = q.stats 

237 assert s["total_enqueued"] == 2 

238 assert s["total_dequeued"] == 1 

239 assert s["size"] == 1 

240 

241 def test_stats_includes_expired(self): 

242 q = PriorityTaskQueue() 

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

244 q.enqueue(t) 

245 q.dequeue() 

246 s = q.stats 

247 assert s["total_expired"] == 1 

248 

249 

250# ============================================================================ 

251# TaskScheduler — Submit & Immediate Execution 

252# ============================================================================ 

253 

254class TestTaskSchedulerSubmit: 

255 def test_submit_and_run_once(self): 

256 sched = TaskScheduler() 

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

258 assert sched.submit(t) is True 

259 result = sched.run_once() 

260 assert result is t 

261 assert t.status == TaskStatus.COMPLETED 

262 assert t.result == 99 

263 

264 def test_submit_queue_full(self): 

265 sched = TaskScheduler() 

266 sched._queue = PriorityTaskQueue(max_size=1) 

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

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

269 assert sched.submit(t1) is True 

270 assert sched.submit(t2) is False 

271 

272 def test_run_once_empty(self): 

273 sched = TaskScheduler() 

274 assert sched.run_once() is None 

275 

276 def test_pending_count(self): 

277 sched = TaskScheduler() 

278 assert sched.pending == 0 

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

280 assert sched.pending == 1 

281 sched.run_once() 

282 assert sched.pending == 0 

283 

284 def test_run_loop(self): 

285 sched = TaskScheduler() 

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

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

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

289 

290 count = sched.run_loop() 

291 assert count == 3 

292 assert sched.pending == 0 

293 

294 def test_run_loop_max_tasks(self): 

295 sched = TaskScheduler() 

296 for i in range(5): 

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

298 

299 count = sched.run_loop(max_tasks=2) 

300 assert count == 2 

301 assert sched.pending == 3 

302 

303 

304# ============================================================================ 

305# TaskScheduler — Schedule After 

306# ============================================================================ 

307 

308class TestTaskSchedulerScheduleAfter: 

309 def test_schedule_after_fires(self): 

310 sched = TaskScheduler() 

311 sched.start() 

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

313 time.sleep(0.15) 

314 sched.stop() 

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

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

317 

318 def test_schedule_after_task_attributes(self): 

319 sched = TaskScheduler() 

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

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

322 assert t.deadline is not None 

323 assert t.args == (10,) 

324 

325 

326# ============================================================================ 

327# TaskScheduler — Start / Stop / Stats 

328# ============================================================================ 

329 

330class TestTaskSchedulerLifecycle: 

331 def test_start_stop(self): 

332 sched = TaskScheduler() 

333 assert sched._running is False 

334 sched.start() 

335 assert sched._running is True 

336 sched.stop() 

337 assert sched._running is False 

338 

339 def test_stats(self): 

340 sched = TaskScheduler() 

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

342 sched.run_once() 

343 

344 s = sched.stats 

345 assert s["scheduled_count"] == 1 

346 assert s["executed_count"] == 1 

347 assert s["pending"] == 0 

348 

349 def test_schedule_at_interval(self): 

350 results = [] 

351 

352 def put(): 

353 results.append(1) 

354 

355 sched = TaskScheduler() 

356 sched.start() 

357 sched.schedule_at_interval(put, 0.05) 

358 time.sleep(0.15) 

359 sched.stop() 

360 

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

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

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

364 

365 

366# ============================================================================ 

367# WorkerPool 

368# ============================================================================ 

369 

370class TestWorkerPool: 

371 def test_creation(self): 

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

373 assert pool._num_workers == 2 

374 assert pool._running is False 

375 assert pool.pending == 0 

376 

377 def test_submit_before_start(self): 

378 pool = WorkerPool(num_workers=2) 

379 t = pool.submit(lambda: 42) 

380 assert t.status == TaskStatus.PENDING 

381 assert pool.pending == 1 

382 

383 def test_submit_full_queue(self): 

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

385 pool.submit(lambda: 1) 

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

387 pool.submit(lambda: 2) 

388 

389 def test_start_stop_workers(self): 

390 results = [] 

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

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

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

394 

395 pool.start() 

396 time.sleep(0.15) 

397 pool.stop() 

398 

399 assert len(results) == 2 

400 assert "a" in results 

401 assert "b" in results 

402 

403 def test_multiple_tasks_workers(self): 

404 results = [] 

405 lock = threading.Lock() 

406 

407 def add(n): 

408 time.sleep(0.02) 

409 with lock: 

410 results.append(n) 

411 

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

413 for i in range(9): 

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

415 

416 pool.start() 

417 time.sleep(0.3) 

418 pool.stop() 

419 

420 assert len(results) == 9 

421 

422 def test_stop_with_wait(self): 

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

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

425 pool.start() 

426 time.sleep(0.02) 

427 pool.stop(wait=True) 

428 assert pool._running is False 

429 

430 def test_stop_no_wait(self): 

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

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

433 pool.start() 

434 time.sleep(0.02) 

435 pool.stop(wait=False) 

436 assert pool._running is False 

437 

438 def test_stats(self): 

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

440 s = pool.stats 

441 assert s["workers"] == 3 

442 assert s["running"] is False 

443 assert s["size"] == 0 

444 

445 def test_stats_after_execution(self): 

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

447 for _ in range(5): 

448 pool.submit(lambda: None) 

449 pool.start() 

450 time.sleep(0.15) 

451 pool.stop() 

452 

453 s = pool.stats 

454 assert s["total_dequeued"] == 5 

455 

456 def test_concurrent_submit(self): 

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

458 pool.start() 

459 

460 errors = [] 

461 

462 def submitter(): 

463 try: 

464 for _ in range(10): 

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

466 except Exception as e: 

467 errors.append(e) 

468 

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

470 for t in threads: 

471 t.start() 

472 for t in threads: 

473 t.join() 

474 

475 pool.stop() 

476 assert len(errors) == 0 

477 s = pool.stats 

478 assert s["total_enqueued"] == 40