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;
  }
Пример #2
0
  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);
  }