test_streaming_pull_manager.py

# Copyright 2018, Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import threading
import time
import types as stdlib_types

import mock
import pytest

from google.api_core import bidi
from google.api_core import exceptions
from google.cloud.pubsub_v1 import types
from google.cloud.pubsub_v1.subscriber import client
from google.cloud.pubsub_v1.subscriber import message
from google.cloud.pubsub_v1.subscriber import scheduler
from google.cloud.pubsub_v1.subscriber._protocol import dispatcher
from google.cloud.pubsub_v1.subscriber._protocol import heartbeater
from google.cloud.pubsub_v1.subscriber._protocol import leaser
from google.cloud.pubsub_v1.subscriber._protocol import messages_on_hold
from google.cloud.pubsub_v1.subscriber._protocol import requests
from google.cloud.pubsub_v1.subscriber._protocol import streaming_pull_manager
from google.pubsub_v1 import types as gapic_types
import grpc


@pytest.mark.parametrize(
    "exception,expected_cls",
    [
        (ValueError("meep"), ValueError),
        (
            mock.create_autospec(grpc.RpcError, instance=True),
            exceptions.GoogleAPICallError,
        ),
        ({"error": "RPC terminated"}, Exception),
        ("something broke", Exception),
    ],
)
def test__wrap_as_exception(exception, expected_cls):
    assert isinstance(
        streaming_pull_manager._wrap_as_exception(exception), expected_cls
    )


def test__wrap_callback_errors_no_error():
    msg = mock.create_autospec(message.Message, instance=True)
    callback = mock.Mock()
    on_callback_error = mock.Mock()

    streaming_pull_manager._wrap_callback_errors(callback, on_callback_error, msg)

    callback.assert_called_once_with(msg)
    msg.nack.assert_not_called()
    on_callback_error.assert_not_called()


def test__wrap_callback_errors_error():
    callback_error = ValueError("meep")

    msg = mock.create_autospec(message.Message, instance=True)
    callback = mock.Mock(side_effect=callback_error)
    on_callback_error = mock.Mock()

    streaming_pull_manager._wrap_callback_errors(callback, on_callback_error, msg)

    msg.nack.assert_called_once()
    on_callback_error.assert_called_once_with(callback_error)


def test_constructor_and_default_state():
    manager = streaming_pull_manager.StreamingPullManager(
        mock.sentinel.client, mock.sentinel.subscription
    )

    # Public state
    assert manager.is_active is False
    assert manager.flow_control == types.FlowControl()
    assert manager.dispatcher is None
    assert manager.leaser is None
    assert manager.ack_histogram is not None
    assert manager.ack_deadline == 10
    assert manager.load == 0

    # Private state
    assert manager._client == mock.sentinel.client
    assert manager._subscription == mock.sentinel.subscription
    assert manager._scheduler is not None
    assert manager._messages_on_hold is not None
    assert manager._client_id is not None


def test_constructor_with_options():
    manager = streaming_pull_manager.StreamingPullManager(
        mock.sentinel.client,
        mock.sentinel.subscription,
        flow_control=mock.sentinel.flow_control,
        scheduler=mock.sentinel.scheduler,
    )

    assert manager.flow_control == mock.sentinel.flow_control
    assert manager._scheduler == mock.sentinel.scheduler


def make_manager(**kwargs):
    client_ = mock.create_autospec(client.Client, instance=True)
    scheduler_ = mock.create_autospec(scheduler.Scheduler, instance=True)
    return streaming_pull_manager.StreamingPullManager(
        client_, "subscription-name", scheduler=scheduler_, **kwargs
    )


def fake_leaser_add(leaser, init_msg_count=0, assumed_msg_size=10):
    """Add a simplified fake add() method to a leaser instance.

    The fake add() method actually increases the leaser's internal message count
    by one for each message, and the total bytes by ``assumed_msg_size`` for
    each message (regardless of the actual message size).
    """

    def fake_add(self, items):
        self.message_count += len(items)
        self.bytes += len(items) * assumed_msg_size

    leaser.message_count = init_msg_count
    leaser.bytes = init_msg_count * assumed_msg_size
    leaser.add = stdlib_types.MethodType(fake_add, leaser)


