无法Mock通过new创建的KafkaConsumer对象（不可修改源码）

DevOps 小满

2026-7-16

问题描述

我有如下类定义：

@Component
@EnableKafka
@Slf4j
@ConditionalOnProperty(prefix = "load", name = "kafka.config", havingValue = "True")
public class InsightApiKafkaScheduler extends InsightsAPISchedulerService {

    @Value("${kafka.source.topic}")
    String sourceTopic;

    private final ConsumerFactory<String,String> consumerFactory;
    List<String> insightsMessages;

    private final CompanyDeviceService deviceServiceImpl;
    private final CompanyDeviceService companyServiceImpl;
    private final WaterMarkUpdationServiceImpl waterMarkUpdationServiceImpl;
    private final RestService restService;
    private final KafkaService kafkaService;

    @Autowired
    public InsightApiKafkaScheduler(ConsumerFactory<String, String> consumerFactory, CompanyDeviceService deviceServiceImpl, CompanyDeviceService companyServiceImpl, WaterMarkUpdationServiceImpl waterMarkUpdationServiceImpl, RestService restService, KafkaService kafkaService) {
        this.consumerFactory = consumerFactory;
        this.deviceServiceImpl = deviceServiceImpl;
        this.companyServiceImpl = companyServiceImpl;
        this.restService = restService;
        this.kafkaService = kafkaService;
        this.auditPayload = new ArrayList<>();
        this.waterMarkUpdationServiceImpl = waterMarkUpdationServiceImpl;
    }

    @Override
    @Scheduled(cron = "${insights.api.schedule}")
    public void consumeMessages(){
        insightsMessages = new ArrayList<>();
        log.debug(String.valueOf(consumerFactory.getConfigurationProperties()));
        KafkaConsumer<String,String> kafkaConsumer = new KafkaConsumer<>(consumerFactory.getConfigurationProperties());
        ...
        ConsumerRecords<String,String> consumerRecords = kafkaConsumer.poll(Duration.ofMillis(20000L));
        ...
    }
...
}

现在要为consumeMessages()编写测试用例，需要MockKafkaConsumer对象来模拟kafkaConsumer.poll()返回的consumerRecords数据。由于KafkaConsumer是通过new关键字创建的，我尝试了以下代码来获取Mock对象：

KafkaConsumer<String, String> kafkaConsumer = Mockito.mock(KafkaConsumer.class);
PowerMockito.whenNew(KafkaConsumer.class).withAnyArguments().thenReturn(kafkaConsumer);

但这并不生效，程序执行到实际代码时仍会创建真实的KafkaConsumer对象。

在无法对实际代码进行任何修改/重构的前提下，我该如何创建KafkaConsumer的Mock对象？

补充：以下是我编写测试用例的方式：

@RunWith(SpringRunner.class)
public class InsightApiKafkaSchedulerTest {

    
    @SuppressWarnings("rawtypes")
    @Mock
    private ConsumerFactory<String,String> consumerFactory;

    @Mock
    private WaterMarkUpdationServiceImpl waterMarkUpdationServiceImpl;

    @Mock
    private KafkaConsumer<String, String> kafkaConsumer;

    @Mock
    private CompanyServiceImpl companyServiceImpl;

    @Mock
    private DeviceServiceImpl deviceServiceImpl;

    @Mock
    private RestService restService;

    @Mock
    private KafkaService kafkaService;
    
    @InjectMocks
    private InsightApiKafkaScheduler insightApiKafkaScheduler;
    
    @Before
    public void setup(){
        ReflectionTestUtils.setField(insightApiKafkaScheduler, "deviceServiceImpl", deviceServiceImpl);
        ReflectionTestUtils.setField(insightApiKafkaScheduler, "companyServiceImpl", companyServiceImpl);
        ReflectionTestUtils.setField(insightApiKafkaScheduler,"waterMarkUpdationServiceImpl", waterMarkUpdationServiceImpl);
        ReflectionTestUtils.setField(insightApiKafkaScheduler,"consumerFactory",consumerFactory);
        ReflectionTestUtils.setField(insightApiKafkaScheduler,"restService", restService);
        ReflectionTestUtils.setField(insightApiKafkaScheduler,"kafkaService", kafkaService);
    }

    @Test
    public void testConsumeMessages() throws Exception {
        long startOffset = 10L;
        ConsumerRecords<String, String> consumerRecords = buildConsumerRecords();

        HashMap<String, Object> config = new HashMap<>();
        config.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
        config.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG,StringDeserializer.class);
        config.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG,"my-server.com:9092");
        
        when(consumerFactory.getConfigurationProperties()).thenReturn(config);
        PowerMockito.whenNew(KafkaConsumer.class).withAnyArguments().thenReturn(kafkaConsumer);
        TopicPartition topicPartition = new TopicPartition("my-source-topic", 0);
        PowerMockito.whenNew(TopicPartition.class).withAnyArguments().thenReturn(topicPartition);
        
        when(kafkaConsumer.poll(any(Duration.class))).thenReturn(consumerRecords);
        when(waterMarkUpdationServiceImpl.getWatermark()).thenReturn(startOffset);
        doNothing().when(companyServiceImpl).sendMessages(any(), any());
        doNothing().when(kafkaService).commitOffsetsToKafka(any(), any(), any());
        doNothing().when(waterMarkUpdationServiceImpl).updateProcessTracking(any());

        insightApiKafkaScheduler.consumeMessages();
    }

    private ConsumerRecords<String, String> buildConsumerRecords() {
        List<String> messages = new ArrayList<>();
        
        // populates messages with test data
        
        return records;
    }
}

