alluxio-master-2.x/job/server/src/test/java/alluxio/master/job/plan/PlanCoordinatorTest.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.master.job.plan;

import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import alluxio.AlluxioMockUtil;
import alluxio.client.file.FileSystem;
import alluxio.client.file.FileSystemContext;
import alluxio.collections.Pair;
import alluxio.grpc.JobCommand;
import alluxio.job.JobConfig;
import alluxio.job.JobServerContext;
import alluxio.job.SelectExecutorsContext;
import alluxio.job.plan.PlanDefinition;
import alluxio.job.plan.PlanDefinitionRegistry;
import alluxio.job.wire.Status;
import alluxio.job.wire.TaskInfo;
import alluxio.master.job.command.CommandManager;
import alluxio.underfs.UfsManager;
import alluxio.wire.WorkerInfo;

import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.mockito.Mockito;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;

/**
 * Tests {@link PlanCoordinator}.
 */
public final class PlanCoordinatorTest {
  private WorkerInfo mWorkerInfo;
  private long mJobId;
  private JobConfig mJobconfig;
  private JobServerContext mJobServerContext;
  private CommandManager mCommandManager;
  private PlanDefinition<JobConfig, Serializable, Serializable> mPlanDefinition;
  private List<WorkerInfo> mWorkerInfoList;

  @Before
  public void before() throws Exception {
    mCommandManager = new CommandManager();

    // Create mock JobServerContext
    FileSystem fs = mock(FileSystem.class);
    FileSystemContext fsCtx = mock(FileSystemContext.class);
    UfsManager ufsManager = Mockito.mock(UfsManager.class);
    mJobServerContext = new JobServerContext(fs, fsCtx, ufsManager);

    // Create mock job info.
    mJobconfig = mock(JobConfig.class, Mockito.withSettings().serializable());
    when(mJobconfig.getName()).thenReturn("mock");
    mJobId = 1;

    // Create mock job definition.
    @SuppressWarnings("unchecked")
    PlanDefinition<JobConfig, Serializable, Serializable> mockPlanDefinition =
        mock(PlanDefinition.class);
    PlanDefinitionRegistry singleton = mock(PlanDefinitionRegistry.class);
    AlluxioMockUtil.setInternalState(PlanDefinitionRegistry.class, "INSTANCE", singleton);

    when(singleton.getJobDefinition(mJobconfig)).thenReturn(mockPlanDefinition);
    mPlanDefinition = mockPlanDefinition;

    // Create test worker.
    mWorkerInfo = new WorkerInfo();
    mWorkerInfo.setId(0);
    mWorkerInfoList = Lists.newArrayList(mWorkerInfo);
  }

  @Test
  public void createJobCoordinator() throws Exception {
    mockSelectExecutors(mWorkerInfo);
    PlanCoordinator.create(
        mCommandManager, mJobServerContext, mWorkerInfoList, mJobId, mJobconfig, null);

    List<JobCommand> commands = mCommandManager.pollAllPendingCommands(mWorkerInfo.getId());
    Assert.assertEquals(1, commands.size());
    Assert.assertEquals(mJobId, commands.get(0).getRunTaskCommand().getJobId());
    Assert.assertEquals(0, commands.get(0).getRunTaskCommand().getTaskId());
  }

  @Test
  public void updateStatusFailure() throws Exception {
    mockSelectExecutors(mWorkerInfo);
    PlanCoordinator planCoordinator = PlanCoordinator.create(mCommandManager, mJobServerContext,
        mWorkerInfoList, mJobId, mJobconfig, null);
    setTasksWithStatuses(planCoordinator, Status.RUNNING, Status.FAILED, Status.COMPLETED);

    Assert.assertEquals(Status.FAILED, planCoordinator.getPlanInfoWire(true).getStatus());
    Assert.assertTrue(
        planCoordinator.getPlanInfoWire(true).getErrorMessage().contains("Task execution failed"));
  }

  @Test
  public void updateStatusFailureOverCancel() throws Exception {
    mockSelectExecutors(mWorkerInfo);
    PlanCoordinator planCoordinator = PlanCoordinator.create(mCommandManager, mJobServerContext,
        mWorkerInfoList, mJobId, mJobconfig, null);
    setTasksWithStatuses(planCoordinator, Status.RUNNING, Status.FAILED, Status.COMPLETED);

    Assert.assertEquals(Status.FAILED, planCoordinator.getPlanInfoWire(true).getStatus());
  }

