Using Kafka MockConsumer

1. Overview

In this tutorial, we'll explore the MockConsumer, one of Kafka‘s Consumer implementations.

First, we'll discuss what are the main things to be considered when testing a Kafka Consumer. Then, we'll see how we can use MockConsumer to implement tests.

2. Testing a Kafka Consumer

Consuming data from Kafka consists of two main steps. Firstly, we have to subscribe to topics or assign topic partitions manually. Secondly, we poll batches of records using the poll method.

The polling is usually done in an infinite loop. That's because we typically want to consume data continuously.

For example, let's consider the simple consuming logic consisting of just the subscription and the polling loop:

void consume() {
    try {
        consumer.subscribe(Arrays.asList("foo", "bar"));
        while (true) {
            ConsumerRecords<String, String> records = consumer.poll(Duration.ofSeconds(1));
            records.forEach(record -> processRecord(record));
        }
    } catch (WakeupException ex) {
        // ignore for shutdown
    } catch (RuntimeException ex) {
        // exception handling
    } finally {
        consumer.close();
    }
}

Looking at the code above, we can see that there are a few things we can test:

• the subscription
• the polling loop
• the exception handling
• if the Consumer was closed correctly

We have multiple options to test the consuming logic.

We can use an in-memory Kafka instance. But, this approach has some disadvantages. In general, an in-memory Kafka instance makes tests very heavy and slow. Moreover, setting it up is not a simple task and can lead to unstable tests.

Alternatively, we can use a mocking framework to mock the Consumer. Although using this approach makes tests lightweight, setting it up can be somewhat tricky.

The final option, and perhaps the best, is to use the MockConsumer, which is a Consumer implementation meant for testing. Not only does it help us to build lightweight tests, but it's also easy to set up.

Let's have a look at the features it provides.

3. Using MockConsumer

MockConsumer implements the Consumer interface that the kafka-clients library provides. Therefore, it mocks the entire behavior of a real Consumer without us needing to write a lot of code.

Let's look at some usage examples of the MockConsumer. In particular, we'll take a few common scenarios that we may come across while testing a consumer application, and implement them using the MockConsumer.

For our example, let's consider an application that consumes country population updates from a Kafka topic. The updates contain only the name of the country and its current population:

class CountryPopulation {

    private String country;
    private Integer population;

    // standard constructor, getters and setters
}

Our consumer just polls for updates using a Kafka Consumer instance, processes them, and at the end, commits the offset using the commitSync method:

public class CountryPopulationConsumer {

    private Consumer<String, Integer> consumer;
    private java.util.function.Consumer<Throwable> exceptionConsumer;
    private java.util.function.Consumer<CountryPopulation> countryPopulationConsumer;

    // standard constructor

    void startBySubscribing(String topic) {
        consume(() -> consumer.subscribe(Collections.singleton(topic)));
    }

    void startByAssigning(String topic, int partition) {
        consume(() -> consumer.assign(Collections.singleton(new TopicPartition(topic, partition))));
    }

    private void consume(Runnable beforePollingTask) {
        try {
            beforePollingTask.run();
            while (true) {
                ConsumerRecords<String, Integer> records = consumer.poll(Duration.ofMillis(1000));
                StreamSupport.stream(records.spliterator(), false)
                    .map(record -> new CountryPopulation(record.key(), record.value()))
                    .forEach(countryPopulationConsumer);
                consumer.commitSync();
            }
        } catch (WakeupException e) {
            System.out.println("Shutting down...");
        } catch (RuntimeException ex) {
            exceptionConsumer.accept(ex);
        } finally {
            consumer.close();
        }
    }

    public void stop() {
        consumer.wakeup();
    }
}

3.1. Creating a MockConsumer Instance

Next, let's see how we can create an instance of MockConsumer:

@BeforeEach
void setUp() {
    consumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
    updates = new ArrayList<>();
    countryPopulationConsumer = new CountryPopulationConsumer(consumer, 
      ex -> this.pollException = ex, updates::add);
}

Basically, all we need to provide is the offset reset strategy.

Note that we use updates to collect the records countryPopulationConsumer will receive. This will help us to assert the expected results.

In the same way, we use pollException to collect and assert the exceptions.

For all the test cases, we'll use the above set up method. Now, let's look at a few test cases for the consumer application.

3.2. Assigning Topic Partitions

To begin, let's create a test for the startByAssigning method:

