def test_routing_and_task_success(self):
# Register two types of tasks.
foo_task_processor = mock.Mock()
bar_task_processor = mock.Mock()
self._worker.register(test_task_pb2.FooTaskArgs, foo_task_processor)
self._worker.register(test_task_pb2.BarTaskArgs, bar_task_processor)
# Create pubsub messages.
foo_task_args = test_task_pb2.FooTaskArgs(widget='Water is made of water.')
foo_message = _make_mock_pubsub_message(foo_task_args)
bar_task_args = test_task_pb2.BarTaskArgs(best_number=42)
bar_message = _make_mock_pubsub_message(bar_task_args)
# Subscribe, and process the messages.
subscribe_retval = 'grapes of testing'
self._mock_subscribe([foo_message, bar_message], retval=subscribe_retval)
self.assertIs(subscribe_retval, self._worker.subscribe('kumquat'))
self._client.subscribe.assert_called_once_with('kumquat', mock.ANY)
# Test that the messages were routed to the correct tasks, and that both
# messages were ACKed.
foo_task_processor.assert_called_once_with(foo_task_args)
foo_message.ack.assert_called_once_with()
bar_task_processor.assert_called_once_with(bar_task_args)
bar_message.ack.assert_called_once_with()
def test_callback_error(self):
foo_task_processor = mock.Mock(side_effect=RuntimeError('foo error'))
self._worker.register(test_task_pb2.FooTaskArgs, foo_task_processor)
foo_message = _make_mock_pubsub_message(test_task_pb2.FooTaskArgs())
self._mock_subscribe([foo_message])
self._worker.subscribe('kumquat')
foo_message.nack.assert_called_once_with()
def test_task_to_string_error(self):
self._task_to_string.side_effect = RuntimeError(
'Unable to convert proto to string.')
self._worker.register(test_task_pb2.FooTaskArgs, mock.Mock())
foo_message = _make_mock_pubsub_message(test_task_pb2.FooTaskArgs())
self._mock_subscribe([foo_message])
self._worker.subscribe('kumquat')
foo_message.nack.assert_called_once_with()
def test_publish_is_called_by(self, request_by_args):
args = test_task_pb2.FooTaskArgs()
task = task_pb2.Task()
task.args.Pack(args)
publish_future = mock.Mock()
self._client.publish.return_value = publish_future
if request_by_args:
returned_future = self._requestor.request('kumquat', args)
else:
returned_future = self._requestor.request_task('kumquat', task)
self.assertIs(publish_future, returned_future)
self._client.publish.assert_called_once_with('kumquat',
task.SerializeToString())
def test_retry_failure_then_succeed(self):
# Fail once, then succeed.
foo_task_processor = mock.Mock(
side_effect=[RuntimeError('retry this'), None])
self._worker.register(test_task_pb2.FooTaskArgs, foo_task_processor)
subscriber_future = self._worker.subscribe(_SUBSCRIPTION_NAME)
# Request that foo be worked on.
foo_task_args = test_task_pb2.FooTaskArgs()
foo_task_args.widget = 'widget maker maker'
request_future = self._requestor.request(_TOPIC_NAME, foo_task_args)
request_future.result(timeout=10)
# Wait up to a minute for the messages to go through, then stop waiting for
# more.
for _ in range(60):
time.sleep(1)
if foo_task_processor.call_count >= 2:
break
subscriber_future.cancel()
self.assertEqual(2, foo_task_processor.call_count)
foo_task_processor.assert_has_calls([mock.call(foo_task_args)] * 2)
class WorkerTest(parameterized.TestCase):
def setUp(self):
super().setUp()
self._client = mock.create_autospec(client.Client)
self._task_to_string = mock.Mock(side_effect=text_format.MessageToString)
self._worker = worker.Worker(
pubsub_subscriber_client=self._client,
task_to_string=self._task_to_string,
)
def test_routing_and_task_success(self):
# Register two types of tasks.
foo_task_processor = mock.Mock()
bar_task_processor = mock.Mock()
self._worker.register(test_task_pb2.FooTaskArgs, foo_task_processor)
self._worker.register(test_task_pb2.BarTaskArgs, bar_task_processor)
# Create pubsub messages.
foo_task_args = test_task_pb2.FooTaskArgs(widget='Water is made of water.')
foo_message = _make_mock_pubsub_message(foo_task_args)
bar_task_args = test_task_pb2.BarTaskArgs(best_number=42)
bar_message = _make_mock_pubsub_message(bar_task_args)
# Subscribe, and process the messages.
subscribe_retval = 'grapes of testing'
self._mock_subscribe([foo_message, bar_message], retval=subscribe_retval)
self.assertIs(subscribe_retval, self._worker.subscribe('kumquat'))
self._client.subscribe.assert_called_once_with('kumquat', mock.ANY)
# Test that the messages were routed to the correct tasks, and that both
# messages were ACKed.
foo_task_processor.assert_called_once_with(foo_task_args)
foo_message.ack.assert_called_once_with()
bar_task_processor.assert_called_once_with(bar_task_args)
bar_message.ack.assert_called_once_with()
def test_register_after_subscribe_error(self):
self._worker.subscribe('kumquat')
with self.assertRaisesRegex(RuntimeError,
'register.*after a subscriber is started'):
self._worker.register(test_task_pb2.FooTaskArgs, mock.Mock())
@parameterized.named_parameters(
('invalid_task_proto', b'this is probably not a valid binary proto'),
('unregistered_type', _make_task_bytes(test_task_pb2.FooTaskArgs())))
def test_unhandled_message(self, message_data):
self._worker.register(test_task_pb2.BarTaskArgs, mock.Mock())
msg = mock.create_autospec(message.Message)
msg.data = message_data
msg.message_id = str(uuid.uuid4())
self._mock_subscribe([msg])
with warnings.catch_warnings():
# Ignore a warning from trying to parse an invalid binary proto.
warnings.filterwarnings('ignore', 'Unexpected end-group tag:')
self._worker.subscribe('kumquat')
msg.nack.assert_called_once_with()
def test_task_to_string_error(self):
self._task_to_string.side_effect = RuntimeError(
'Unable to convert proto to string.')
self._worker.register(test_task_pb2.FooTaskArgs, mock.Mock())
foo_message = _make_mock_pubsub_message(test_task_pb2.FooTaskArgs())
self._mock_subscribe([foo_message])
self._worker.subscribe('kumquat')
foo_message.nack.assert_called_once_with()
def test_callback_error(self):
foo_task_processor = mock.Mock(side_effect=RuntimeError('foo error'))
self._worker.register(test_task_pb2.FooTaskArgs, foo_task_processor)
foo_message = _make_mock_pubsub_message(test_task_pb2.FooTaskArgs())
self._mock_subscribe([foo_message])
self._worker.subscribe('kumquat')
foo_message.nack.assert_called_once_with()
def _mock_subscribe(self, messages, retval=None):
"""Set up self._client.subscribe() to mock messages being published.
Note: This processes messages synchronously in the calling thread, unlike a
real pubsub subscriber.
Args:
messages: Iterable of pubsub messages to treat as if they were received
from pubsub.
retval: Value to return from subscribe(). Normally that function would
return a Future, but that isn't needed for this synchronous
implementation.
"""
def side_effect(subscription, callback):
del subscription # Unused.
for msg in messages:
callback(msg)
return retval
self._client.subscribe.side_effect = side_effect