StepScopeなTaskletをテストする

0. わかればどうということはないのだが

これも使い古された内容だけに、多分に自分用メモ。

Spring Batchでいろんな社内処理を書いているんだが、先行するStepの処理結果を後続Stepで使いたい、というシーンはよくあると思う。

そんなとき、

1. 独自のワークテーブル、もしくはファイルに処理結果を出力
2. Beanにして持ちまわる
3. ExecutionContextに突っ込む

...あたりが候補にあげられるんじゃなかろうか。

上記のうち、1は一番柔軟に対応できるが設計が面倒だったりするし、Beanの場合は完全なインメモリ処理になるので、ちょっとした値のやり取りであればExecutionContextでお手軽に済ませる場合がほとんどだ。

1. AutowiredできないComponent

今回話題にするバッチのシナリオはこんな感じ。

1. APIアクセスでデータを取得するStep
2. 先行Stepの結果を加工するStep
3. API側にページング処理がないのでタスクレットモデルで処理

ひとまずソースを貼っておく。

BatchConfig

@Configuration
@EnableBatchProcessing
public class BatchConfig {

    private final JobRepository jobRepository;
    private final PlatformTransactionManager platformTransactionManager;

    public BatchConfig(JobRepository jobRepository, PlatformTransactionManager platformTransactionManager) {
        this.jobRepository = jobRepository;
        this.platformTransactionManager = platformTransactionManager;
    }

    /**
     * このプロジェクトの {@link Job } 。
     *
     * @param fooStep 先に実行される {@link Step }
     * @param barStep あとで実行される {@link Step }
     * @return {@link Job } のインスタンス
     */
    @Bean
    public Job job1(Step fooStep, Step barStep) {
        return new JobBuilder("job1", jobRepository)
                .start(fooStep)
                .next(barStep)
                .build();
    }

    /**
     * 先に実行される {@link Step } 。
     *
     * @param fooTasklet 先に実行されるタスクレット
     * @return {@link Step } のインスタンス
     */
    @Bean
    public Step fooStep(FooTasklet fooTasklet) {
        return new StepBuilder("foo_step", jobRepository)
                .tasklet(fooTasklet, platformTransactionManager)
                .listener(makeListener(FooTasklet.PASS_PARAM1))
                .build();
    }

    /**
     * あとで実行される {@link Step } 。
     *
     * @param barTasklet あとで実行されるタスクレット
     * @return {@link Step } のインスタンス
     */
    @Bean
    public Step barStep(BarTasklet barTasklet) {
        return new StepBuilder("bar_step", jobRepository)
                .tasklet(barTasklet, platformTransactionManager)
                .build();
    }

    /**
     * JobのExecutionContextへプッシュするリスナを作る。
     *
     * @param keys プッシュ対象を監視するキー
     * @return {@link ExecutionContextPromotionListener } のインスタンス
     */
    private ExecutionContextPromotionListener makeListener(String... keys) {
        ExecutionContextPromotionListener listener = new ExecutionContextPromotionListener();
        listener.setKeys(keys);

        return listener;
    }
}

先行するタスクレット

@Component
public class FooTasklet implements Tasklet {

    /**
     * 後続タスクへ引き渡す値のキー
     */
    public static final String PASS_PARAM1 = "PassParam1FromFooTasklet";

    @Override
    public RepeatStatus execute(StepContribution contribution, ChunkContext chunkContext) throws Exception {
        contribution.setExitStatus(ExitStatus.COMPLETED);
        ExecutionContext ec = chunkContext.getStepContext().getStepExecution().getExecutionContext();

        ec.put(PASS_PARAM1, "次に引き継ぐ値");

        Thread.sleep(15000L);

        return RepeatStatus.FINISHED;
    }
}

後続のタスクレット

@Component
@StepScope
public class BarTasklet implements Tasklet {

    private final Some1Service some1Service;
    private final Some2Service some2Service;

    public BarTasklet(Some1Service some1Service, Some2Service some2Service) {
        this.some1Service = some1Service;
        this.some2Service = some2Service;
    }

    @Override
    public RepeatStatus execute(StepContribution contribution, ChunkContext chunkContext) throws Exception {

        String val = (String) chunkContext.getStepContext().getStepExecution().getJobExecution().getExecutionContext().get(FooTasklet.PASS_PARAM1);

        System.out.println(val);
        some1Service.doSomething1();
        some2Service.doSomething2();

        contribution.setExitStatus(ExitStatus.COMPLETED);

        return RepeatStatus.FINISHED;
    }
}

後続のStepでは前の処理の値を使うので、通常のBean化戦略では実行時の値が決まっていないので実行の瞬間までBeanにならないようにしたい。
これは、StepScopeアノテーションを付加することで実現しているのだが、StepScopeになっているとテストコードにタスクレットをAutowiredでインジェクトできない、という制約が生じる。

まぁ、Bean化するとき中身が決まっていないのだから当然といえば当然なのだが、これでは都合が悪いので何らかの対策が必要になる。

ということで、テストはこうなる。

@SpringBootTest
class BarTaskletTest {

    @Autowired
    private Some1Service some1Service;
    @Autowired
    private Some2Service some2Service;

    @BeforeEach
    void setUp() {
    }

    @AfterEach
    void tearDown() {
    }

    @Test
    @DisplayName("BarTaskletを実行できる")
    void runsTasklet() throws Exception {
        // ChunkContextをモック化
        StepExecution se = Mockito.mock(StepExecution.class);
        StepContext sc = Mockito.mock(StepContext.class);
        ChunkContext cc = Mockito.mock(ChunkContext.class);

        JobExecution je = MetaDataInstanceFactory.createJobExecution();
        je.getExecutionContext().put(FooTasklet.PASS_PARAM1, "次に引き継ぐ値");

        Mockito.when(cc.getStepContext()).thenReturn(sc);
        Mockito.when(sc.getStepExecution()).thenReturn(se);
        Mockito.when(se.getJobExecution()).thenReturn(je);

        // StepContributionをモック化
        StepContribution contrib = Mockito.mock(StepContribution.class, Mockito.CALLS_REAL_METHODS);

        // StepScopeはAutowiredできないので自力でnewする
        BarTasklet barTasklet = new BarTasklet(some1Service, some2Service);

        RepeatStatus rs = StepScopeTestUtils.doInStepScope(se, () -> barTasklet.execute(contrib, cc));
        assertEquals(RepeatStatus.FINISHED, rs);
        assertEquals(ExitStatus.COMPLETED, contrib.getExitStatus());
    }
}

ミソは以下のとおり。
なお今回は、ExitStatusの検証もしたかったので、StepContributionのモック時にMockito.CALLS_REAL_METHODSも付加している。

• StepContributionとChunkContextを適切にモック化する
• StepScopeなTaskletは自力でnewしてインスタンス化する
• StepScopeTestUtils#doInStepScope(StepExecution, Callable<T>)内でタスクレットを実行して検証

2. 本日のソースコード

全文はこちらをどうぞ。