  @Test
  public void updateStatusCancel() throws Exception {
    mockSelectExecutors(mWorkerInfo);
    PlanCoordinator planCoordinator = PlanCoordinator.create(mCommandManager, mJobServerContext,
        mWorkerInfoList, mJobId, mJobconfig, null);
    setTasksWithStatuses(planCoordinator, Status.CANCELED, Status.RUNNING, Status.COMPLETED);

    Assert.assertEquals(Status.CANCELED, planCoordinator.getPlanInfoWire(true).getStatus());
  }

  @Test
  public void updateStatusRunning() throws Exception {
    mockSelectExecutors(mWorkerInfo);
    PlanCoordinator planCoordinator = PlanCoordinator.create(mCommandManager, mJobServerContext,
        mWorkerInfoList, mJobId, mJobconfig, null);
    setTasksWithStatuses(planCoordinator, Status.COMPLETED, Status.RUNNING, Status.COMPLETED);

    Assert.assertEquals(Status.RUNNING, planCoordinator.getPlanInfoWire(true).getStatus());
  }

  @Test
  public void updateStatusCompleted() throws Exception {
    mockSelectExecutors(mWorkerInfo);
    PlanCoordinator planCoordinator = PlanCoordinator.create(mCommandManager, mJobServerContext,
        mWorkerInfoList, mJobId, mJobconfig, null);
    setTasksWithStatuses(planCoordinator, Status.COMPLETED, Status.COMPLETED, Status.COMPLETED);

    Assert.assertEquals(Status.COMPLETED, planCoordinator.getPlanInfoWire(true).getStatus());
  }

  @Test
  public void updateStatusJoinFailure() throws Exception {
    mockSelectExecutors(mWorkerInfo);
    when(mPlanDefinition.join(Mockito.eq(mJobconfig),
        Mockito.anyMapOf(WorkerInfo.class, Serializable.class)))
        .thenThrow(new UnsupportedOperationException("test exception"));
    PlanCoordinator planCoordinator = PlanCoordinator.create(mCommandManager, mJobServerContext,
        mWorkerInfoList, mJobId, mJobconfig, null);
    setTasksWithStatuses(planCoordinator, Status.COMPLETED, Status.COMPLETED, Status.COMPLETED);

    Assert.assertEquals(Status.FAILED, planCoordinator.getPlanInfoWire(true).getStatus());
    Assert.assertEquals("test exception", planCoordinator.getPlanInfoWire(true).getErrorMessage());
  }

  @Test
  public void noTasks() throws Exception {
    mockSelectExecutors();
    PlanCoordinator planCoordinator = PlanCoordinator.create(mCommandManager, mJobServerContext,
        mWorkerInfoList, mJobId, mJobconfig, null);
    Assert.assertEquals(Status.COMPLETED, planCoordinator.getPlanInfoWire(true).getStatus());
  }

  @Test
  public void failWorker() throws Exception {
    mockSelectExecutors(mWorkerInfo);
    PlanCoordinator planCoordinator = PlanCoordinator.create(mCommandManager, mJobServerContext,
        mWorkerInfoList, mJobId, mJobconfig, null);
    planCoordinator.failTasksForWorker(mWorkerInfo.getId());
    Assert.assertEquals(Status.FAILED, planCoordinator.getPlanInfoWire(true).getStatus());
  }

  /**
   * @param workerInfos the worker infos to return from the mocked selectExecutors method
   */
  private void mockSelectExecutors(WorkerInfo... workerInfos) throws Exception {
    Set<Pair<WorkerInfo, Serializable>> taskAddressToArgs = Sets.newHashSet();
    for (WorkerInfo workerInfo : workerInfos) {
      taskAddressToArgs.add(new Pair<>(workerInfo, null));
    }
    when(mPlanDefinition.selectExecutors(Mockito.eq(mJobconfig),
        Mockito.eq(Lists.newArrayList(mWorkerInfo)), Mockito.any(SelectExecutorsContext.class)))
        .thenReturn(taskAddressToArgs);
  }

  private void setTasksWithStatuses(PlanCoordinator planCoordinator, Status... statuses)
      throws Exception {
    List<TaskInfo> taskInfos = new ArrayList<>();
    int taskId = 0;
    for (Status status : statuses) {
      taskInfos.add(new TaskInfo().setTaskId(taskId++).setJobId(mJobId).setStatus(status));
    }
    planCoordinator.updateTasks(taskInfos);
  }
}
