File tree Expand file tree Collapse file tree 1 file changed +4
-0
lines changed Expand file tree Collapse file tree 1 file changed +4
-0
lines changed Original file line number Diff line number Diff line change
1
+ import asyncio
1
2
import unittest
2
3
3
4
from collections .abc import AsyncIterator
@@ -252,6 +253,7 @@ async def mock_consume_generator():
252
253
253
254
# Mock _continue_consuming to check if it's called by create_task
254
255
self .aggregator ._continue_consuming = AsyncMock ()
256
+ mock_create_task .side_effect = lambda coro : asyncio .ensure_future (coro )
255
257
256
258
(
257
259
result ,
@@ -303,6 +305,7 @@ async def mock_consume_generator():
303
305
current_task_state_after_update
304
306
)
305
307
self .aggregator ._continue_consuming = AsyncMock ()
308
+ mock_create_task .side_effect = lambda coro : asyncio .ensure_future (coro )
306
309
307
310
(
308
311
result ,
@@ -417,6 +420,7 @@ async def initial_consume_generator():
417
420
self .mock_task_manager .get_task .return_value = (
418
421
auth_event # Task state at interrupt
419
422
)
423
+ mock_create_task .side_effect = lambda coro : asyncio .ensure_future (coro )
420
424
421
425
# Call the main method that triggers _continue_consuming via create_task
422
426
_ , _ = await self .aggregator .consume_and_break_on_interrupt (
You can’t perform that action at this time.
0 commit comments