Spring Batch: How to cast stepContext.getJobExecutionContext() into custom Map implementation in mockito when then

Question

This question is related to This. Why I am asking here and not updating the question, because this can be helpful question and not related to the previous one which I have asked.

In a Test class how I can cast the below implementation: Lets say I have a map like this: TreeMap<String, Map<Integer, Set<Student>>>

and I want to cast the default chunkContext.getStepContext().getJobExecutionContext() which is Map<String, Object> to TreeMap<String, Map<Integer, Set<Student>>> in my Test class.

Existing implementation is:

Mockito.when(chunkContext.getStepContext().getJobExecutionContext().get("keyOfStudentMap"))
            .thenReturn((TreeMap<String, Map<Integer, Set<Student>>>)studentMap);

when I hover in getJobExecutionContext(), it shows Map<String, Object> and want to change in a way where this can changed to TreeMap<String, Map<Integer, Set<Student>>>

Sorry if anything is unclear. I can update the question based on your comments. :)

Answer 1

Instead of doing:

chunkContext.getStepContext().getJobExecutionContext()

you need to use:

chunkContext.getStepContext().getStepExecution().getJobExecution().getExecutionContext()

which returns an instance of ExecutionContext.

Now whatever you put as key in ExecutionContext, you get is as is when calling get and hence you can cast it to the original type (TreeMap<String, Map<Integer, Set<Student>>> in your case).