async def test_unassigned_items():
q = SmartQueue([1, 2, 3])
with q.new_consumer() as c:
async for handle in c:
assert q.has_new_items() == q.has_unassigned_items()
if not q.has_unassigned_items():
assert handle.data == 3
break
assert not q.has_unassigned_items()
await q.reschedule_all(c)
assert q.has_unassigned_items()
assert not q.has_new_items()
async def test_unassigned_items():
"""Test the `SmartQueue.has_unassigned_items()` method."""
q = SmartQueue(async_iter([1, 2, 3]))
with q.new_consumer() as c:
await asyncio.sleep(0.1)
assert q.has_unassigned_items()
async for handle in c:
assert not q.finished()
await asyncio.sleep(0.1)
if not q.has_unassigned_items():
assert handle.data == 3
break
# All items are in progress, `has_unassigned_items()` should return `False`
assert not q.has_unassigned_items()
await q.reschedule_all(c)
# Now the items are unassigned again
assert q.has_unassigned_items()
# Queue is still not finished
assert not q.finished()
await q.close()
async def test_async_task_iterator(task_iterator_interval, worker_interval, executor_interval):
"""Check that the queue waits until new items appear."""
inputs = list(range(5))
yielded = []
async def tasks():
for n in inputs:
print(f"Yielding {n}")
yield n
yielded.append(n)
await asyncio.sleep(task_iterator_interval)
q = SmartQueue(tasks())
loop = asyncio.get_event_loop()
# A sample worker that accepts a task quickly and then exits
async def worker():
print("Started new worker")
with q.new_consumer() as consumer:
async for handle in consumer:
await asyncio.sleep(worker_interval)
await q.mark_done(handle)
print(f"Exiting after task {handle.data}")
return
# Simulate how Executor works: spawn new workers until all items are handled
worker_task = None
done_task = loop.create_task(q.wait_until_done())
while not done_task.done():
if (worker_task is None or worker_task.done()) and q.has_unassigned_items():
worker_task = loop.create_task(worker())
await asyncio.sleep(executor_interval)
assert yielded == inputs
await q.close()