From 7aa5115165d18383a7281d1473278a1ab06c2598 Mon Sep 17 00:00:00 2001 From: "Alex Davies (aider)" Date: Wed, 8 Jan 2025 11:36:00 -0400 Subject: [PATCH] test: Add test for capturing output from multiple asyncio tasks in parallel --- tests/test_capture_output.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/tests/test_capture_output.py b/tests/test_capture_output.py index 8888821..f5e034f 100644 --- a/tests/test_capture_output.py +++ b/tests/test_capture_output.py @@ -16,6 +16,33 @@ def test_sync_capture_line_by_line(): assert "[stdout] Line 2" in output assert "[stderr] Error!" in output +@pytest.mark.asyncio +async def test_async_capture_parallel_tasks(): + output = [] + def callback(data): + output.append(data) + + async def task1(): + async with CaptureOutput(callback).async_context(): + print("Task 1 - Line 1") + print("Task 1 - Line 2") + print("Task 1 - Error!", file=sys.stderr) + + async def task2(): + async with CaptureOutput(callback).async_context(): + print("Task 2 - Line 1") + print("Task 2 - Line 2") + print("Task 2 - Error!", file=sys.stderr) + + await asyncio.gather(task1(), task2()) + + assert "[stdout] Task 1 - Line 1" in output + assert "[stdout] Task 1 - Line 2" in output + assert "[stderr] Task 1 - Error!" in output + assert "[stdout] Task 2 - Line 1" in output + assert "[stdout] Task 2 - Line 2" in output + assert "[stderr] Task 2 - Error!" in output + @pytest.mark.asyncio async def test_async_capture_line_by_line(): output = []