wehub-resource-sync
55 lines
1.5 KiB
Python
55 lines
1.5 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import Awaitable, Callable
|
|
|
|
import pytest
|
|
|
|
from agents.sandbox.materialization import gather_in_order
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_gather_in_order_limits_concurrency_and_preserves_order() -> None:
|
|
active_tasks = 0
|
|
max_active_tasks = 0
|
|
release_tasks = asyncio.Event()
|
|
started_tasks: list[int] = []
|
|
|
|
def task_factory(index: int) -> Callable[[], Awaitable[str]]:
|
|
async def run() -> str:
|
|
nonlocal active_tasks
|
|
nonlocal max_active_tasks
|
|
active_tasks += 1
|
|
max_active_tasks = max(max_active_tasks, active_tasks)
|
|
started_tasks.append(index)
|
|
try:
|
|
await release_tasks.wait()
|
|
return f"result-{index}"
|
|
finally:
|
|
active_tasks -= 1
|
|
|
|
return run
|
|
|
|
gather_task = asyncio.create_task(
|
|
gather_in_order([task_factory(index) for index in range(5)], max_concurrency=2)
|
|
)
|
|
while len(started_tasks) < 2:
|
|
await asyncio.sleep(0)
|
|
|
|
assert started_tasks == [0, 1]
|
|
assert max_active_tasks == 2
|
|
|
|
release_tasks.set()
|
|
result = await gather_task
|
|
|
|
assert result == ["result-0", "result-1", "result-2", "result-3", "result-4"]
|
|
assert max_active_tasks == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_gather_in_order_rejects_invalid_concurrency() -> None:
|
|
with pytest.raises(ValueError) as exc_info:
|
|
await gather_in_order([], max_concurrency=0)
|
|
|
|
assert str(exc_info.value) == "max_concurrency must be at least 1"
|