JUnit：Spring Batch中JDBC写入方法Mock测试失败排查

开发了一个向SQL数据库写入数据的Spring Batch应用，数据库写入逻辑被包裹在功能开关条件中（代码如下）。编写单元测试时，希望Mock数据库写入操作并验证功能开关逻辑，但测试始终失败，提示jdbcBatchItemWriter.write()方法从未被调用。尝试添加doNothing().when(jdbcBatchItemWriter).write(any(List.class));也未解决问题，寻求可行方案。

业务代码（TreatmentWriter.java） #

public class TreatmentWriter implements ItemWriter<TreatmentData> {

    JdbcBatchItemWriter<TreatmentData> jdbcBatchItemWriter;

    @Autowired
    @Qualifier("dataSource")
    private DataSource dataSource;

    @Autowired
    @Qualifier("shawQuery")
    private String shawQuery;

    @Autowired
    FeatureManager featureManager;

    @Override
    public void write(List<? extends TreatmentData> processedData) {
        for (TreatmentData data : processedData) {
            log.debug("Writing processed data: " + data.getTreatment());
        }
        try {
            if (featureManager.isActive(FeaturesToToggle.FEATURE_FLAG)) {
                log.debug("Query to update staging table: " + shawQuery);
                jdbcBatchItemWriter = new JdbcBatchItemWriter<>();
                jdbcBatchItemWriter.setItemSqlParameterSourceProvider(new BeanPropertyItemSqlParameterSourceProvider<>());
                jdbcBatchItemWriter.setDataSource(dataSource);
                jdbcBatchItemWriter.setSql(shawQuery);
                jdbcBatchItemWriter.write(processedData);
            } else {
                log.info("feature flag is disabled");
            }
        } catch (Exception ex) {
            log.error("Error writing to DB: {}", ex);
        }
    }
}

测试代码（TreatmentWriterTest.java） #

@ExtendWith(MockitoExtension.class)
@TestInstance(TestInstance.Lifecycle.PER_CLASS)
@ActiveProfiles("test")
@ContextConfiguration(classes = { TreatmentWriter.class})
public class TreatmentWriterTest {

    private MemoryAppender memoryAppender;

    @InjectMocks
    TreatmentWriter treatmentWriter;

    @Mock
    private FeatureManager featureManager;

    @Spy
    private JdbcBatchItemWriter<TreatmentData> jdbcBatchItemWriter;

    @Mock
    private SecretsBeans secretsBeans;

    @Mock
    private DataSource shawDataSourceCollectEtl;

    @BeforeEach
    public void init() throws NoSuchFieldException, IllegalAccessException {
        Logger logger = (Logger) LoggerFactory.getLogger("com.ally.cr");
        memoryAppender = new MemoryAppender();
        memoryAppender.setContext((LoggerContext) LoggerFactory.getILoggerFactory());
        logger.setLevel(Level.DEBUG);
        logger.addAppender(memoryAppender);
        memoryAppender.start();
    }

    @Test
    public void UT_whenFeatureFlagIsOn_writeToShaw() throws Exception {
        mockFeatureFlag(true); //setting flag to true

        Field shawQuery;
        shawQuery = treatmentWriter.getClass().getDeclaredField("shawQuery");
        shawQuery.setAccessible(true);
        shawQuery.set(treatmentWriter, "");

        TreatmentData treatmentData = createTreatmentData(); //generating test data
    
        treatmentWriter.write(Arrays.asList(treatmentData));
        assertThat(memoryAppender.search("Query to update Shaw staging table: ", Level.INFO).size()).isEqualTo(1);
        verify(jdbcBatchItemWriter).write(any(List.class));
    }
}

测试失败的核心原因是：在TreatmentWriter的write方法中，每次功能开关开启时，都会新建一个JdbcBatchItemWriter实例，并覆盖类成员变量jdbcBatchItemWriter。而测试代码中通过@Spy创建的是类初始的成员变量实例，实际执行时调用的是新创建的对象，因此verify的是未被使用过的初始实例，自然会提示方法从未被调用。

需要重构TreatmentWriter，避免在write方法中重复创建JdbcBatchItemWriter实例，改用依赖注入的方式管理该对象，并提前完成初始化配置。

重构后的业务代码（推荐构造注入） #

构造注入更利于依赖管理和单元测试，同时将JdbcBatchItemWriter的初始化逻辑移到构造方法中：

public class TreatmentWriter implements ItemWriter<TreatmentData> {

