Examples with RunTaskContext alluxio.job.RunTaskContext used on opensource projects

Example 1 with RunTaskContext

use of alluxio.job.RunTaskContext in project alluxio by Alluxio.

the class ReplicateDefinitionTest method runTaskReplicateTestHelper.

/**
 * Helper function to run a replicate task.
 *
 * @param blockWorkers available block workers
 * @param mockInStream mock blockInStream returned by the Block Store
 * @param mockOutStream mock blockOutStream returned by the Block Store
 */
private void runTaskReplicateTestHelper(List<BlockWorkerInfo> blockWorkers, BlockInStream mockInStream, BlockOutStream mockOutStream) throws Exception {
    when(mMockFileSystem.getStatus(any(AlluxioURI.class))).thenReturn(mTestStatus);
    when(mMockFileSystemContext.getCachedWorkers()).thenReturn(blockWorkers);
    when(mMockBlockStore.getInStream(anyLong(), any(InStreamOptions.class))).thenReturn(mockInStream);
    when(mMockBlockStore.getInStream(any(BlockInfo.class), any(InStreamOptions.class), any(Map.class))).thenReturn(mockInStream);
    PowerMockito.mockStatic(BlockInStream.class);
    when(BlockInStream.create(any(FileSystemContext.class), any(BlockInfo.class), any(WorkerNetAddress.class), any(BlockInStreamSource.class), any(InStreamOptions.class))).thenReturn(mockInStream);
    when(mMockBlockStore.getOutStream(eq(TEST_BLOCK_ID), eq(TEST_BLOCK_SIZE), eq(LOCAL_ADDRESS), any(OutStreamOptions.class))).thenReturn(mockOutStream);
    when(mMockBlockStore.getInfo(TEST_BLOCK_ID)).thenReturn(mTestBlockInfo.setLocations(Lists.newArrayList(new BlockLocation().setWorkerAddress(ADDRESS_1))));
    PowerMockito.mockStatic(AlluxioBlockStore.class);
    when(AlluxioBlockStore.create(any(FileSystemContext.class))).thenReturn(mMockBlockStore);
    ReplicateConfig config = new ReplicateConfig(TEST_PATH, TEST_BLOCK_ID, 1);
    ReplicateDefinition definition = new ReplicateDefinition();
    definition.runTask(config, null, new RunTaskContext(1, 1, mMockJobServerContext));
}

Example 2 with RunTaskContext

use of alluxio.job.RunTaskContext in project alluxio by Alluxio.

the class TaskExecutorTest method runCompletion.

@Test
public void runCompletion() throws Exception {
    long jobId = 1;
    long taskId = 2;
    JobConfig jobConfig = mock(JobConfig.class);
    Serializable taskArgs = Lists.newArrayList(1);
    RunTaskContext context = mock(RunTaskContext.class);
    Integer taskResult = 1;
    @SuppressWarnings("unchecked") PlanDefinition<JobConfig, Serializable, Serializable> planDefinition = mock(PlanDefinition.class);
    when(mRegistry.getJobDefinition(any(JobConfig.class))).thenReturn(planDefinition);
    when(planDefinition.runTask(any(JobConfig.class), eq(taskArgs), any(RunTaskContext.class))).thenReturn(taskResult);
    RunTaskCommand command = RunTaskCommand.newBuilder().setJobConfig(ByteString.copyFrom(SerializationUtils.serialize(jobConfig))).setTaskArgs(ByteString.copyFrom(SerializationUtils.serialize(taskArgs))).build();
    TaskExecutor executor = new TaskExecutor(jobId, taskId, command, context, mTaskExecutorManager);
    executor.run();
    verify(planDefinition).runTask(any(JobConfig.class), eq(taskArgs), eq(context));
    verify(mTaskExecutorManager).notifyTaskCompletion(jobId, taskId, taskResult);
}

Example 3 with RunTaskContext

use of alluxio.job.RunTaskContext in project alluxio by Alluxio.

the class TaskExecutorTest method runFailure.

@Test
public void runFailure() throws Exception {
    long jobId = 1;
    long taskId = 2;
    JobConfig jobConfig = new SleepJobConfig(10);
    Serializable taskArgs = Lists.newArrayList(1);
    RunTaskContext context = mock(RunTaskContext.class);
    @SuppressWarnings("unchecked") PlanDefinition<JobConfig, Serializable, Serializable> planDefinition = mock(PlanDefinition.class);
    when(mRegistry.getJobDefinition(eq(jobConfig))).thenReturn(planDefinition);
    when(planDefinition.runTask(eq(jobConfig), any(Serializable.class), any(RunTaskContext.class))).thenThrow(new UnsupportedOperationException("failure"));
    RunTaskCommand command = RunTaskCommand.newBuilder().setJobConfig(ByteString.copyFrom(SerializationUtils.serialize(jobConfig))).setTaskArgs(ByteString.copyFrom(SerializationUtils.serialize(taskArgs))).build();
    TaskExecutor executor = new TaskExecutor(jobId, taskId, command, context, mTaskExecutorManager);
    executor.run();
    verify(mTaskExecutorManager).notifyTaskFailure(eq(jobId), eq(taskId), any());
}