wehub-resource-sync
104 lines
2.7 KiB
Python
104 lines
2.7 KiB
Python
import asyncio
|
|
|
|
import pytest
|
|
|
|
from livekit.agents.utils.aio.itertools import Tee
|
|
|
|
pytestmark = [pytest.mark.unit, pytest.mark.virtual_time, pytest.mark.no_concurrent]
|
|
|
|
|
|
async def _async_iter(items):
|
|
for item in items:
|
|
yield item
|
|
|
|
|
|
async def _failing_iter(items, error):
|
|
for item in items:
|
|
yield item
|
|
raise error
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tee_basic():
|
|
tee = Tee(_async_iter([1, 2, 3]), n=3)
|
|
for peer in tee:
|
|
result = [item async for item in peer]
|
|
assert result == [1, 2, 3]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tee_exception_propagation():
|
|
err = ValueError("upstream error")
|
|
tee = Tee(_failing_iter([1, 2], err), n=2)
|
|
|
|
for peer in tee:
|
|
result = []
|
|
with pytest.raises(ValueError, match="upstream error"):
|
|
async for item in peer:
|
|
result.append(item)
|
|
assert result == [1, 2]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tee_exception_concurrent():
|
|
err = ValueError("concurrent error")
|
|
tee = Tee(_failing_iter([1], err), n=2)
|
|
|
|
results: list[list[int]] = [[], []]
|
|
errors: list[BaseException | None] = [None, None]
|
|
|
|
async def consume(idx: int, peer):
|
|
try:
|
|
async for item in peer:
|
|
results[idx].append(item)
|
|
except ValueError as e:
|
|
errors[idx] = e
|
|
|
|
await asyncio.gather(consume(0, tee[0]), consume(1, tee[1]))
|
|
|
|
for idx in range(2):
|
|
assert results[idx] == [1]
|
|
assert isinstance(errors[idx], ValueError)
|
|
assert str(errors[idx]) == "concurrent error"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tee_cancelled_error_does_not_propagate_to_peers():
|
|
"""CancelledError in one peer's task should not cascade to other peers."""
|
|
items = [1, 2, 3]
|
|
tee = Tee(_async_iter(items), n=2)
|
|
|
|
peer1_results: list[int] = []
|
|
|
|
async def consume_peer0():
|
|
async for _item in tee[0]:
|
|
# Block forever after first item so we can cancel this task
|
|
await asyncio.sleep(1000)
|
|
|
|
async def consume_peer1():
|
|
async for item in tee[1]:
|
|
peer1_results.append(item)
|
|
|
|
task0 = asyncio.create_task(consume_peer0())
|
|
task1 = asyncio.create_task(consume_peer1())
|
|
|
|
# Give tasks a chance to start
|
|
await asyncio.sleep(0.01)
|
|
|
|
# Cancel task0 — this should NOT affect task1
|
|
task0.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task0
|
|
|
|
# task1 should complete normally without CancelledError
|
|
await asyncio.wait_for(task1, timeout=2.0)
|
|
assert peer1_results == items
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tee_empty_iterator():
|
|
tee = Tee(_async_iter([]), n=2)
|
|
for peer in tee:
|
|
result = [item async for item in peer]
|
|
assert result == []
|