Python Kafka mocking Kafka Consumer 的返回类型

Question

I want to test a script when I use kafka-python package. I want to test type of return object for function

def _get_kafka_consumer() -> KafkaConsumer:
    consumer = KafkaConsumer(bootstrap_servers=_KAFKA_BOOTSTRAP_SERVICE,
                             auto_offset_reset='earliest')
    consumer.subscribe([_KAFKA_TOPIC_INPUT])
    return consumer

My test class looks like

class TestVoiceToText(unittest.TestCase):

    def test_get_kafka_consumer_output_type(self):
        result = _get_kafka_consumer()
        self.assertIsInstance(result, KafkaConsumer)

and of course it does not pass because there is no Kafka Cluster running so KafkaConsumer cannot be created. How can I mock that returning of KafkaConsumer(...) is of type KafkaConsumer without actual need of calling the actual constructor?

Answer 1

I managed to solve the problem by using patch function from unittest.mock package:

    def test_get_kafka_consumer_output_type(self):
        with patch('voice_to_text.kafka_connector.KafkaConsumer') as kafka_consumer_class_mock:
            kafka_consumer_instance = _get_kafka_consumer()
            kafka_consumer_class_mock_instance = kafka_consumer_class_mock.return_value

            kafka_consumer_class_mock.assert_called_once()
            self.assertEquals(kafka_consumer_class_mock_instance, kafka_consumer_instance)

It only checks if a result of _get_kafka_consumer() is actually an object returned by calling a function KafkaConsumer() . We don't care what actually that function is doing.

Answer 2

You would want the concept of spec in mocking. You can use autospec here I believe with success. Speccing will limit your mock to the original API of the object you're mocking, and allow it to pass isinstance tests as well. Here's the documentation on Autospeccing: https://docs.python.org/3.3/library/unittest.mock.html#autospeccing

Here's how I'd do your test:

 @mock.patch('voice_to_text.kafka_connector.KafkaConsumer', autospec=True)
 def test_get_kafka_consumer_output_type(self, kafka_consumer_mock):
        kafka_consumer_instance = _get_kafka_consumer()
        
        kafka_consumer_mock.assert_called_once_with(
            bootstrap_servers=_KAFKA_BOOTSTRAP_SERVICE,
            auto_offset_reset='earliest')
        kafka_consumer_mock.return_value.subscribe.assert_called_once_with(
            [_KAFKA_TOPIC_INPUT])
        kafka_consumer_mock.assert_called_once()

        self.assertEquals(kafka_consumer_mock.return_value, kafka_consumer_instance)
        self.assertIsInstance(kafka_consumer_instance, KafkaConsumer)