def test_ack_deadline():
    manager = make_manager()
    assert manager.ack_deadline == 10
    manager.ack_histogram.add(20)
    assert manager.ack_deadline == 20
    manager.ack_histogram.add(10)
    assert manager.ack_deadline == 20


def test_client_id():
    manager1 = make_manager()
    request1 = manager1._get_initial_request(stream_ack_deadline_seconds=10)
    client_id_1 = request1.client_id
    assert client_id_1

    manager2 = make_manager()
    request2 = manager2._get_initial_request(stream_ack_deadline_seconds=10)
    client_id_2 = request2.client_id
    assert client_id_2

    assert client_id_1 != client_id_2


def test_streaming_flow_control():
    manager = make_manager(
        flow_control=types.FlowControl(max_messages=10, max_bytes=1000)
    )
    request = manager._get_initial_request(stream_ack_deadline_seconds=10)
    assert request.max_outstanding_messages == 10
    assert request.max_outstanding_bytes == 1000


def test_streaming_flow_control_use_legacy_flow_control():
    manager = make_manager(
        flow_control=types.FlowControl(max_messages=10, max_bytes=1000),
        use_legacy_flow_control=True,
    )
    request = manager._get_initial_request(stream_ack_deadline_seconds=10)
    assert request.max_outstanding_messages == 0
    assert request.max_outstanding_bytes == 0


def test_ack_deadline_with_max_duration_per_lease_extension():
    manager = make_manager()
    manager._flow_control = types.FlowControl(max_duration_per_lease_extension=5)

    assert manager.ack_deadline == 5
    for _ in range(5):
        manager.ack_histogram.add(20)

    assert manager.ack_deadline == 5


def test_maybe_pause_consumer_wo_consumer_set():
    manager = make_manager(
        flow_control=types.FlowControl(max_messages=10, max_bytes=1000)
    )
    manager.maybe_pause_consumer()  # no raise
    # Ensure load > 1
    _leaser = manager._leaser = mock.create_autospec(leaser.Leaser)
    _leaser.message_count = 100
    _leaser.bytes = 10000
    manager.maybe_pause_consumer()  # no raise


def test_lease_load_and_pause():
    manager = make_manager(
        flow_control=types.FlowControl(max_messages=10, max_bytes=1000)
    )
    manager._leaser = leaser.Leaser(manager)
    manager._consumer = mock.create_autospec(bidi.BackgroundConsumer, instance=True)
    manager._consumer.is_paused = False

    # This should mean that our messages count is at 10%, and our bytes
    # are at 15%; load should return the higher (0.15), and shouldn't cause
    # the consumer to pause.
    manager.leaser.add(
        [requests.LeaseRequest(ack_id="one", byte_size=150, ordering_key="")]
    )
    assert manager.load == 0.15
    manager.maybe_pause_consumer()
    manager._consumer.pause.assert_not_called()

    # After this message is added, the messages should be higher at 20%
    # (versus 16% for bytes).
    manager.leaser.add(
        [requests.LeaseRequest(ack_id="two", byte_size=10, ordering_key="")]
    )
    assert manager.load == 0.2

    # Returning a number above 100% is fine, and it should cause this to pause.
    manager.leaser.add(
        [requests.LeaseRequest(ack_id="three", byte_size=1000, ordering_key="")]
    )
    assert manager.load == 1.16
    manager.maybe_pause_consumer()
    manager._consumer.pause.assert_called_once()


