def test_write_to_pubsub_many_chunks(self):
number_of_elements = 83
chunk_size = 11
mock_pubsub = mock.Mock()
topic_path = "project/fakeproj/topics/faketopic"
data_list = [
'data {}'.format(i).encode("utf-8")
for i in range(number_of_elements)
]
utils.write_to_pubsub(mock_pubsub,
topic_path,
data_list,
chunk_size=chunk_size)
call_list = []
for start in range(0, number_of_elements, chunk_size):
# Publish a batch of messages
call_list += [
mock.call(topic_path, data)
for data in data_list[start:start + chunk_size]
]
# Wait for those messages to be received
call_list += [
mock.call().result()
for _ in data_list[start:start + chunk_size]
]
mock_pubsub.publish.assert_has_calls(call_list)
def test_write_to_pubsub(self):
mock_pubsub = mock.Mock()
topic_path = "project/fakeproj/topics/faketopic"
data = b'data'
utils.write_to_pubsub(mock_pubsub, topic_path, [data])
mock_pubsub.publish.assert_has_calls(
[mock.call(topic_path, data),
mock.call().result()])
def test_write_to_pubsub_with_attributes(self):
mock_pubsub = mock.Mock()
topic_path = "project/fakeproj/topics/faketopic"
data = b'data'
attributes = {'key': 'value'}
message = PubsubMessage(data, attributes)
utils.write_to_pubsub(
mock_pubsub, topic_path, [message], with_attributes=True)
mock_pubsub.publish.assert_has_calls(
[mock.call(topic_path, data, **attributes),
mock.call().result()])
def test_write_to_pubsub_delay(self):
number_of_elements = 2
chunk_size = 1
mock_pubsub = mock.Mock()
topic_path = "project/fakeproj/topics/faketopic"
data = b'data'
with mock.patch('apache_beam.io.gcp.tests.utils.time') as mock_time:
utils.write_to_pubsub(mock_pubsub,
topic_path, [data] * number_of_elements,
chunk_size=chunk_size,
delay_between_chunks=123)
mock_time.sleep.assert_called_with(123)
mock_pubsub.publish.assert_has_calls(
[mock.call(topic_path, data),
mock.call().result()] * number_of_elements)