This commit is contained in:
@@ -1,6 +1,7 @@
|
||||
import unittest
|
||||
|
||||
from src.pwo import retry, async_retry, async_test
|
||||
from pwo import retry, async_retry, async_test, AsyncQueueIterator, aenumerate
|
||||
from asyncio import Queue
|
||||
|
||||
|
||||
class PrivateTest(unittest.TestCase):
|
||||
@@ -70,5 +71,25 @@ class PrivateTest(unittest.TestCase):
|
||||
await bar()
|
||||
self.assertEqual(max_attempts, attempt)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@async_test
|
||||
async def test_async_queue_iterator(self):
|
||||
queue = Queue()
|
||||
queue_size = 10
|
||||
objects = [object() for _ in range(queue_size)]
|
||||
|
||||
async def poll() -> int:
|
||||
completed = 0
|
||||
async for i, obj in aenumerate(AsyncQueueIterator(queue)):
|
||||
self.assertIs(objects[i], obj)
|
||||
completed += 1
|
||||
return completed
|
||||
|
||||
handle = poll()
|
||||
|
||||
for o in objects:
|
||||
queue.put_nowait(o)
|
||||
queue.put_nowait(None)
|
||||
processed = await handle
|
||||
self.assertEqual(queue_size, processed)
|
||||
|
||||
|
||||
|
Reference in New Issue
Block a user