def test_drop_and_resume():
    manager = make_manager(
        flow_control=types.FlowControl(max_messages=10, max_bytes=1000)
    )
    manager._leaser = leaser.Leaser(manager)
    manager._consumer = mock.create_autospec(bidi.BackgroundConsumer, instance=True)
    manager._consumer.is_paused = True

    # Add several messages until we're over the load threshold.
    manager.leaser.add(
        [
            requests.LeaseRequest(ack_id="one", byte_size=750, ordering_key=""),
            requests.LeaseRequest(ack_id="two", byte_size=250, ordering_key=""),
        ]
    )

    assert manager.load == 1.0

    # Trying to resume now should have no effect as we're over the threshold.
    manager.maybe_resume_consumer()
    manager._consumer.resume.assert_not_called()

    # Drop the 200 byte message, which should put us under the resume
    # threshold.
    manager.leaser.remove(
        [requests.DropRequest(ack_id="two", byte_size=250, ordering_key="")]
    )
    manager.maybe_resume_consumer()
    manager._consumer.resume.assert_called_once()


def test_resume_not_paused():
    manager = make_manager()
    manager._consumer = mock.create_autospec(bidi.BackgroundConsumer, instance=True)
    manager._consumer.is_paused = False

    # Resuming should have no effect is the consumer is not actually paused.
    manager.maybe_resume_consumer()
    manager._consumer.resume.assert_not_called()


def test_maybe_resume_consumer_wo_consumer_set():
    manager = make_manager(
        flow_control=types.FlowControl(max_messages=10, max_bytes=1000)
    )
    manager.maybe_resume_consumer()  # no raise


def test__maybe_release_messages_on_overload():
    manager = make_manager(
        flow_control=types.FlowControl(max_messages=10, max_bytes=1000)
    )

    msg = mock.create_autospec(message.Message, instance=True, ack_id="ack", size=11)
    manager._messages_on_hold.put(msg)
    manager._on_hold_bytes = msg.size

    # Ensure load is exactly 1.0 (to verify that >= condition is used)
    _leaser = manager._leaser = mock.create_autospec(leaser.Leaser)
    _leaser.message_count = 10
    _leaser.bytes = 1000 + msg.size

    manager._maybe_release_messages()

    assert manager._messages_on_hold.size == 1
    manager._leaser.add.assert_not_called()
    manager._scheduler.schedule.assert_not_called()


def test__maybe_release_messages_below_overload():
    manager = make_manager(
        flow_control=types.FlowControl(max_messages=10, max_bytes=1000)
    )
    manager._callback = mock.sentinel.callback

    # Init leaser message count to 11, so that when subtracting the 3 messages
    # that are on hold, there is still room for another 2 messages before the
    # max load is hit.
    _leaser = manager._leaser = mock.create_autospec(leaser.Leaser)
    fake_leaser_add(_leaser, init_msg_count=11, assumed_msg_size=10)

    messages = [
        mock.create_autospec(message.Message, instance=True, ack_id="ack_foo", size=10),
        mock.create_autospec(message.Message, instance=True, ack_id="ack_bar", size=10),
        mock.create_autospec(message.Message, instance=True, ack_id="ack_baz", size=10),
    ]
    for msg in messages:
        manager._messages_on_hold.put(msg)
        manager._on_hold_bytes = 3 * 10

    # the actual call of MUT
    manager._maybe_release_messages()

    assert manager._messages_on_hold.size == 1
    msg = manager._messages_on_hold.get()
    assert msg.ack_id == "ack_baz"

    schedule_calls = manager._scheduler.schedule.mock_calls
    assert len(schedule_calls) == 2
    for _, call_args, _ in schedule_calls:
        assert call_args[0] == mock.sentinel.callback
        assert isinstance(call_args[1], message.Message)
        assert call_args[1].ack_id in ("ack_foo", "ack_bar")


