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_get_nowait(tasks, get_size, expected_tasks):
q = TaskQueue()
await q.add(tasks)
batch, tasks = q.get_nowait(get_size)
assert tasks == expected_tasks
q.complete(batch, tasks)
assert all(task not in q for task in tasks)
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