async def test_subscriber_failure(subscriber, assigner, subscriber_factory):
assign_queues = wire_queues(assigner.get_assignment)
async with subscriber:
await assign_queues.called.get()
sub1 = mock_async_context_manager(
MagicMock(spec=AsyncSingleSubscriber))
sub1_queues = wire_queues(sub1.read)
subscriber_factory.return_value = sub1
await assign_queues.results.put({Partition(1)})
await sub1_queues.called.get()
await sub1_queues.results.put(FailedPrecondition("sub failed"))
with pytest.raises(FailedPrecondition):
await subscriber.read()
async def test_iterator(
default_subscriber,
subscriber_factory,
multiplexed_client: AsyncSubscriberClientInterface,
):
read_queues = wire_queues(default_subscriber.read)
subscription = SubscriptionPath(1, CloudZone.parse("us-central1-a"), "abc")
message = Message(PubsubMessage(message_id="1")._pb, "", 0, None)
async with multiplexed_client:
iterator = await multiplexed_client.subscribe(
subscription, DISABLED_FLOW_CONTROL
)
subscriber_factory.assert_has_calls(
[call(subscription, None, DISABLED_FLOW_CONTROL)]
)
read_fut_1 = asyncio.ensure_future(iterator.__anext__())
assert not read_fut_1.done()
await read_queues.called.get()
default_subscriber.read.assert_has_calls([call()])
await read_queues.results.put(message)
assert await read_fut_1 is message
read_fut_2 = asyncio.ensure_future(iterator.__anext__())
assert not read_fut_2.done()
await read_queues.called.get()
default_subscriber.read.assert_has_calls([call(), call()])
await read_queues.results.put(FailedPrecondition(""))
with pytest.raises(FailedPrecondition):
await read_fut_2
default_subscriber.__aexit__.assert_called_once()
async def test_assigner_failure(subscriber, assigner, subscriber_factory):
assign_queues = wire_queues(assigner.get_assignment)
async with subscriber:
await assign_queues.called.get()
await assign_queues.results.put(FailedPrecondition("bad assign"))
with pytest.raises(FailedPrecondition):
await subscriber.read()
async def test_init(subscriber, assigner):
assign_queues = wire_queues(assigner.get_assignment)
async with subscriber:
assigner.__aenter__.assert_called_once()
await assign_queues.called.get()
assigner.get_assignment.assert_called_once()
assigner.__aexit__.assert_called_once()
async def test_assignment_change(subscriber, assigner, subscriber_factory):
assign_queues = wire_queues(assigner.get_assignment)
async with subscriber:
await assign_queues.called.get()
sub1 = mock_async_context_manager(
MagicMock(spec=AsyncSingleSubscriber))
sub2 = mock_async_context_manager(
MagicMock(spec=AsyncSingleSubscriber))
sub3 = mock_async_context_manager(
MagicMock(spec=AsyncSingleSubscriber))
subscriber_factory.side_effect = (
lambda partition: sub1 if partition == Partition(1) else sub2
if partition == Partition(2) else sub3)
await assign_queues.results.put({Partition(1), Partition(2)})
await assign_queues.called.get()
subscriber_factory.assert_has_calls(
[call(Partition(1)), call(Partition(2))], any_order=True)
sub1.__aenter__.assert_called_once()
sub2.__aenter__.assert_called_once()
await assign_queues.results.put({Partition(1), Partition(3)})
await assign_queues.called.get()
subscriber_factory.assert_has_calls(
[call(Partition(1)),
call(Partition(2)),
call(Partition(3))],
any_order=True)
sub3.__aenter__.assert_called_once()
sub2.__aexit__.assert_called_once()
sub1.__aexit__.assert_called_once()
sub2.__aexit__.assert_called_once()
sub3.__aexit__.assert_called_once()
async def test_delivery_from_multiple(subscriber, assigner,
subscriber_factory):
assign_queues = wire_queues(assigner.get_assignment)
async with subscriber:
await assign_queues.called.get()
sub1 = mock_async_context_manager(
MagicMock(spec=AsyncSingleSubscriber))
sub2 = mock_async_context_manager(
MagicMock(spec=AsyncSingleSubscriber))
sub1_queues = wire_queues(sub1.read)
sub2_queues = wire_queues(sub2.read)
subscriber_factory.side_effect = (
lambda partition: sub1 if partition == Partition(1) else sub2)
await assign_queues.results.put({Partition(1), Partition(2)})
await sub1_queues.results.put(
Message(PubsubMessage(message_id="1")._pb, "", 0, None))
await sub2_queues.results.put(
Message(PubsubMessage(message_id="2")._pb, "", 0, None))
message_ids: Set[str] = set()
message_ids.add((await subscriber.read()).message_id)
message_ids.add((await subscriber.read()).message_id)
assert message_ids == {"1", "2"}
async def test_reinitialize_after_retryable(
retrying_connection: Connection[int, int],
reinitializer,
default_connection,
asyncio_sleep,
):
reinit_queues = wire_queues(reinitializer.reinitialize)
default_connection.read.return_value = 1
await reinit_queues.results.put(InternalServerError("abc"))
await reinit_queues.results.put(None)
async with retrying_connection as _:
asyncio_sleep.assert_called_once_with(_MIN_BACKOFF_SECS)
assert reinitializer.reinitialize.call_count == 2
assert await retrying_connection.read() == 1
assert (
default_connection.read.call_count == 2
) # re-call to read once first completes
async def test_publish_after_increase(
mock_publishers, mock_policies, mock_watcher, publisher
):
get_queues = wire_queues(mock_watcher.get_partition_count)
await get_queues.results.put(2)
async with publisher:
get_queues.called.get_nowait()
mock_policies[2].route.return_value = Partition(1)
mock_publishers[Partition(1)].publish.return_value = "a"
await publisher.publish(PubSubMessage())
mock_policies[2].route.assert_called_with(PubSubMessage())
mock_publishers[Partition(1)].publish.assert_called()
await get_queues.called.get()
await get_queues.results.put(3)
await get_queues.called.get()
mock_policies[3].route.return_value = Partition(2)
mock_publishers[Partition(2)].publish.return_value = "a"
await publisher.publish(PubSubMessage())
mock_policies[3].route.assert_called_with(PubSubMessage())
mock_publishers[Partition(2)].publish.assert_called()