def test__maybe_release_messages_negative_on_hold_bytes_warning(caplog):
    manager = make_manager(
        flow_control=types.FlowControl(max_messages=10, max_bytes=1000)
    )

    msg = mock.create_autospec(message.Message, instance=True, ack_id="ack", size=17)
    manager._messages_on_hold.put(msg)
    manager._on_hold_bytes = 5  # too low for some reason

    _leaser = manager._leaser = mock.create_autospec(leaser.Leaser)
    _leaser.message_count = 3
    _leaser.bytes = 150

    with caplog.at_level(logging.WARNING):
        manager._maybe_release_messages()

    expected_warnings = [
        record.message.lower()
        for record in caplog.records
        if "unexpectedly negative" in record.message
    ]
    assert len(expected_warnings) == 1
    assert "on hold bytes" in expected_warnings[0]
    assert "-12" in expected_warnings[0]

    assert manager._on_hold_bytes == 0  # should be auto-corrected


def test_send_unary():
    manager = make_manager()
    manager._UNARY_REQUESTS = True

    manager.send(
        gapic_types.StreamingPullRequest(
            ack_ids=["ack_id1", "ack_id2"],
            modify_deadline_ack_ids=["ack_id3", "ack_id4", "ack_id5"],
            modify_deadline_seconds=[10, 20, 20],
        )
    )

    manager._client.acknowledge.assert_called_once_with(
        subscription=manager._subscription, ack_ids=["ack_id1", "ack_id2"]
    )

    manager._client.modify_ack_deadline.assert_has_calls(
        [
            mock.call(
                subscription=manager._subscription,
                ack_ids=["ack_id3"],
                ack_deadline_seconds=10,
            ),
            mock.call(
                subscription=manager._subscription,
                ack_ids=["ack_id4", "ack_id5"],
                ack_deadline_seconds=20,
            ),
        ],
        any_order=True,
    )


def test_send_unary_empty():
    manager = make_manager()
    manager._UNARY_REQUESTS = True

    manager.send(gapic_types.StreamingPullRequest())

    manager._client.acknowledge.assert_not_called()
    manager._client.modify_ack_deadline.assert_not_called()


def test_send_unary_api_call_error(caplog):
    caplog.set_level(logging.DEBUG)

    manager = make_manager()
    manager._UNARY_REQUESTS = True

    error = exceptions.GoogleAPICallError("The front fell off")
    manager._client.acknowledge.side_effect = error

    manager.send(gapic_types.StreamingPullRequest(ack_ids=["ack_id1", "ack_id2"]))

    assert "The front fell off" in caplog.text


def test_send_unary_retry_error(caplog):
    caplog.set_level(logging.DEBUG)

    manager, _, _, _, _, _ = make_running_manager()
    manager._UNARY_REQUESTS = True

    error = exceptions.RetryError(
        "Too long a transient error", cause=Exception("Out of time!")
    )
    manager._client.acknowledge.side_effect = error

    with pytest.raises(exceptions.RetryError):
        manager.send(gapic_types.StreamingPullRequest(ack_ids=["ack_id1", "ack_id2"]))

    assert "RetryError while sending unary RPC" in caplog.text
    assert "signaled streaming pull manager shutdown" in caplog.text


def test_send_streaming():
    manager = make_manager()
    manager._UNARY_REQUESTS = False
    manager._rpc = mock.create_autospec(bidi.BidiRpc, instance=True)

    manager.send(mock.sentinel.request)

    manager._rpc.send.assert_called_once_with(mock.sentinel.request)


def test_heartbeat():
    manager = make_manager()
    manager._rpc = mock.create_autospec(bidi.BidiRpc, instance=True)
    manager._rpc.is_active = True

    manager.heartbeat()

    manager._rpc.send.assert_called_once_with(gapic_types.StreamingPullRequest())


def test_heartbeat_inactive():
    manager = make_manager()
    manager._rpc = mock.create_autospec(bidi.BidiRpc, instance=True)
    manager._rpc.is_active = False

    manager.heartbeat()

    manager._rpc.send.assert_not_called()