    private final JdbcBatchItemWriter<TreatmentData> jdbcBatchItemWriter;
    private final DataSource dataSource;
    private final String shawQuery;
    private final FeatureManager featureManager;

    @Autowired
    public TreatmentWriter(JdbcBatchItemWriter<TreatmentData> jdbcBatchItemWriter,
                          @Qualifier("dataSource") DataSource dataSource,
                          @Qualifier("shawQuery") String shawQuery,
                          FeatureManager featureManager) {
        this.jdbcBatchItemWriter = jdbcBatchItemWriter;
        this.dataSource = dataSource;
        this.shawQuery = shawQuery;
        this.featureManager = featureManager;
        
        // 初始化JdbcBatchItemWriter配置
        this.jdbcBatchItemWriter.setItemSqlParameterSourceProvider(new BeanPropertyItemSqlParameterSourceProvider<>());
        this.jdbcBatchItemWriter.setDataSource(dataSource);
        this.jdbcBatchItemWriter.setSql(shawQuery);
        
        // 必须调用afterPropertiesSet完成初始化（Spring Batch要求）
        try {
            this.jdbcBatchItemWriter.afterPropertiesSet();
        } catch (Exception e) {
            throw new RuntimeException("初始化JdbcBatchItemWriter失败", e);
        }
    }

    @Override
    public void write(List<? extends TreatmentData> processedData) {
        for (TreatmentData data : processedData) {
            log.debug("Writing processed data: " + data.getTreatment());
        }
        try {
            if (featureManager.isActive(FeaturesToToggle.FEATURE_FLAG)) {
                log.debug("Query to update staging table: " + shawQuery);
                jdbcBatchItemWriter.write(processedData);
            } else {
                log.info("feature flag is disabled");
            }
        } catch (Exception ex) {
            log.error("Error writing to DB: {}", ex);
        }
    }
}

调整后的测试代码 #

现在JdbcBatchItemWriter是注入的实例，测试时可以直接Mock，无需担心被覆盖：

@ExtendWith(MockitoExtension.class)
@TestInstance(TestInstance.Lifecycle.PER_CLASS)
@ActiveProfiles("test")
public class TreatmentWriterTest {

    private MemoryAppender memoryAppender;

    private TreatmentWriter treatmentWriter;

    @Mock
    private FeatureManager featureManager;

    @Mock
    private JdbcBatchItemWriter<TreatmentData> jdbcBatchItemWriter;

    @Mock
    @Qualifier("dataSource")
    private DataSource dataSource;

    @BeforeEach
    public void init() {
        // 初始化日志Appender
        Logger logger = (Logger) LoggerFactory.getLogger("com.ally.cr");
        memoryAppender = new MemoryAppender();
        memoryAppender.setContext((LoggerContext) LoggerFactory.getILoggerFactory());
        logger.setLevel(Level.DEBUG);
        logger.addAppender(memoryAppender);
        memoryAppender.start();

        // 手动创建TreatmentWriter实例，传入所有Mock对象
        treatmentWriter = new TreatmentWriter(jdbcBatchItemWriter, dataSource, "", featureManager);
    }

    @Test
    public void UT_whenFeatureFlagIsOn_writeToShaw() throws Exception {
        // 模拟功能开关开启
        when(featureManager.isActive(FeaturesToToggle.FEATURE_FLAG)).thenReturn(true);

        TreatmentData treatmentData = createTreatmentData();
        treatmentWriter.write(Arrays.asList(treatmentData));

        // 验证日志输出
        assertThat(memoryAppender.search("Query to update staging table: ", Level.DEBUG).size()).isEqualTo(1);
        // 验证write方法被调用
        verify(jdbcBatchItemWriter).write(any(List.class));
    }

    @Test
    public void UT_whenFeatureFlagIsOff_doNotWrite() throws Exception {
        // 模拟功能开关关闭
        when(featureManager.isActive(FeaturesToToggle.FEATURE_FLAG)).thenReturn(false);

        TreatmentData treatmentData = createTreatmentData();
        treatmentWriter.write(Arrays.asList(treatmentData));

        // 验证日志输出
        assertThat(memoryAppender.search("feature flag is disabled", Level.INFO).size()).isEqualTo(1);
        // 验证write方法从未被调用
        verify(jdbcBatchItemWriter, never()).write(any(List.class));
    }

    private TreatmentData createTreatmentData() {
        TreatmentData data = new TreatmentData();
        data.setTreatment("test-treatment");
        return data;
    }
}

内容的提问来源于stack exchange，提问作者shoaib-jpeg