def setUp(self): self.tup_dict = {'id': 14, 'comp': 'some_spout', 'stream': 'default', 'task': 'some_spout', 'tuple': [1, 2, 3]} self.tup = Tuple(self.tup_dict['id'], self.tup_dict['comp'], self.tup_dict['stream'], self.tup_dict['task'], self.tup_dict['tuple'],) self.spout = ReliableSpout(input_stream=BytesIO(), output_stream=BytesIO()) self.spout.initialize({}, {}) self.spout.logger = logging.getLogger(__name__)
class ReliableSpoutTests(unittest.TestCase): def setUp(self): self.tup_dict = {'id': 14, 'comp': 'some_spout', 'stream': 'default', 'task': 'some_spout', 'tuple': [1, 2, 3]} self.tup = Tuple(self.tup_dict['id'], self.tup_dict['comp'], self.tup_dict['stream'], self.tup_dict['task'], self.tup_dict['tuple'],) self.spout = ReliableSpout(input_stream=BytesIO(), output_stream=BytesIO()) self.spout.initialize({}, {}) self.spout.logger = logging.getLogger(__name__) @patch.object(ReliableSpout, 'send_message', autospec=True) def test_emit_unreliable(self, send_message_mock): # An unreliable emit is not allowed with ReliableSpout with self.assertRaises(ValueError): self.spout.emit([1, 2, 3], need_task_ids=False) @patch.object(ReliableSpout, 'send_message', autospec=True) def test_emit_reliable(self, send_message_mock): # Reliable emit self.spout.emit([1, 2, 3], tup_id='foo', need_task_ids=False) send_message_mock.assert_called_with(self.spout, {'command': 'emit', 'tuple': [1, 2, 3], 'need_task_ids': False, 'id': 'foo'}) self.assertEqual(self.spout.unacked_tuples.get('foo'), ([1, 2, 3], None, None, False)) @patch.object(ReliableSpout, 'send_message', autospec=True) def test_emit_reliable_direct(self, send_message_mock): # Reliable emit as direct task self.spout.emit([1, 2, 3], tup_id='foo', direct_task='other_spout') send_message_mock.assert_called_with(self.spout, {'command': 'emit', 'tuple': [1, 2, 3], 'task': 'other_spout', 'id': 'foo', 'need_task_ids': False}) self.assertEqual(self.spout.unacked_tuples.get('foo'), ([1, 2, 3], None, 'other_spout', False)) def test_ack(self): self.spout.failed_tuples['foo'] = 2 self.spout.unacked_tuples['foo'] = ([1, 2, 3], None, None, True) # Make sure ack cleans up failed_tuples and unacked_tuples self.spout.ack('foo') self.assertNotIn('foo', self.spout.failed_tuples) self.assertNotIn('foo', self.spout.unacked_tuples) @patch.object(ReliableSpout, 'emit', autospec=True) def test_fail_below_limit(self, emit_mock): # Make sure fail increments failed_tuples and calls emit self.spout.max_fails = 3 self.spout.failed_tuples['foo'] = 2 self.spout.unacked_tuples['foo'] = ([1, 2, 3], None, None, True) self.spout.fail('foo') self.assertEqual(self.spout.failed_tuples['foo'], 3) emit_mock.assert_called_with(self.spout, [1, 2, 3], direct_task=None, need_task_ids=True, stream=None, tup_id='foo') @patch.object(ReliableSpout, 'ack', autospec=True) @patch.object(ReliableSpout, 'emit', autospec=True) def test_fail_above_limit(self, emit_mock, ack_mock): # Make sure fail increments failed_tuples self.spout.max_fails = 3 self.spout.failed_tuples['foo'] = 3 emit_args = ([1, 2, 3], None, None, True) self.spout.unacked_tuples['foo'] = emit_args self.spout.fail('foo') self.assertEqual(emit_mock.call_count, 0) ack_mock.assert_called_with(self.spout, 'foo')