async def test_unfinished_tasks_readded():
q = TaskQueue()
await wait(q.add((2, 1, 3)))
batch, tasks = await wait(q.get())
q.complete(batch, (2, ))
batch, tasks = await wait(q.get())
assert tasks == (1, 3)
async def test_wait_empty_queue():
q = TaskQueue()
try:
await wait(q.get())
except asyncio.TimeoutError:
pass
else:
assert False, "should not return from get() when nothing is available on queue"
async def test_queue_get_cap(start_tasks, get_max, expected, remainder):
q = TaskQueue()
await wait(q.add(start_tasks))
batch, tasks = await wait(q.get(get_max))
assert tasks == expected
if remainder:
_, tasks2 = await wait(q.get())
assert tasks2 == remainder
else:
try:
_, tasks2 = await wait(q.get())
except asyncio.TimeoutError:
pass
else:
assert False, f"No more tasks to get, but got {tasks2!r}"
async def test_cannot_complete_batch_with_wrong_task():
q = TaskQueue()
await wait(q.add((1, 2)))
batch, tasks = await wait(q.get())
# cannot complete a valid task with a task it wasn't given
with pytest.raises(ValidationError):
q.complete(batch, (3, 4))
# partially invalid completion calls leave the valid task in an incomplete state
with pytest.raises(ValidationError):
q.complete(batch, (1, 3))
assert 1 in q
async def test_cannot_complete_batch_unless_pending():
q = TaskQueue()
await wait(q.add((1, 2)))
# cannot complete a valid task without a batch id
with pytest.raises(ValidationError):
q.complete(None, (1, 2))
assert 1 in q
batch, tasks = await wait(q.get())
# cannot complete a valid task with an invalid batch id
with pytest.raises(ValidationError):
q.complete(batch + 1, (1, 2))
assert 1 in q
async def test_queue_contains_task_until_complete(tasks):
q = TaskQueue(order_fn=id)
first_task = tasks[0]
assert first_task not in q
await wait(q.add(tasks))
assert first_task in q
batch, pending_tasks = await wait(q.get())
assert first_task in q
q.complete(batch, pending_tasks)
assert first_task not in q
async def test_queue_size_reset_after_complete():
q = TaskQueue(maxsize=2)
await wait(q.add((1, 2)))
batch, tasks = await wait(q.get())
# there should not be room to add another task
try:
await wait(q.add((3, )))
except asyncio.TimeoutError:
pass
else:
assert False, "should not be able to add task past maxsize"
# do imaginary work here, then complete it all
q.complete(batch, tasks)
# there should be room to add more now
await wait(q.add((3, )))
async def test_two_pending_adds_one_release():
q = TaskQueue(2)
asyncio.ensure_future(q.add((3, 1, 2)))
# wait for ^ to run and pause
await asyncio.sleep(0)
# note that the highest-priority items are queued first
assert 1 in q
assert 2 in q
assert 3 not in q
# two tasks are queued, none are started
assert len(q) == 2
assert q.num_in_progress() == 0
asyncio.ensure_future(q.add((0, 4)))
# wait for ^ to run and pause
await asyncio.sleep(0)
# task consumer 1 completes the first two pending
batch, tasks = await wait(q.get())
assert tasks == (1, 2)
# both tasks started
assert len(q) == 2
assert q.num_in_progress() == 2
q.complete(batch, tasks)
# tasks are drained, but new ones aren't added yet...
assert q.num_in_progress() == 0
assert len(q) == 0
await asyncio.sleep(0.01)
# Now the tasks are added
assert q.num_in_progress() == 0
assert len(q) == 2
# task consumer 2 gets the next two, in priority order
batch, tasks = await wait(q.get())
assert len(tasks) == 2
assert tasks == (0, 3)
assert q.num_in_progress() == 2
assert len(q) == 2
# clean up, so the pending get() call can complete
q.complete(batch, tasks)
# All current tasks finished
assert q.num_in_progress() == 0
await asyncio.sleep(0)
# only task 4 remains
assert q.num_in_progress() == 0
assert len(q) == 1
async def test_custom_priority_order():
q = TaskQueue(maxsize=4, order_fn=lambda x: 0 - x)
await wait(q.add((2, 1, 3)))
(batch, tasks) = await wait(q.get())
assert tasks == (3, 2, 1)
async def test_default_priority_order():
q = TaskQueue(maxsize=4)
await wait(q.add((2, 1, 3)))
(batch, tasks) = await wait(q.get())
assert tasks == (1, 2, 3)