Coverage for agentos/tests/test_task_scheduler.py: 0%
348 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-09 09:19 +0800
1"""Tests for agentos.tools.task_scheduler — Task, PriorityTaskQueue, TaskScheduler, WorkerPool."""
3import pytest
4import threading
5import time
6from agentos.tools.task_scheduler import (
7 PriorityTaskQueue,
8 Task,
9 TaskScheduler,
10 TaskStatus,
11 WorkerPool,
12)
15# ============================================================================
16# TaskStatus
17# ============================================================================
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
28 def test_unique_values(self):
29 values = [s.value for s in TaskStatus]
30 assert len(values) == len(set(values))
33# ============================================================================
34# Task
35# ============================================================================
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
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
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
68 def test_run_failure(self):
69 def boom():
70 raise ValueError("kaboom")
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
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
86 def test_elapsed_before_start(self):
87 t = Task(task_id="t6", func=lambda: 1)
88 assert t.elapsed is None
91# ============================================================================
92# PriorityTaskQueue — Basic Operations
93# ============================================================================
95class TestPriorityTaskQueueBasic:
96 def test_creation_defaults(self):
97 q = PriorityTaskQueue()
98 assert q.size == 0
99 assert q._max_size == 1000
101 def test_creation_custom_size(self):
102 q = PriorityTaskQueue(max_size=10)
103 assert q._max_size == 10
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
111 result = q.dequeue()
112 assert result is t
113 assert q.size == 0
115 def test_dequeue_empty(self):
116 q = PriorityTaskQueue()
117 assert q.dequeue() is None
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
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)
133 first = q.dequeue()
134 second = q.dequeue()
135 assert first.task_id == "high"
136 assert second.task_id == "low"
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)
145 assert q.dequeue().task_id == "first"
146 assert q.dequeue().task_id == "second"
149# ============================================================================
150# PriorityTaskQueue — Cancellation
151# ============================================================================
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
162 def test_cancel_nonexistent(self):
163 q = PriorityTaskQueue()
164 assert q.cancel("no-such") is False
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
178# ============================================================================
179# PriorityTaskQueue — Expiration
180# ============================================================================
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
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
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
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
216# ============================================================================
217# PriorityTaskQueue — Stats
218# ============================================================================
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
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
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
249# ============================================================================
250# TaskScheduler — Submit & Immediate Execution
251# ============================================================================
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
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
271 def test_run_once_empty(self):
272 sched = TaskScheduler()
273 assert sched.run_once() is None
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
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))
289 count = sched.run_loop()
290 assert count == 3
291 assert sched.pending == 0
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))
298 count = sched.run_loop(max_tasks=2)
299 assert count == 2
300 assert sched.pending == 3
303# ============================================================================
304# TaskScheduler — Schedule After
305# ============================================================================
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
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,)
325# ============================================================================
326# TaskScheduler — Start / Stop / Stats
327# ============================================================================
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
338 def test_stats(self):
339 sched = TaskScheduler()
340 sched.submit(Task(task_id="x", func=lambda: 1))
341 sched.run_once()
343 s = sched.stats
344 assert s["scheduled_count"] == 1
345 assert s["executed_count"] == 1
346 assert s["pending"] == 0
348 def test_schedule_at_interval(self):
349 results = []
351 def put():
352 results.append(1)
354 sched = TaskScheduler()
355 sched.start()
356 sched.schedule_at_interval(put, 0.05)
357 time.sleep(0.15)
358 sched.stop()
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
365# ============================================================================
366# WorkerPool
367# ============================================================================
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
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
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)
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"))
394 pool.start()
395 time.sleep(0.15)
396 pool.stop()
398 assert len(results) == 2
399 assert "a" in results
400 assert "b" in results
402 def test_multiple_tasks_workers(self):
403 results = []
404 lock = threading.Lock()
406 def add(n):
407 time.sleep(0.02)
408 with lock:
409 results.append(n)
411 pool = WorkerPool(num_workers=3, max_queue_size=20)
412 for i in range(9):
413 pool.submit(lambda n=i: add(n))
415 pool.start()
416 time.sleep(0.3)
417 pool.stop()
419 assert len(results) == 9
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
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
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
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()
452 s = pool.stats
453 assert s["total_dequeued"] == 5
455 def test_concurrent_submit(self):
456 pool = WorkerPool(num_workers=4, max_queue_size=100)
457 pool.start()
459 errors = []
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)
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()
474 pool.stop()
475 assert len(errors) == 0
476 s = pool.stats
477 assert s["total_enqueued"] == 40