@mock.patch("google.api_core.bidi.ResumableBidiRpc", autospec=True)
@mock.patch("google.api_core.bidi.BackgroundConsumer", autospec=True)
@mock.patch("google.cloud.pubsub_v1.subscriber._protocol.leaser.Leaser", autospec=True)
@mock.patch(
    "google.cloud.pubsub_v1.subscriber._protocol.dispatcher.Dispatcher", autospec=True
)
@mock.patch(
    "google.cloud.pubsub_v1.subscriber._protocol.heartbeater.Heartbeater", autospec=True
)
def test_open(heartbeater, dispatcher, leaser, background_consumer, resumable_bidi_rpc):
    manager = make_manager()

    manager.open(mock.sentinel.callback, mock.sentinel.on_callback_error)

    heartbeater.assert_called_once_with(manager)
    heartbeater.return_value.start.assert_called_once()
    assert manager._heartbeater == heartbeater.return_value

    dispatcher.assert_called_once_with(manager, manager._scheduler.queue)
    dispatcher.return_value.start.assert_called_once()
    assert manager._dispatcher == dispatcher.return_value

    leaser.assert_called_once_with(manager)
    leaser.return_value.start.assert_called_once()
    assert manager.leaser == leaser.return_value

    background_consumer.assert_called_once_with(manager._rpc, manager._on_response)
    background_consumer.return_value.start.assert_called_once()
    assert manager._consumer == background_consumer.return_value

    resumable_bidi_rpc.assert_called_once_with(
        start_rpc=manager._client.api.streaming_pull,
        initial_request=mock.ANY,
        should_recover=manager._should_recover,
        should_terminate=manager._should_terminate,
        throttle_reopen=True,
    )
    initial_request_arg = resumable_bidi_rpc.call_args.kwargs["initial_request"]
    assert initial_request_arg.func == manager._get_initial_request
    assert initial_request_arg.args[0] == 10  # the default stream ACK timeout
    assert not manager._client.api.get_subscription.called

    resumable_bidi_rpc.return_value.add_done_callback.assert_called_once_with(
        manager._on_rpc_done
    )
    assert manager._rpc == resumable_bidi_rpc.return_value

    manager._consumer.is_active = True
    assert manager.is_active is True


def test_open_already_active():
    manager = make_manager()
    manager._consumer = mock.create_autospec(bidi.BackgroundConsumer, instance=True)
    manager._consumer.is_active = True

    with pytest.raises(ValueError, match="already open"):
        manager.open(mock.sentinel.callback, mock.sentinel.on_callback_error)


def test_open_has_been_closed():
    manager = make_manager()
    manager._closed = True

    with pytest.raises(ValueError, match="closed"):
        manager.open(mock.sentinel.callback, mock.sentinel.on_callback_error)


def make_running_manager():
    manager = make_manager()
    manager._consumer = mock.create_autospec(bidi.BackgroundConsumer, instance=True)
    manager._consumer.is_active = True
    manager._dispatcher = mock.create_autospec(dispatcher.Dispatcher, instance=True)
    manager._leaser = mock.create_autospec(leaser.Leaser, instance=True)
    manager._heartbeater = mock.create_autospec(heartbeater.Heartbeater, instance=True)

    return (
        manager,
        manager._consumer,
        manager._dispatcher,
        manager._leaser,
        manager._heartbeater,
        manager._scheduler,
    )


def test_close():
    (
        manager,
        consumer,
        dispatcher,
        leaser,
        heartbeater,
        scheduler,
    ) = make_running_manager()

    manager.close()

    consumer.stop.assert_called_once()
    leaser.stop.assert_called_once()
    dispatcher.stop.assert_called_once()
    heartbeater.stop.assert_called_once()
    scheduler.shutdown.assert_called_once()

    assert manager.is_active is False


def test_close_inactive_consumer():
    (
        manager,
        consumer,
        dispatcher,
        leaser,
        heartbeater,
        scheduler,
    ) = make_running_manager()
    consumer.is_active = False

    manager.close()

    consumer.stop.assert_not_called()
    leaser.stop.assert_called_once()
    dispatcher.stop.assert_called_once()
    heartbeater.stop.assert_called_once()
    scheduler.shutdown.assert_called_once()


