private static Environment getMockEnvironment(File[] tempDirs) {
IOManager ioMan = mock(IOManager.class);
when(ioMan.getSpillingDirectories()).thenReturn(tempDirs);
Environment env = mock(Environment.class);
when(env.getJobID()).thenReturn(new JobID());
when(env.getUserClassLoader()).thenReturn(RocksDBStateBackendConfigTest.class.getClassLoader());
when(env.getIOManager()).thenReturn(ioMan);
when(env.getTaskKvStateRegistry())
.thenReturn(new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()));
TaskInfo taskInfo = mock(TaskInfo.class);
when(env.getTaskInfo()).thenReturn(taskInfo);
when(taskInfo.getIndexOfThisSubtask()).thenReturn(0);
return env;
}
public void doMocking(AbstractInvokable taskMock) throws Exception {
TaskInfo taskInfoMock = mock(TaskInfo.class);
when(taskInfoMock.getTaskNameWithSubtasks()).thenReturn("dummyName");
TaskManagerRuntimeInfo tmRuntimeInfo = mock(TaskManagerRuntimeInfo.class);
when(tmRuntimeInfo.getConfiguration()).thenReturn(new Configuration());
task =
new Task(
mock(JobInformation.class),
new TaskInformation(
new JobVertexID(), "test task name", 1, 1, "foobar", new Configuration()),
mock(ExecutionAttemptID.class),
0,
0,
Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
Collections.<InputGateDeploymentDescriptor>emptyList(),
0,
mock(TaskStateHandles.class),
mock(MemoryManager.class),
mock(IOManager.class),
mock(NetworkEnvironment.class),
mock(BroadcastVariableManager.class),
mock(TaskManagerConnection.class),
mock(InputSplitProvider.class),
mock(CheckpointResponder.class),
mock(LibraryCacheManager.class),
mock(FileCache.class),
tmRuntimeInfo,
mock(TaskMetricGroup.class),
mock(ResultPartitionConsumableNotifier.class),
mock(PartitionStateChecker.class),
mock(Executor.class));
Field f = task.getClass().getDeclaredField("invokable");
f.setAccessible(true);
f.set(task, taskMock);
Field f2 = task.getClass().getDeclaredField("executionState");
f2.setAccessible(true);
f2.set(task, ExecutionState.RUNNING);
}