如何用Mockito对Kafka RecordIntercept接口的intercept及afterRecord方法做单元测试

问题背景 #

我有一个实现了RecordInterceptor<String, Message<Object>>接口的KafkaRecordInterceptor类，核心代码如下：

类定义 #

@Component
public class KafkaRecordInterceptor implements RecordInterceptor<String, Message<Object>> {
    public static final ThreadLocal<MessageContext> messageContextThreadLocal = new ThreadLocal<>();
    private final List<MessageContextListener> messageContextListener;
    private final MetricsUtils metricsUtils;
    private final SystemClock systemClock;
    // 省略其他代码
}

intercept 方法实现 #

@Override
public ConsumerRecord intercept(@NonNull ConsumerRecord consumerRecord, @NonNull Consumer consumer) {

    MessageContext messageContext = new MessageContext();
    messageContext.setStartClock(systemClock.now());
    messageContextThreadLocal.set(messageContext);

    Optional<MessagePriority> messagePriorityOptional =
        Objects.requireNonNull(metricsUtils.getMessagePriority(consumerRecord));

    messageContext.setCorrelationId(KafkaUtils.extractCorrelationId(consumerRecord));
    messageContext
        .setOperationType(((GenericMessage<?>) consumerRecord.value()).getPayload().getClass().getSimpleName());
    messageContext.setPriority(messagePriorityOptional.map(MessagePriority::name).orElse(""));

    messageContextListener.forEach(listener -> listener.onStart(messageContext));

    return Objects.requireNonNull(consumerRecord);
}

afterRecord 方法实现 #

@Override
public void afterRecord(ConsumerRecord record, Consumer consumer) {
    MessageContext messageContext = messageContextThreadLocal.get();
    messageContextListener.forEach(listener -> listener.onClear(messageContext));
    messageContextThreadLocal.remove();
}

请问对该类的intercept方法和afterRecord方法进行单元测试的最佳方式是什么？

单元测试方案 #

推荐使用JUnit 5 + Mockito编写单元测试，核心思路是mock所有外部依赖，验证方法行为是否符合预期，同时处理ThreadLocal的线程隔离问题。

一、测试前置准备 #

初始化测试类，注入被测试对象并mock所有依赖，每次测试前清理ThreadLocal避免数据污染：

@ExtendWith(MockitoExtension.class)
class KafkaRecordInterceptorTest {

    @Mock
    private MetricsUtils metricsUtils;

    @Mock
    private SystemClock systemClock;

    @Mock
    private MessageContextListener listener1;

    @Mock
    private MessageContextListener listener2;

    private KafkaRecordInterceptor interceptor;

    @BeforeEach
    void setUp() {
        List<MessageContextListener> listeners = Arrays.asList(listener1, listener2);
        interceptor = new KafkaRecordInterceptor(listeners, metricsUtils, systemClock);
        KafkaRecordInterceptor.messageContextThreadLocal.remove();
    }
}

二、测试 intercept 方法 #

需要验证以下核心行为：

• ThreadLocal是否正确设置MessageContext
• MessageContext的属性赋值是否符合预期
• 依赖方法（metricsUtils.getMessagePriority、KafkaUtils.extractCorrelationId）是否被正确调用
• 所有MessageContextListener的onStart方法是否触发
• 返回值为原ConsumerRecord

示例测试代码：

@Test
void intercept_shouldSetUpMessageContextAndNotifyListeners() {
    // 1. 准备测试数据
    String correlationId = "test-correlation-id";
    long mockTime = 1620000000000L;
    TestPayload payload = new TestPayload();
    GenericMessage<TestPayload> message = new GenericMessage<>(payload);
    ConsumerRecord<String, Message<Object>> consumerRecord = new ConsumerRecord<>(
        "test-topic", 0, 0, "key", message
    );
    Consumer<String, Message<Object>> mockConsumer = mock(Consumer.class);

    // 2. mock依赖行为
    when(systemClock.now()).thenReturn(mockTime);
    when(metricsUtils.getMessagePriority(consumerRecord))
        .thenReturn(Optional.of(MessagePriority.HIGH));
    
    // mock静态方法（若用Mockito需借助MockedStatic）
    try (MockedStatic<KafkaUtils> mockedKafkaUtils = Mockito.mockStatic(KafkaUtils.class)) {
        mockedKafkaUtils.when(() -> KafkaUtils.extractCorrelationId(consumerRecord))
            .thenReturn(correlationId);

        // 3. 执行测试方法
        ConsumerRecord<String, Message<Object>> result = interceptor.intercept(consumerRecord, mockConsumer);

        // 4. 验证结果
        assertSame(consumerRecord, result);
        MessageContext context = KafkaRecordInterceptor.messageContextThreadLocal.get();
        assertNotNull(context);
        assertEquals(mockTime, context.getStartClock());
        assertEquals(correlationId, context.getCorrelationId());
        assertEquals("TestPayload", context.getOperationType());
        assertEquals("HIGH", context.getPriority());
        
        verify(listener1).onStart(context);
        verify(listener2).onStart(context);
        verify(metricsUtils).getMessagePriority(consumerRecord);
    }
}

// 测试用Payload类
static class TestPayload {}

注意：若无法避免静态方法调用，可使用PowerMock扩展Mockito，或重构代码将静态逻辑封装为可注入服务，提升可测试性。

三、测试 afterRecord 方法 #

需要验证以下核心行为：

• 正确获取ThreadLocal中的MessageContext
• 所有MessageContextListener的onClear方法触发
• ThreadLocal被清空

示例测试代码：

@Test
void afterRecord_shouldNotifyListenersAndClearThreadLocal() {
    // 1. 准备测试数据
    ConsumerRecord<String, Message<Object>> consumerRecord = mock(ConsumerRecord.class);
    Consumer<String, Message<Object>> mockConsumer = mock(Consumer.class);
    MessageContext testContext = new MessageContext();
    KafkaRecordInterceptor.messageContextThreadLocal.set(testContext);

    // 2. 执行测试方法
    interceptor.afterRecord(consumerRecord, mockConsumer);

    // 3. 验证结果
    verify(listener1).onClear(testContext);
    verify(listener2).onClear(testContext);
    assertNull(KafkaRecordInterceptor.messageContextThreadLocal.get());
}

四、边界情况补充测试 #

• 当metricsUtils.getMessagePriority返回空Optional时，验证messageContext.getPriority()为空字符串
• 若业务允许consumerRecord.value()非GenericMessage类型，验证异常处理逻辑
• 当messageContextListener为空列表时，验证方法无异常抛出

内容的提问来源于stack exchange，提问作者Bogdan Alin Muresan