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."""
3import threading
4import time
6import pytest
8from agentos.tools.task_scheduler import (
9 PriorityTaskQueue,
10 Task,
11 TaskScheduler,
12 TaskStatus,
13 WorkerPool,
14)
16# ============================================================================
17# TaskStatus
18# ============================================================================
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
29 def test_unique_values(self):
30 values = [s.value for s in TaskStatus]
31 assert len(values) == len(set(values))
34# ============================================================================
35# Task
36# ============================================================================
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
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
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
69 def test_run_failure(self):
70 def boom():
71 raise ValueError("kaboom")
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
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
87 def test_elapsed_before_start(self):
88 t = Task(task_id="t6", func=lambda: 1)
89 assert t.elapsed is None
92# ============================================================================
93# PriorityTaskQueue — Basic Operations
94# ============================================================================
96class TestPriorityTaskQueueBasic:
97 def test_creation_defaults(self):
98 q = PriorityTaskQueue()
99 assert q.size == 0
100 assert q._max_size == 1000
102 def test_creation_custom_size(self):
103 q = PriorityTaskQueue(max_size=10)
104 assert q._max_size == 10
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
112 result = q.dequeue()
113 assert result is t
114 assert q.size == 0
116 def test_dequeue_empty(self):
117 q = PriorityTaskQueue()
118 assert q.dequeue() is None
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
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)
134 first = q.dequeue()
135 second = q.dequeue()
136 assert first.task_id == "high"
137 assert second.task_id == "low"
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)
146 assert q.dequeue().task_id == "first"
147 assert q.dequeue().task_id == "second"
150# ============================================================================
151# PriorityTaskQueue — Cancellation
152# ============================================================================
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
163 def test_cancel_nonexistent(self):
164 q = PriorityTaskQueue()
165 assert q.cancel("no-such") is False
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
179# ============================================================================
180# PriorityTaskQueue — Expiration
181# ============================================================================
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
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
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
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
217# ============================================================================
218# PriorityTaskQueue — Stats
219# ============================================================================
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
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
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
250# ============================================================================
251# TaskScheduler — Submit & Immediate Execution
252# ============================================================================
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
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
272 def test_run_once_empty(self):
273 sched = TaskScheduler()
274 assert sched.run_once() is None
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
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))
290 count = sched.run_loop()
291 assert count == 3
292 assert sched.pending == 0
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))
299 count = sched.run_loop(max_tasks=2)
300 assert count == 2
301 assert sched.pending == 3
304# ============================================================================
305# TaskScheduler — Schedule After
306# ============================================================================
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
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,)
326# ============================================================================
327# TaskScheduler — Start / Stop / Stats
328# ============================================================================
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
339 def test_stats(self):
340 sched = TaskScheduler()
341 sched.submit(Task(task_id="x", func=lambda: 1))
342 sched.run_once()
344 s = sched.stats
345 assert s["scheduled_count"] == 1
346 assert s["executed_count"] == 1
347 assert s["pending"] == 0
349 def test_schedule_at_interval(self):
350 results = []
352 def put():
353 results.append(1)
355 sched = TaskScheduler()
356 sched.start()
357 sched.schedule_at_interval(put, 0.05)
358 time.sleep(0.15)
359 sched.stop()
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
366# ============================================================================
367# WorkerPool
368# ============================================================================
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
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
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)
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"))
395 pool.start()
396 time.sleep(0.15)
397 pool.stop()
399 assert len(results) == 2
400 assert "a" in results
401 assert "b" in results
403 def test_multiple_tasks_workers(self):
404 results = []
405 lock = threading.Lock()
407 def add(n):
408 time.sleep(0.02)
409 with lock:
410 results.append(n)
412 pool = WorkerPool(num_workers=3, max_queue_size=20)
413 for i in range(9):
414 pool.submit(lambda n=i: add(n))
416 pool.start()
417 time.sleep(0.3)
418 pool.stop()
420 assert len(results) == 9
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
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
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
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()
453 s = pool.stats
454 assert s["total_dequeued"] == 5
456 def test_concurrent_submit(self):
457 pool = WorkerPool(num_workers=4, max_queue_size=100)
458 pool.start()
460 errors = []
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)
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()
475 pool.stop()
476 assert len(errors) == 0
477 s = pool.stats
478 assert s["total_enqueued"] == 40