def test_dequeue_then_enqueue():
pool = CompletionOrderedAsyncWorkPool()
work = asyncio.Future()
result = pool.__anext__()
pool.include_work(work)
cirq.testing.assert_asyncio_still_running(result)
work.set_result(5)
cirq.testing.assert_asyncio_will_have_result(result, 5)
async def test_dequeue_then_enqueue():
pool = CompletionOrderedAsyncWorkPool()
work = asyncio.Future()
result = pool.__anext__()
pool.include_work(work)
assert await cirq.testing.asynchronous.asyncio_pending(result)
work.set_result(5)
assert await result == 5
def test_async_all_done_flag_then_finish_work():
pool = CompletionOrderedAsyncWorkPool()
done = pool.async_all_done()
work = asyncio.Future()
pool.include_work(work)
pool.set_all_work_received_flag()
cirq.testing.assert_asyncio_still_running(done)
work.set_result(5)
cirq.testing.assert_asyncio_will_have_result(done, None)
async def test_async_all_done_flag_then_finish_work():
pool = CompletionOrderedAsyncWorkPool()
done = pool.async_all_done()
work = asyncio.Future()
pool.include_work(work)
pool.set_all_work_received_flag()
assert await cirq.testing.asynchronous.asyncio_pending(done)
work.set_result(5)
assert await done is None
def test_enqueue_then_dequeue_with_failure():
pool = CompletionOrderedAsyncWorkPool()
assert pool.num_active == 0
work = asyncio.Future()
pool.include_work(work)
assert pool.num_active == 1
result = pool.__anext__()
cirq.testing.assert_asyncio_still_running(result)
assert pool.num_active == 1
work.set_exception(ValueError('test'))
cirq.testing.assert_asyncio_will_raise(result, ValueError, match='test')
assert pool.num_active == 0
async def test_enqueue_then_dequeue_with_failure():
pool = CompletionOrderedAsyncWorkPool()
assert pool.num_active == 0
work = asyncio.Future()
pool.include_work(work)
assert pool.num_active == 1
result = pool.__anext__()
assert await cirq.testing.asynchronous.asyncio_pending(result)
assert pool.num_active == 1
work.set_exception(ValueError('test'))
with pytest.raises(ValueError, match='test'):
await result
assert pool.num_active == 0
def test_ordering():
pool = CompletionOrderedAsyncWorkPool()
w1 = asyncio.Future()
w2 = asyncio.Future()
w3 = asyncio.Future()
pool.include_work(w1)
pool.include_work(w2)
r1 = pool.__anext__()
pool.include_work(w3)
r2 = pool.__anext__()
r3 = pool.__anext__()
cirq.testing.assert_asyncio_still_running(r1)
w2.set_result(6)
cirq.testing.assert_asyncio_will_have_result(r1, 6)
cirq.testing.assert_asyncio_still_running(r2)
w1.set_result(7)
cirq.testing.assert_asyncio_will_have_result(r2, 7)
cirq.testing.assert_asyncio_still_running(r3)
w3.set_result(8)
cirq.testing.assert_asyncio_will_have_result(r3, 8)
async def test_ordering():
pool = CompletionOrderedAsyncWorkPool()
w1 = asyncio.Future()
w2 = asyncio.Future()
w3 = asyncio.Future()
pool.include_work(w1)
pool.include_work(w2)
r1 = pool.__anext__()
pool.include_work(w3)
r2 = pool.__anext__()
r3 = pool.__anext__()
assert await cirq.testing.asynchronous.asyncio_pending(r1)
w2.set_result(6)
assert await r1 == 6
assert await cirq.testing.asynchronous.asyncio_pending(r2)
w1.set_result(7)
assert await r2 == 7
assert await cirq.testing.asynchronous.asyncio_pending(r3)
w3.set_result(8)
assert await r3 == 8