diff --git a/sigflow/workers/pool.py b/sigflow/workers/pool.py index 1b85211..e473620 100644 --- a/sigflow/workers/pool.py +++ b/sigflow/workers/pool.py @@ -1,4 +1,4 @@ -from concurrent.futures import ThreadPoolExecutor, as_completed +from concurrent.futures import ThreadPoolExecutor class WorkerPool: @@ -7,5 +7,4 @@ def __init__(self, workers: int = 4): def map(self, fn, items): with ThreadPoolExecutor(max_workers=self.workers) as pool: - futures = [pool.submit(fn, item) for item in items] - return [future.result() for future in as_completed(futures)] + return list(pool.map(fn, items)) \ No newline at end of file diff --git a/tests/test_workers.py b/tests/test_workers.py index df2f57b..9fc77f0 100644 --- a/tests/test_workers.py +++ b/tests/test_workers.py @@ -2,4 +2,16 @@ def test_worker_pool_maps_items(): - assert sorted(WorkerPool(2).map(lambda x: x + 1, [1, 2, 3])) == [2, 3, 4] + assert WorkerPool(2).map(lambda x: x + 1, [1, 2, 3]) == [2, 3, 4] + + +def test_worker_pool_preserves_order_with_varying_times(): + import time + + def slow_if_first(x): + if x == 0: + time.sleep(0.1) + return x + + results = WorkerPool(workers=2).map(slow_if_first, [0, 1, 2]) + assert results == [0, 1, 2] \ No newline at end of file