class TestKeyedTopic(TestCase): def setUp(self): self.keyed_producer = Mock() self.keyed_producer_provider = Mock(return_value=self.keyed_producer) self.consumer = Mock() self.consumer_provider = Mock(return_value=self.consumer) self.topic = KeyedTopic( 'fruit', self.keyed_producer_provider, self.consumer_provider, FakeKeyCodec(), FakeValueCodec()) def test_send_message(self): msg = 'banana' self.topic.send_messages('2', msg) self.keyed_producer.send_messages.assert_called_once_with('fruit', '2?', 'banana!') self.assertEqual(1, self.keyed_producer_provider.call_count) def test_send_messages(self): messages = ['apple', 'banana', 'lemon'] self.topic.send_messages('3', *messages) self.keyed_producer.send_messages.assert_called_once_with( 'fruit', '3?', 'apple!', 'banana!', 'lemon!') self.assertEqual(1, self.keyed_producer_provider.call_count) def test_multiple_send_messages(self): messages_a = ['apple', 'banana', 'lemon'] messages_b = ['lime', 'peach', 'grapes'] self.topic.send_messages('4', *messages_a) self.topic.send_messages('5', *messages_b) self.keyed_producer.send_messages.assert_has_calls([ call('fruit', '4?', 'apple!', 'banana!', 'lemon!'), call('fruit', '5?', 'lime!', 'peach!', 'grapes!')]) self.assertEqual(1, self.keyed_producer_provider.call_count) def test_consume_a_message(self): msg = KafkaMessage('fruit', 0, 1, '3?', 'banana!') self.consumer.next.return_value = msg decoded_msg = self.topic.next() self.assertEqual(decoded_msg.value, 'banana') self.assertEqual(decoded_msg.key, '3') self.consumer_provider.assert_called_once_with('fruit') def test_consume_messages_via_iterator(self): self.consumer.next.side_effect = [ KafkaMessage('fruit', 0, n, item[0], item[1]) for n, item in enumerate([ ('5?', 'apple!'), ('6?', 'banana!'), ('7?', 'lemon!')]) ] messages = [] for message in self.topic: messages.append(message) if len(messages) >= 3: break self.assertEqual( ['apple', 'banana', 'lemon'], [msg.value for msg in messages]) self.assertEqual( ['5', '6', '7'], [msg.key for msg in messages]) self.consumer_provider.assert_called_once_with('fruit') def test_task_done(self): msg = Mock() self.topic.task_done(msg) self.consumer.task_done.assert_called_once_with(msg) def test_commit(self): self.topic.commit() self.consumer.commit.assert_called_once_with()