def test_handle_produce_response_on_nonretriable(self, fake_decode,
fake_retry):
m1 = Message('test_driver', {'val': 1}, None, None, 0, 1)
m2 = Message('test_driver', {'val': 2}, None, None, 0, 1)
response = Mock()
response.code = 500
response.error = False
response.request = Mock()
response.request._batch = [m1, m2]
response.request._id = 1
body = {
'offsets': [{}, {}],
'value_schema_id': 1,
'key_schema_id': 2,
'error_code': 50101,
'message': 'Nonretriable'
}
fake_decode.return_value = body
self.producer._handle_produce_response('test_driver', response)
calls = [
call('test_driver', m1, 'nonretriable'),
call('test_driver', m2, 'nonretriable')
]
self.client.mock_for('drop_message').assert_has_calls(calls)
def test_message_batches_from_retry_queue(self):
m1 = Message('test_driver', {'val': 1}, None, None, sys.maxsize, 1)
m2 = Message('test_driver', {'val': 2}, None, None, 0, 1)
queue = self.client.retry_queues['test_driver']
for m in [m1, m2]:
queue.put(m)
result = list(self.producer._message_batches_from_queue(queue))
self.assertEqual(len(result), 1)
self.assertEqual(result[0], [m2])
def test_message_batches_from_queue_multiple_batches(self):
self.client.flush_max_batch_size = 1
m1 = Message('test_driver', {'val': 1}, None, None, 0, 1)
m2 = Message('test_driver', {'val': 2}, None, None, 0, 1)
queue = self.client.message_queues['test_driver']
for m in [m1, m2]:
queue.put(m)
result = list(self.producer._message_batches_from_queue(queue))
self.assertEqual(len(result), 2)
self.assertEqual(result[0], [m1])
self.assertEqual(result[1], [m2])
def test_evaluate_queue_on_length(self, fake_flush, fake_flush_timer):
self.client.flush_length_threshold = 1
m1 = Message('test_driver', {'val': 1}, None, None, 0, 1)
m2 = Message('test_driver', {'val': 2}, None, None, 0, 1)
queue = self.client.message_queues['test_driver']
for m in [m1, m2]:
queue.put(m)
self.producer.evaluate_queue('test_driver',
self.client.message_queues['test_driver'])
fake_flush.assert_called_with('test_driver', FlushReason.LENGTH)
self.assertFalse(fake_flush_timer.called)
def test_send_batch_produce_request(self, fake_request):
m1 = Message('test_driver', {'val': 1}, None, None, 0, 1)
m2 = Message('test_driver', {'val': 2}, None, None, 0, 1)
batch = [m1, m2]
fake_request._id = 1
self.producer._send_batch_produce_request('test_driver', batch)
call_args = fake_request.mock_calls[0][1]
self.assertEqual(call_args[2], 10)
self.assertEqual(call_args[3], 60)
self.assertEqual(call_args[5], 'test_driver')
self.assertEqual(call_args[6], batch)
def test_send_batch_produce_request_in_shutdown(self, fake_request):
m1 = Message('test_driver', {'val': 1}, None, None, 0, 1)
m2 = Message('test_driver', {'val': 2}, None, None, 0, 1)
batch = [m1, m2]
fake_request._id = 1
self.client.in_shutdown = True
self.producer._send_batch_produce_request('test_driver', batch)
call_args = fake_request.mock_calls[0][1]
self.assertEqual(call_args[2], 0)
self.assertEqual(call_args[3], 0)
self.assertEqual(call_args[5], 'test_driver')
for idx, m in enumerate(batch):
self.assertTrue(m.true_equals(call_args[6][idx]))
def test_produce_through_evaluate_into_flush(self):
self.client.flush_length_threshold = 0
self.client.produce('test_driver', self.test_value, self.test_schema)
self.assertEqual(
self.io_loop.next_callback,
Callback(
self.producer.evaluate_queue,
('test_driver', self.client.message_queues['test_driver']),
{}))
self.io_loop.run_next_callback()
self.assertEqual(self.io_loop.next_callback.fn,
self.producer._send_batch_produce_request)
expected_message = Message('test_driver', self.test_value, None, None,
0, 1)
self.assertTrue(self.io_loop.next_callback.args[1][0].true_equals(
expected_message))
self.assertEqual(self.io_loop.next_later.seconds,
self.client.flush_time_threshold_seconds)
self.assertEqual(
self.io_loop.next_later.callback,
Callback(self.producer._flush_topic, ('test_driver', 'time'), {}))
self.io_loop.pop_callback()
self.io_loop.pop_later()
self.assertTrue(self.io_loop.finished)
def test_flush_through_send(self):
m1 = Message('test_driver', {'val': 1}, None, None, 0, 1)
queue = self.client.message_queues['test_driver']
queue.put(m1)
self.producer._flush_topic('test_driver', 'test')
self.assertEqual(self.io_loop.next_callback.fn,
self.producer._send_batch_produce_request)
self.assertTrue(self.io_loop.next_callback.args[1][0].true_equals(m1))
self.io_loop.run_next_callback()
# Response handler callback from Tornado's HTTP client, this function
# is defined inline so we can't test against it very well
self.assertEqual(self.io_loop.next_callback.fn.__name__,
'handle_response')
self.io_loop.pop_callback()
self.assertEqual(
self.io_loop.next_later.callback,
Callback(self.producer._flush_topic, ('test_driver', 'time'), {}))
self.io_loop.pop_later()
self.assertTrue(self.io_loop.finished)
def test_handle_produce_response_transport_error(self, fake_retry):
m1 = Message('test_driver', {'val': 1}, None, None, 0, 1)
m2 = Message('test_driver', {'val': 2}, None, None, 0, 1)
response = Mock()
response.code = 599
response.error = object()
response.request = Mock()
response.request._batch = [m1, m2]
response.request._id = 1
self.producer._handle_produce_response('test_driver', response)
self.client.mock_for('response_5xx').assert_called_once_with(
'test_driver', response)
fake_retry.assert_has_calls(
[call('test_driver', m1),
call('test_driver', m2)])
def test_queue_message_for_retry_when_full(self):
m1 = Message('test_driver', {'val': 1}, None, None, 0, 1)
self.client.retry_queues['test_driver'] = PriorityQueue(maxsize=1)
self.client.retry_queues['test_driver'].put('whatever')
self.producer._queue_message_for_retry('test_driver', m1)
self.client.mock_for('drop_message').assert_called_once_with(
'test_driver', m1, 'retry_queue_full')
def test_handle_produce_response_on_200(self, fake_decode,
fake_handle_success):
m1 = Message('test_driver', {'val': 1}, None, None, 0, 1)
m2 = Message('test_driver', {'val': 2}, None, None, 0, 1)
response = Mock()
response.code = 200
response.error = False
response.request = Mock()
response.request._batch = [m1, m2]
response.request._id = 1
body = {'offsets': [{}, {}], 'value_schema_id': 1, 'key_schema_id': 2}
fake_decode.return_value = body
self.producer._handle_produce_response('test_driver', response)
fake_handle_success.assert_called_once_with('test_driver', response,
body)
def test_flush_topic_sends_batches(self):
self.client.flush_max_batch_size = 1
m1 = Message('test_driver', {'val': 1}, None, None, 0, 1)
m2 = Message('test_driver', {'val': 2}, None, None, 0, 1)
queue = self.client.message_queues['test_driver']
for m in [m1, m2]:
queue.put(m)
self.producer._flush_topic('test_driver', 'testing')
self.client.mock_for('flush_topic').assert_called_with(
'test_driver', 'testing')
self.callback_mock.assert_has_calls([
call(self.producer._send_batch_produce_request, 'test_driver',
[m1]),
call(self.producer._send_batch_produce_request, 'test_driver',
[m2])
])
def test_queue_message_for_retry(self):
m1 = Message('test_driver', {'val': 1}, None, None, 0, 1)
self.producer._queue_message_for_retry('test_driver', m1)
retry_message = self.client.retry_queues['test_driver'].get()
self.assertEqual(retry_message.attempt_number, m1.attempt_number + 1)
self.assertNotEqual(retry_message.retry_after_time,
m1.retry_after_time)
self.client.mock_for('retry_message').assert_called_once_with(
'test_driver', retry_message)
def test_handle_produce_success_simple_success(self):
m1 = Message('test_driver', {'val': 1}, None, None, 0, 1)
response = Mock()
response.request = Mock()
response.request._batch = [m1]
body = {'offsets': [{}], 'value_schema_id': 1, 'key_schema_id': 2}
self.producer._handle_produce_success('test_driver', response, body)
self.client.mock_for('produce_success').assert_called_once_with(
'test_driver', [(m1, {})], [])
def test_produce_with_full_queue(self):
self.client.message_queues['test_driver'] = Queue(maxsize=1)
self.client.message_queues['test_driver'].put('whatever')
self.client.produce('test_driver', self.test_value, self.test_schema)
expected_message = Message('test_driver', self.test_value, None, None,
0, 1)
self.assertEqual(1, self.client.message_queues['test_driver'].qsize())
self.assertTrue(
self.client.mock_for('drop_message').called_with(
'test_driver', expected_message, 'primary_queue_full'))
def test_initial_produce(self):
self.client.produce('test_driver', self.test_value, self.test_schema)
expected_message = Message('test_driver', self.test_value, None, None,
0, 1)
self.assertEqual(self.client.schema_cache['test_driver']['value'],
self.test_schema)
self.assertEqual(None,
self.client.schema_cache['test_driver'].get('key'))
self.assertEqual(1, self.client.message_queues['test_driver'].qsize())
self.client.mock_for('produce').assert_called_once_with(
expected_message)
self.callback_mock.assert_called_once_with(
self.client.producer.evaluate_queue, 'test_driver',
self.client.message_queues['test_driver'])
def test_handle_produce_success_mixed_success(self, fake_retry):
m1 = Message('test_driver', {'val': 1}, None, None, 0, 1)
m2 = Message('test_driver', {'val': 2}, None, None, 0, 1)
m3 = Message('test_driver', {'val': 3}, None, None, 0, 1)
offsets = [{}, {
'error_code': 1,
'message': 'Nonretriable'
}, {
'error_code': 2,
'message': 'Retriable'
}]
response = Mock()
response.request = Mock()
response.request._batch = [m1, m2, m3]
response.request._id = 1
body = {'offsets': offsets, 'value_schema_id': 1, 'key_schema_id': 2}
self.producer._handle_produce_success('test_driver', response, body)
self.client.mock_for('produce_success').assert_called_once_with(
'test_driver', [(m1, {})], [(m2, offsets[1]), (m3, offsets[2])])
self.client.mock_for('drop_message').assert_called_once_with(
'test_driver', m2, 'nonretriable')
fake_retry.assert_called_once_with('test_driver', m3)
def make_response(self, response_attrs=None, request_attrs=None):
response = Mock()
request = Mock()
response.code = 200
response.error = None
response.body = json.dumps({'offsets': [{}]})
request._id = 1
request._topic = 'test_driver'
request._batch = [
Message('test_driver', self.test_value, None, None, 0, 1)
]
for k, v in (response_attrs or {}).items():
if k == 'body':
v = json.dumps(v)
setattr(response, k, v)
for k, v in (request_attrs or {}).items():
setattr(request, k, v)
response.request = request
return response
def test_queue_message_for_retry_nonretriable(self):
m1 = Message('test_driver', {'val': 1}, None, None, 0, 10)
self.producer._queue_message_for_retry('test_driver', m1)
self.client.mock_for('drop_message').assert_called_once_with(
'test_driver', m1, 'max_retries_exceeded')