async def test_async_all_done_pre_flag():
pool = CompletionOrderedAsyncWorkPool()
done = pool.async_all_done()
assert await cirq.testing.asynchronous.asyncio_pending(done)
pool.set_all_work_received_flag()
assert await done is None
def test_done_then_dequeue():
pool = CompletionOrderedAsyncWorkPool()
pool.set_all_work_received_flag()
result = pool.__anext__()
cirq.testing.assert_asyncio_will_raise(result,
StopAsyncIteration,
match='no_more_work')
def test_async_all_done_pre_flag():
pool = CompletionOrderedAsyncWorkPool()
done = pool.async_all_done()
cirq.testing.assert_asyncio_still_running(done)
pool.set_all_work_received_flag()
cirq.testing.assert_asyncio_will_have_result(done, None)
async def test_dequeue_then_done():
pool = CompletionOrderedAsyncWorkPool()
result = pool.__anext__()
assert await cirq.testing.asynchronous.asyncio_pending(result)
pool.set_all_work_received_flag()
with pytest.raises(StopAsyncIteration, match='no_more_work'):
await result
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_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_async_for():
pool = CompletionOrderedAsyncWorkPool()
async def consume():
results = []
async for result in pool:
results.append(result)
return results
async def produce():
w1 = asyncio.Future()
w2 = asyncio.Future()
w3 = asyncio.Future()
pool.include_work(w1)
pool.include_work(w2)
w1.set_result(1)
pool.include_work(w3)
pool.set_all_work_received_flag()
w2.set_result(2)
w3.set_result(3)
a = asyncio.ensure_future(produce())
b = asyncio.ensure_future(consume())
cirq.testing.assert_asyncio_will_have_result(asyncio.wait([a, b]))
cirq.testing.assert_asyncio_will_have_result(a, None)
cirq.testing.assert_asyncio_will_have_result(b, [1, 2, 3])
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
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
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)
def _get_pool() -> CompletionOrderedAsyncWorkPool:
with cirq.testing.assert_deprecated("Use duet.AsyncCollector",
deadline="v0.14"):
return CompletionOrderedAsyncWorkPool()
def test_async_all_done_post_flag():
pool = CompletionOrderedAsyncWorkPool()
pool.set_all_work_received_flag()
cirq.testing.assert_asyncio_will_have_result(pool.async_all_done(), None)
cirq.testing.assert_asyncio_will_have_result(pool.async_all_done(), None)
async def test_async_all_done_post_flag():
pool = CompletionOrderedAsyncWorkPool()
pool.set_all_work_received_flag()
assert await pool.async_all_done() is None
assert await pool.async_all_done() is None
async def test_done_then_dequeue():
pool = CompletionOrderedAsyncWorkPool()
pool.set_all_work_received_flag()
result = pool.__anext__()
with pytest.raises(StopAsyncIteration, match='no_more_work'):
await result