def test_blocking__commit():
batch = create_batch()
futures = (
batch.publish({"data": b"This is my message."}),
batch.publish({"data": b"This is another message."}),
)
# Set up the underlying API publish method to return a PublishResponse.
publish_response = types.PublishResponse(message_ids=["a", "b"])
patch = mock.patch.object(type(batch.client.api),
"publish",
return_value=publish_response)
with patch as publish:
batch._commit()
# Establish that the underlying API call was made with expected
# arguments.
publish.assert_called_once_with(
"topic_name",
[
types.PubsubMessage(data=b"This is my message."),
types.PubsubMessage(data=b"This is another message."),
],
)
# Establish that all of the futures are done, and that they have the
# expected values.
assert futures[0].done()
assert futures[0].result() == "a"
assert futures[1].done()
assert futures[1].result() == "b"
def test_blocking_commit():
batch = create_batch()
futures = (
batch.publish({'data': b'This is my message.'}),
batch.publish({'data': b'This is another message.'}),
)
# Set up the underlying API publish method to return a PublishResponse.
with mock.patch.object(type(batch.client.api), 'publish') as publish:
publish.return_value = types.PublishResponse(message_ids=['a', 'b'])
# Actually commit the batch.
batch._commit()
# Establish that the underlying API call was made with expected
# arguments.
publish.assert_called_once_with('topic_name', [
types.PubsubMessage(data=b'This is my message.'),
types.PubsubMessage(data=b'This is another message.'),
])
# Establish that all of the futures are done, and that they have the
# expected values.
assert all([f.done() for f in futures])
assert futures[0].result() == 'a'
assert futures[1].result() == 'b'
def test_blocking_commit_wrong_messageid_length():
batch = create_batch()
futures = (
batch.publish({'data': b'blah blah blah'}),
batch.publish({'data': b'blah blah blah blah'}),
)
# Set up a PublishResponse that only returns one message ID.
with mock.patch.object(type(batch.client.api), 'publish') as publish:
publish.return_value = types.PublishResponse(message_ids=['a'])
batch._commit()
for future in futures:
assert future.done()
assert isinstance(future.exception(), exceptions.PublishError)
def test_batch_done_callback_called_on_publish_response_invalid():
batch_done_callback_tracker = BatchDoneCallbackTracker()
batch = create_batch(batch_done_callback=batch_done_callback_tracker)
# Ensure messages exist.
message = types.PubsubMessage(data=b"foobarbaz")
batch.publish(message)
# No message ids returned in successful publish response -> invalid.
publish_response = types.PublishResponse(message_ids=[])
with mock.patch.object(type(batch.client.api),
"publish",
return_value=publish_response):
batch._commit()
assert batch_done_callback_tracker.called
assert not batch_done_callback_tracker.success
def test_batch_done_callback_called_on_success():
batch_done_callback_tracker = BatchDoneCallbackTracker()
batch = create_batch(batch_done_callback=batch_done_callback_tracker)
# Ensure messages exist.
message = types.PubsubMessage(data=b"foobarbaz")
batch.publish(message)
# One response for one published message.
publish_response = types.PublishResponse(message_ids=["a"])
with mock.patch.object(type(batch.client.api),
"publish",
return_value=publish_response):
batch._commit()
assert batch_done_callback_tracker.called
assert batch_done_callback_tracker.success
def test_blocking__commit_wrong_messageid_length():
batch = create_batch()
futures = (
batch.publish({"data": b"blah blah blah"}),
batch.publish({"data": b"blah blah blah blah"}),
)
# Set up a PublishResponse that only returns one message ID.
publish_response = types.PublishResponse(message_ids=["a"])
patch = mock.patch.object(type(batch.client.api),
"publish",
return_value=publish_response)
with patch:
batch._commit()
for future in futures:
assert future.done()
assert isinstance(future.exception(), exceptions.PublishError)
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 = types.PubsubMessage(data=b"foobarbaz")
batch.publish(message)
# One response for one published message.
publish_response = 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 api_publish_delay(_, messages):
api_publish_called.set()
time.sleep(1.0)
message_ids = [str(i) for i in range(len(messages))]
return types.PublishResponse(message_ids=message_ids)
