alluxio-master-2.x/job/server/src/test/java/alluxio/worker/job/task/TaskExecutorTest.java
/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio.worker.job.task;

import static org.mockito.Mockito.any;
import static org.mockito.Mockito.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import alluxio.AlluxioMockUtil;
import alluxio.grpc.RunTaskCommand;
import alluxio.job.JobConfig;
import alluxio.job.RunTaskContext;
import alluxio.job.SleepJobConfig;
import alluxio.job.plan.PlanDefinition;
import alluxio.job.plan.PlanDefinitionRegistry;
import alluxio.job.util.SerializationUtils;

import com.google.common.collect.Lists;
import com.google.protobuf.ByteString;
import org.junit.Before;
import org.junit.Test;

import java.io.Serializable;

/**
 * Tests {@link TaskExecutor}.
 */
public final class TaskExecutorTest {
  private TaskExecutorManager mTaskExecutorManager;
  private PlanDefinitionRegistry mRegistry;

  @Before
  public void before() {
    mTaskExecutorManager = mock(TaskExecutorManager.class);
    mRegistry = mock(PlanDefinitionRegistry.class);
    AlluxioMockUtil.setInternalState(PlanDefinitionRegistry.class, "INSTANCE", mRegistry);
  }

  @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);
  }

  @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());
  }
}