解决方案

测试用例未生效是因为缺少PowerMockito的必要配置，要拦截new关键字创建对象，必须满足以下要求：

替换测试运行器
将@RunWith(SpringRunner.class)替换为@RunWith(PowerMockRunner.class)，同时添加@PowerMockRunnerDelegate(SpringRunner.class)，这样既能保留Spring测试的能力，又能启用PowerMockito的字节码增强特性。
添加@PrepareForTest注解
在测试类上添加@PrepareForTest(InsightApiKafkaScheduler.class)，PowerMockito需要对创建KafkaConsumer的类进行字节码修改，才能拦截其中的new操作。
精准匹配构造方法（可选）
可以指定构造参数类型来替代withAnyArguments()，让Mock逻辑更严谨：

PowerMockito.whenNew(KafkaConsumer.class)
            .withParameterTypes(Map.class)
            .withArguments(anyMap())
            .thenReturn(kafkaConsumer);

完整修改后的测试类示例

@RunWith(PowerMockRunner.class)
@PowerMockRunnerDelegate(SpringRunner.class)
@PrepareForTest(InsightApiKafkaScheduler.class)
public class InsightApiKafkaSchedulerTest {

    @Mock
    private ConsumerFactory<String,String> consumerFactory;

    @Mock
    private WaterMarkUpdationServiceImpl waterMarkUpdationServiceImpl;

    @Mock
    private KafkaConsumer<String, String> kafkaConsumer;

    @Mock
    private CompanyServiceImpl companyServiceImpl;

    @Mock
    private DeviceServiceImpl deviceServiceImpl;

    @Mock
    private RestService restService;

    @Mock
    private KafkaService kafkaService;
    
    @InjectMocks
    private InsightApiKafkaScheduler insightApiKafkaScheduler;
    
    @Before
    public void setup(){
        ReflectionTestUtils.setField(insightApiKafkaScheduler, "deviceServiceImpl", deviceServiceImpl);
        ReflectionTestUtils.setField(insightApiKafkaScheduler, "companyServiceImpl", companyServiceImpl);
        ReflectionTestUtils.setField(insightApiKafkaScheduler,"waterMarkUpdationServiceImpl", waterMarkUpdationServiceImpl);
        ReflectionTestUtils.setField(insightApiKafkaScheduler,"consumerFactory",consumerFactory);
        ReflectionTestUtils.setField(insightApiKafkaScheduler,"restService", restService);
        ReflectionTestUtils.setField(insightApiKafkaScheduler,"kafkaService", kafkaService);
        // 补充设置sourceTopic字段，避免空指针
        ReflectionTestUtils.setField(insightApiKafkaScheduler, "sourceTopic", "my-source-topic");
    }

    @Test
    public void testConsumeMessages() throws Exception {
        long startOffset = 10L;
        ConsumerRecords<String, String> consumerRecords = buildConsumerRecords();

        HashMap<String, Object> config = new HashMap<>();
        config.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
        config.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG,StringDeserializer.class);
        config.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG,"my-server.com:9092");
        
        when(consumerFactory.getConfigurationProperties()).thenReturn(config);
        
        // 精准匹配KafkaConsumer的构造方法
        PowerMockito.whenNew(KafkaConsumer.class)
                    .withParameterTypes(Map.class)
                    .withArguments(config)
                    .thenReturn(kafkaConsumer);
        
        when(kafkaConsumer.poll(any(Duration.class))).thenReturn(consumerRecords);
        when(waterMarkUpdationServiceImpl.getWatermark()).thenReturn(startOffset);
        doNothing().when(companyServiceImpl).sendMessages(any(), any());
        doNothing().when(kafkaService).commitOffsetsToKafka(any(), any(), any());
        doNothing().when(waterMarkUpdationServiceImpl).updateProcessTracking(any());

        insightApiKafkaScheduler.consumeMessages();
        
        // 可添加断言验证方法调用
        verify(kafkaConsumer, times(1)).poll(any(Duration.class));
        verify(companyServiceImpl, times(1)).sendMessages(any(), any());
    }

    private ConsumerRecords<String, String> buildConsumerRecords() {
        // 正确构建测试用ConsumerRecords
        List<ConsumerRecord<String, String>> recordList = new ArrayList<>();
        TopicPartition tp = new TopicPartition("my-source-topic", 0);
        recordList.add(new ConsumerRecord<>("my-source-topic", 0, 10L, "key1", "test-message-1"));
        recordList.add(new ConsumerRecord<>("my-source-topic", 0, 11L, "key2", "test-message-2"));
        
        return new ConsumerRecords<>(Collections.singletonMap(tp, recordList));
    }
}

注意事项
1. 确保依赖中包含powermock-module-junit4和powermock-api-mockito2的正确版本，避免依赖冲突。
2. 如果测试类中还有其他通过new创建的对象需要Mock，需将对应的创建类也加入@PrepareForTest注解中。

内容的提问来源于stack exchange，提问作者Pramit Pakhira