def test_publish_new_batch_needed(creds): client = publisher.Client(credentials=creds) # Use mocks in lieu of the actual batch class. batch1 = mock.Mock(spec=client._batch_class) batch2 = mock.Mock(spec=client._batch_class) # Set the first mock up to claim indiscriminately that it rejects all # messages and the second accepts all. future = mock.sentinel.future future.add_done_callback = mock.Mock(spec=["__call__"]) batch1.publish.return_value = None batch2.publish.return_value = future topic = "topic/path" client._set_batch(topic, batch1) # Actually mock the batch class now. batch_class = mock.Mock(spec=(), return_value=batch2) client._set_batch_class(batch_class) # Publish a message. future = client.publish(topic, b"foo", bar=b"baz") assert future is mock.sentinel.future # Check the mocks. batch_class.assert_called_once_with( client=mock.ANY, topic=topic, settings=client.batch_settings, batch_done_callback=None, commit_when_full=True, commit_retry=gapic_v1.method.DEFAULT, ) message_pb = gapic_types.PubsubMessage(data=b"foo", attributes={"bar": "baz"}) batch1.publish.assert_called_once_with(message_pb) batch2.publish.assert_called_once_with(message_pb)
def test_batch_done_callback_called_on_publish_failure(): batch_done_callback_tracker = BatchDoneCallbackTracker() batch = create_batch(batch_done_callback=batch_done_callback_tracker) # Ensure messages exist. message = gapic_types.PubsubMessage(data=b"foobarbaz") batch.publish(message) # One response for one published message. publish_response = gapic_types.PublishResponse(message_ids=["a"]) # Induce publish error. error = google.api_core.exceptions.InternalServerError("uh oh") with mock.patch.object( type(batch.client.api), "publish", return_value=publish_response, side_effect=error, ): batch._commit() assert batch_done_callback_tracker.called assert not batch_done_callback_tracker.success
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_threads_posting_large_messages_do_not_starve(): settings = types.PublishFlowControl( message_limit=100, byte_limit=110, limit_exceeded_behavior=types.LimitExceededBehavior.BLOCK, ) flow_controller = FlowController(settings) large_msg = grpc_types.PubsubMessage(data=b"x" * 100) # close to entire byte limit adding_initial_done = threading.Event() adding_large_done = threading.Event() adding_busy_done = threading.Event() releasing_busy_done = threading.Event() releasing_large_done = threading.Event() # Occupy some of the flow capacity, then try to add a large message. Releasing # enough messages should eventually allow the large message to come through, even # if more messages are added after it (those should wait for the large message). initial_messages = [grpc_types.PubsubMessage(data=b"x" * 10)] * 5 _run_in_daemon(flow_controller.add, initial_messages, adding_initial_done) assert adding_initial_done.wait(timeout=0.1) _run_in_daemon(flow_controller.add, [large_msg], adding_large_done) # Continuously keep adding more messages after the large one. messages = [grpc_types.PubsubMessage(data=b"x" * 10)] * 10 _run_in_daemon(flow_controller.add, messages, adding_busy_done, action_pause=0.1) # At the same time, gradually keep releasing the messages - the freeed up # capacity should be consumed by the large message, not the other small messages # being added after it. _run_in_daemon(flow_controller.release, messages, releasing_busy_done, action_pause=0.1) # Sanity check - releasing should have completed by now. if not releasing_busy_done.wait(timeout=1.1): pytest.fail( "Releasing messages blocked or errored.") # pragma: NO COVER # Enough messages released, the large message should have come through in # the meantime. if not adding_large_done.wait(timeout=0.1): pytest.fail( "A thread adding a large message starved.") # pragma: NO COVER if adding_busy_done.wait(timeout=0.1): pytest.fail("Adding multiple small messages did not block." ) # pragma: NO COVER # Releasing the large message should unblock adding the remaining "busy" messages # that have not been added yet. _run_in_daemon(flow_controller.release, [large_msg], releasing_large_done) if not releasing_large_done.wait(timeout=0.1): pytest.fail( "Releasing a message blocked or errored.") # pragma: NO COVER if not adding_busy_done.wait(timeout=1.0): pytest.fail("Adding messages blocked or errored.") # pragma: NO COVER
def test_blocking_on_overflow_until_free_capacity(): settings = types.PublishFlowControl( message_limit=1, byte_limit=150, limit_exceeded_behavior=types.LimitExceededBehavior.BLOCK, ) flow_controller = FlowController(settings) msg1 = grpc_types.PubsubMessage(data=b"x" * 100) msg2 = grpc_types.PubsubMessage(data=b"y" * 100) msg3 = grpc_types.PubsubMessage(data=b"z" * 100) msg4 = grpc_types.PubsubMessage(data=b"w" * 100) # If there is a concurrency bug in FlowController, we do not want to block # the main thread running the tests, thus we delegate all add/release # operations to daemon threads and check the outcome (blocked/not blocked) # through Events. adding_1_done = threading.Event() adding_2_done = threading.Event() adding_3_done = threading.Event() adding_4_done = threading.Event() releasing_1_done = threading.Event() releasing_x_done = threading.Event() # Adding a message with free capacity should not block. _run_in_daemon(flow_controller.add, [msg1], adding_1_done) if not adding_1_done.wait(timeout=0.1): pytest.fail( # pragma: NO COVER "Adding a message with enough flow capacity blocked or errored.") # Adding messages when there is not enough capacity should block, even if # added through multiple threads. _run_in_daemon(flow_controller.add, [msg2], adding_2_done) if adding_2_done.wait(timeout=0.1): pytest.fail( "Adding a message on overflow did not block.") # pragma: NO COVER _run_in_daemon(flow_controller.add, [msg3], adding_3_done) if adding_3_done.wait(timeout=0.1): pytest.fail( "Adding a message on overflow did not block.") # pragma: NO COVER _run_in_daemon(flow_controller.add, [msg4], adding_4_done) if adding_4_done.wait(timeout=0.1): pytest.fail( "Adding a message on overflow did not block.") # pragma: NO COVER # After releasing one message, there should be room for a new message, which # should result in unblocking one of the waiting threads. _run_in_daemon(flow_controller.release, [msg1], releasing_1_done) if not releasing_1_done.wait(timeout=0.1): pytest.fail( "Releasing a message blocked or errored.") # pragma: NO COVER done_status = [ adding_2_done.wait(timeout=0.1), adding_3_done.wait(timeout=0.1), adding_4_done.wait(timeout=0.1), ] # In sum() we use the fact that True==1 and False==0, and that Event.wait() # returns False only if it times out, i.e. its internal flag has not been set. done_count = sum(done_status) assert done_count == 1, "Exactly one thread should have been unblocked." # Release another message and verify that yet another thread gets unblocked. added_msg = [msg2, msg3, msg4][done_status.index(True)] _run_in_daemon(flow_controller.release, [added_msg], releasing_x_done) if not releasing_x_done.wait(timeout=0.1): pytest.fail( "Releasing messages blocked or errored.") # pragma: NO COVER released_count = sum(( adding_2_done.wait(timeout=0.1), adding_3_done.wait(timeout=0.1), adding_4_done.wait(timeout=0.1), )) assert released_count == 2, "Exactly two threads should have been unblocked."
def test_len(): batch = create_batch(status=BatchStatus.ACCEPTING_MESSAGES) assert len(batch) == 0 batch.publish(gapic_types.PubsubMessage(data=b"foo")) assert len(batch) == 1
def make_message(ack_id, ordering_key): proto_msg = gapic_types.PubsubMessage(data=b"Q", ordering_key=ordering_key) return message.Message(proto_msg._pb, ack_id, 0, queue.Queue())
def create_message(): return gapic_types.PubsubMessage(data=b"foo", attributes={"bar": u"baz"})