def test_close_idempotent():
    manager, _, _, _, _, scheduler = make_running_manager()

    manager.close()
    manager.close()

    assert scheduler.shutdown.call_count == 1


class FakeDispatcher(object):
    def __init__(self, manager, error_callback):
        self._manager = manager
        self._error_callback = error_callback
        self._thread = None
        self._stop = False

    def start(self):
        self._thread = threading.Thread(target=self._do_work)
        self._thread.daemon = True
        self._thread.start()

    def stop(self):
        self._stop = True
        self._thread.join()
        self._thread = None

    def _do_work(self):
        while not self._stop:
            try:
                self._manager.leaser.add([mock.Mock()])
            except Exception as exc:
                self._error_callback(exc)
            time.sleep(0.1)

        # also try to interact with the leaser after the stop flag has been set
        try:
            self._manager.leaser.remove([mock.Mock()])
        except Exception as exc:
            self._error_callback(exc)


def test_close_no_dispatcher_error():
    manager, _, _, _, _, _ = make_running_manager()
    error_callback = mock.Mock(name="error_callback")
    dispatcher = FakeDispatcher(manager=manager, error_callback=error_callback)
    manager._dispatcher = dispatcher
    dispatcher.start()

    manager.close()

    error_callback.assert_not_called()


def test_close_callbacks():
    manager, _, _, _, _, _ = make_running_manager()

    callback = mock.Mock()

    manager.add_close_callback(callback)
    manager.close(reason="meep")

    callback.assert_called_once_with(manager, "meep")


def test__get_initial_request():
    manager = make_manager()
    manager._leaser = mock.create_autospec(leaser.Leaser, instance=True)
    manager._leaser.ack_ids = ["1", "2"]

    initial_request = manager._get_initial_request(123)

    assert isinstance(initial_request, gapic_types.StreamingPullRequest)
    assert initial_request.subscription == "subscription-name"
    assert initial_request.stream_ack_deadline_seconds == 123
    assert initial_request.modify_deadline_ack_ids == ["1", "2"]
    assert initial_request.modify_deadline_seconds == [10, 10]


def test__get_initial_request_wo_leaser():
    manager = make_manager()
    manager._leaser = None

    initial_request = manager._get_initial_request(123)

    assert isinstance(initial_request, gapic_types.StreamingPullRequest)
    assert initial_request.subscription == "subscription-name"
    assert initial_request.stream_ack_deadline_seconds == 123
    assert initial_request.modify_deadline_ack_ids == []
    assert initial_request.modify_deadline_seconds == []


def test__on_response_delivery_attempt():
    manager, _, dispatcher, leaser, _, scheduler = make_running_manager()
    manager._callback = mock.sentinel.callback

    # Set up the messages.
    response = gapic_types.StreamingPullResponse(
        received_messages=[
            gapic_types.ReceivedMessage(
                ack_id="fack",
                message=gapic_types.PubsubMessage(data=b"foo", message_id="1"),
            ),
            gapic_types.ReceivedMessage(
                ack_id="back",
                message=gapic_types.PubsubMessage(data=b"bar", message_id="2"),
                delivery_attempt=6,
            ),
        ]
    )

    # adjust message bookkeeping in leaser
    fake_leaser_add(leaser, init_msg_count=0, assumed_msg_size=42)

    manager._on_response(response)

    schedule_calls = scheduler.schedule.mock_calls
    assert len(schedule_calls) == 2
    msg1 = schedule_calls[0][1][1]
    assert msg1.delivery_attempt is None
    msg2 = schedule_calls[1][1][1]
    assert msg2.delivery_attempt == 6