@Test
void whenStartingByAssigningTopicPartition_thenExpectUpdatesAreConsumedCorrectly() {
    // GIVEN
    consumer.schedulePollTask(() -> consumer.addRecord(record(TOPIC, PARTITION, "Romania", 19_410_000)));
    consumer.schedulePollTask(() -> countryPopulationConsumer.stop());

    HashMap<TopicPartition, Long> startOffsets = new HashMap<>();
    TopicPartition tp = new TopicPartition(TOPIC, PARTITION);
    startOffsets.put(tp, 0L);
    consumer.updateBeginningOffsets(startOffsets);

    // WHEN
    countryPopulationConsumer.startByAssigning(TOPIC, PARTITION);

    // THEN
    assertThat(updates).hasSize(1);
    assertThat(consumer.closed()).isTrue();
}

At first, we set up the MockConsumer. We start by adding a record to the consumer using the addRecord method.

The first thing to remember is that we cannot add records before assigning or subscribing to a topic. That is why we schedule a poll task using the schedulePollTask method. The task we schedule will run on the first poll before the records are fetched. Thus, the addition of the record will happen after the assignment takes place.

Equally important is that we cannot add to the MockConsumer records that do not belong to the topic and partition assigned to it.

Then, to make sure the consumer does not run indefinitely, we configure it to shut down at the second poll.

Additionally, we must set the beginning offsets. We do this using the updateBeginningOffsets method.

In the end, we check if we consumed the update correctly, and the consumer is closed.

3.3. Subscribing to Topics

Now, let's create a test for our startBySubscribing method:

@Test
void whenStartingBySubscribingToTopic_thenExpectUpdatesAreConsumedCorrectly() {
    // GIVEN
    consumer.schedulePollTask(() -> {
        consumer.rebalance(Collections.singletonList(new TopicPartition(TOPIC, 0)));
        consumer.addRecord(record("Romania", 1000, TOPIC, 0));
    });
    consumer.schedulePollTask(() -> countryPopulationConsumer.stop());

    HashMap<TopicPartition, Long> startOffsets = new HashMap<>();
    TopicPartition tp = new TopicPartition(TOPIC, 0);
    startOffsets.put(tp, 0L);
    consumer.updateBeginningOffsets(startOffsets);

    // WHEN
    countryPopulationConsumer.startBySubscribing(TOPIC);

    // THEN
    assertThat(updates).hasSize(1);
    assertThat(consumer.closed()).isTrue();
}

In this case, the first thing to do before adding a record is a rebalance. We do this by calling the rebalance method, which simulates a rebalance.

The rest is the same as the startByAssigning test case.

3.4. Controlling the Polling Loop

We can control the polling loop in multiple ways.

The first option is to schedule a poll task as we did in the tests above. We do this via schedulePollTask, which takes a Runnable as a parameter. Each task we schedule will run when we call the poll method.

The second option we have is to call the wakeup method. Usually, this is how we interrupt a long poll call. Actually, this is how we implemented the stop method in CountryPopulationConsumer.

Lastly, we can set an exception to be thrown using the setPollException method:

@Test
void whenStartingBySubscribingToTopicAndExceptionOccurs_thenExpectExceptionIsHandledCorrectly() {
    // GIVEN
    consumer.schedulePollTask(() -> consumer.setPollException(new KafkaException("poll exception")));
    consumer.schedulePollTask(() -> countryPopulationConsumer.stop());

    HashMap<TopicPartition, Long> startOffsets = new HashMap<>();
    TopicPartition tp = new TopicPartition(TOPIC, 0);
    startOffsets.put(tp, 0L);
    consumer.updateBeginningOffsets(startOffsets);

    // WHEN
    countryPopulationConsumer.startBySubscribing(TOPIC);

    // THEN
    assertThat(pollException).isInstanceOf(KafkaException.class).hasMessage("poll exception");
    assertThat(consumer.closed()).isTrue();
}

3.5. Mocking End Offsets and Partitions Info

If our consuming logic is based on end offsets or partition information, we can also mock these using MockConsumer.

When we want to mock the end offset, we can use the addEndOffsets and updateEndOffsets methods.

And, in case we want to mock partition information, we can use the updatePartitions method.

4. Conclusion

In this article, we've explored how to use MockConsumer to test a Kafka consumer application.

First, we've looked at an example of consumer logic and which are the essential parts to test. Then, we tested a simple Kafka consumer application using the MockConsumer.

Along the way, we looked at the features of the MockConsumer and how to use it.

As always, all these code samples are available over on GitHub.