Unit Test Kafka Consumer Written Using Spring Kafka

Question

I have created a Kakfa consumer application (using spring kafka) and it seems to be working fine. Now, I am trying to add unit test cases for the same. My consumer is batch acknowledgement consumer and it is started inside a container. The details of the consumer can be found in a below stack overflow post (which is asked by me)

Spring-Kafka Concurrency Property

I am able to write a test case for my consumer by doing some research, but here I had to replicate the creation of container and few other things in order to make it work. I was wondering, if there is any other way for test cases to use the same container which is started during application startup (may by pointing test context) and the consumer which is started during startup to use the embedded kafka instance directly.

The solution I figured out is as below, but not sure this is the right approach or not.

@RunWith(SpringRunner.class)
@SpringBootTest(classes = Launcher.class)
public class BatchConsumerTest {

    @ClassRule
    public static EmbeddedKafkaRule embeddedKafka = new EmbeddedKafkaRule(1, true, "topic1");
    @Autowired
    private KafkaListenerEndpointRegistry registry;

    @Autowired
    private KafkaTemplate<String, String> template;

    @MockBean
    private RestService restService;

    @Before
    public void setup() {
        Mockito.when(restService.invokeService("")).thenReturn("");
    }

    @Test
    public void test() throws Exception {
        ConcurrentMessageListenerContainer<?, ?> container = null;
        for (MessageListenerContainer con : registry.getAllListenerContainers()) {
            container = (ConcurrentMessageListenerContainer<?, ?>) con;
            container.stop();
        }
        Map<String, Object> consumerProps = KafkaTestUtils.consumerProps("group1", "false",
                embeddedKafka.getEmbeddedKafka());
        DefaultKafkaConsumerFactory<String, String> factory = new DefaultKafkaConsumerFactory<String, String>(
                consumerProps);
        ContainerProperties containerProps = new ContainerProperties("ibo-grocerybag-subscription");
        containerProps.setAckMode(AckMode.MANUAL);
        ConcurrentMessageListenerContainer<String, String> messageListContainer = new ConcurrentMessageListenerContainer<String, String>(
                factory, containerProps);
        BatchAcknowledgingConsumerAwareMessageListener<String, String> listener = new BatchConsumer();
        CountDownLatch latch = new CountDownLatch(1);
        // messageListContainer.setupMessageListener(listener);
        messageListContainer.setupMessageListener(new BatchAcknowledgingConsumerAwareMessageListener<String, String>() {

            @Override
            public void onMessage(List<ConsumerRecord<String, String>> data, Acknowledgment acknowledgment,
                    Consumer<?, ?> consumer) {
                System.out.println("*******Data : " + data.get(0).value());
                listener.onMessage(data, acknowledgment, consumer);
                latch.countDown();
            }
        });
        messageListContainer.start();
        ContainerTestUtils.waitForAssignment(messageListContainer,
                embeddedKafka.getEmbeddedKafka().getPartitionsPerTopic());
        template.send("topic1", "Hello");
        latch.await(10000, TimeUnit.MILLISECONDS);
        assertThat(latch.getCount()).isEqualTo(0);
    }

    @After
    public void destroy() {
        embeddedKafka.getEmbeddedKafka().destroy();
    }