def test__on_response_no_leaser_overload():
    manager, _, dispatcher, leaser, _, scheduler = make_running_manager()
    manager._callback = mock.sentinel.callback

    # Set up the messages.
    response = gapic_types.StreamingPullResponse(
        received_messages=[
            gapic_types.ReceivedMessage(
                ack_id="fack",
                message=gapic_types.PubsubMessage(data=b"foo", message_id="1"),
            ),
            gapic_types.ReceivedMessage(
                ack_id="back",
                message=gapic_types.PubsubMessage(data=b"bar", message_id="2"),
            ),
        ]
    )

    # adjust message bookkeeping in leaser
    fake_leaser_add(leaser, init_msg_count=0, assumed_msg_size=42)

    # Actually run the method and prove that modack and schedule
    # are called in the expected way.
    manager._on_response(response)

    dispatcher.modify_ack_deadline.assert_called_once_with(
        [requests.ModAckRequest("fack", 10), requests.ModAckRequest("back", 10)]
    )

    schedule_calls = scheduler.schedule.mock_calls
    assert len(schedule_calls) == 2
    for call in schedule_calls:
        assert call[1][0] == mock.sentinel.callback
        assert isinstance(call[1][1], message.Message)

    # the leaser load limit not hit, no messages had to be put on hold
    assert manager._messages_on_hold.size == 0


def test__on_response_with_leaser_overload():
    manager, _, dispatcher, leaser, _, scheduler = make_running_manager()
    manager._callback = mock.sentinel.callback

    # Set up the messages.
    response = gapic_types.StreamingPullResponse(
        received_messages=[
            gapic_types.ReceivedMessage(
                ack_id="fack",
                message=gapic_types.PubsubMessage(data=b"foo", message_id="1"),
            ),
            gapic_types.ReceivedMessage(
                ack_id="back",
                message=gapic_types.PubsubMessage(data=b"bar", message_id="2"),
            ),
            gapic_types.ReceivedMessage(
                ack_id="zack",
                message=gapic_types.PubsubMessage(data=b"baz", message_id="3"),
            ),
        ]
    )

    # Adjust message bookkeeping in leaser. Pick 999 messages, which is just below
    # the default FlowControl.max_messages limit.
    fake_leaser_add(leaser, init_msg_count=999, assumed_msg_size=10)

    # Actually run the method and prove that modack and schedule
    # are called in the expected way.
    manager._on_response(response)

    # all messages should be added to the lease management and have their ACK
    # deadline extended, even those not dispatched to callbacks
    dispatcher.modify_ack_deadline.assert_called_once_with(
        [
            requests.ModAckRequest("fack", 10),
            requests.ModAckRequest("back", 10),
            requests.ModAckRequest("zack", 10),
        ]
    )

    # one message should be scheduled, the flow control limits allow for it
    schedule_calls = scheduler.schedule.mock_calls
    assert len(schedule_calls) == 1
    call_args = schedule_calls[0][1]
    assert call_args[0] == mock.sentinel.callback
    assert isinstance(call_args[1], message.Message)
    assert call_args[1].message_id == "1"

    # the rest of the messages should have been put on hold
    assert manager._messages_on_hold.size == 2
    while True:
        msg = manager._messages_on_hold.get()
        if msg is None:
            break
        else:
            assert isinstance(msg, message.Message)
            assert msg.message_id in ("2", "3")


def test__on_response_none_data(caplog):
    caplog.set_level(logging.DEBUG)

    manager, _, dispatcher, leaser, _, scheduler = make_running_manager()
    manager._callback = mock.sentinel.callback

    # adjust message bookkeeping in leaser
    fake_leaser_add(leaser, init_msg_count=0, assumed_msg_size=10)

    manager._on_response(response=None)

    scheduler.schedule.assert_not_called()
    assert "callback invoked with None" in caplog.text


def test__on_response_with_ordering_keys():
    manager, _, dispatcher, leaser, _, scheduler = make_running_manager()
    manager._callback = mock.sentinel.callback

    # Set up the messages.
    response = gapic_types.StreamingPullResponse(
        received_messages=[
            gapic_types.ReceivedMessage(
                ack_id="fack",
                message=gapic_types.PubsubMessage(
                    data=b"foo", message_id="1", ordering_key=""
                ),
            ),
            gapic_types.ReceivedMessage(
                ack_id="back",
                message=gapic_types.PubsubMessage(
                    data=b"bar", message_id="2", ordering_key="key1"
                ),
            ),
            gapic_types.ReceivedMessage(
                ack_id="zack",
                message=gapic_types.PubsubMessage(
                    data=b"baz", message_id="3", ordering_key="key1"
                ),
            ),
        ]
    )

    # Make leaser with zero initial messages, so we don't test lease management
    # behavior.
    fake_leaser_add(leaser, init_msg_count=0, assumed_msg_size=10)

    # Actually run the method and prove that modack and schedule are called in
    # the expected way.
    manager._on_response(response)

    # All messages should be added to the lease management and have their ACK
    # deadline extended, even those not dispatched to callbacks.
    dispatcher.modify_ack_deadline.assert_called_once_with(
        [
            requests.ModAckRequest("fack", 10),
            requests.ModAckRequest("back", 10),
            requests.ModAckRequest("zack", 10),
        ]
    )

    # The first two messages should be scheduled, The third should be put on
    # hold because it's blocked by the completion of the second, which has the
    # same ordering key.
    schedule_calls = scheduler.schedule.mock_calls
    assert len(schedule_calls) == 2
    call_args = schedule_calls[0][1]
    assert call_args[0] == mock.sentinel.callback
    assert isinstance(call_args[1], message.Message)
    assert call_args[1].message_id == "1"

    call_args = schedule_calls[1][1]
    assert call_args[0] == mock.sentinel.callback
    assert isinstance(call_args[1], message.Message)
    assert call_args[1].message_id == "2"

    # Message 3 should have been put on hold.
    assert manager._messages_on_hold.size == 1
    # No messages available because message 2 (with "key1") has not completed yet.
    assert manager._messages_on_hold.get() is None

    # Complete message 2 (with "key1").
    manager.activate_ordering_keys(["key1"])

    # Completing message 2 should release message 3.
    schedule_calls = scheduler.schedule.mock_calls
    assert len(schedule_calls) == 3
    call_args = schedule_calls[2][1]
    assert call_args[0] == mock.sentinel.callback
    assert isinstance(call_args[1], message.Message)
    assert call_args[1].message_id == "3"

    # No messages available in the queue.
    assert manager._messages_on_hold.get() is None


def test__should_recover_true():
    manager = make_manager()

    details = "UNAVAILABLE. Service taking nap."
    exc = exceptions.ServiceUnavailable(details)

    assert manager._should_recover(exc) is True


def test__should_recover_false():
    manager = make_manager()

    exc = TypeError("wahhhhhh")

    assert manager._should_recover(exc) is False


def test__should_terminate_true():
    manager = make_manager()

    details = "Cancelled. Go away, before I taunt you a second time."
    exc = exceptions.Cancelled(details)

    assert manager._should_terminate(exc) is True


def test__should_terminate_false():
    manager = make_manager()

    exc = TypeError("wahhhhhh")

    assert manager._should_terminate(exc) is False


@mock.patch("threading.Thread", autospec=True)
def test__on_rpc_done(thread):
    manager = make_manager()

    manager._on_rpc_done(mock.sentinel.error)

    thread.assert_called_once_with(
        name=mock.ANY, target=manager.close, kwargs={"reason": mock.ANY}
    )
    _, kwargs = thread.call_args
    reason = kwargs["kwargs"]["reason"]
    assert isinstance(reason, Exception)
    assert reason.args == (mock.sentinel.error,)  # Exception wraps the original error


def test_activate_ordering_keys():
    manager = make_manager()
    manager._messages_on_hold = mock.create_autospec(
        messages_on_hold.MessagesOnHold, instance=True
    )

    manager.activate_ordering_keys(["key1", "key2"])

    manager._messages_on_hold.activate_ordering_keys.assert_called_once_with(
        ["key1", "key2"], mock.ANY
    )