content: /*
   * 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.tracker;
  
  import static org.mockito.Mockito.mock;
  import static org.mockito.Mockito.when;
  
  import alluxio.AlluxioURI;
  import alluxio.client.WriteType;
  import alluxio.client.file.FileSystem;
  import alluxio.client.file.FileSystemContext;
  import alluxio.exception.JobDoesNotExistException;
  import alluxio.grpc.OperationType;
  import alluxio.job.cmd.load.LoadCliConfig;
  import alluxio.job.cmd.migrate.MigrateCliConfig;
  import alluxio.job.wire.CmdStatusBlock;
  import alluxio.job.wire.JobSource;
  import alluxio.job.wire.SimpleJobStatusBlock;
  import alluxio.job.wire.Status;
  import alluxio.master.job.JobMaster;
  import alluxio.master.job.common.CmdInfo;
  import alluxio.master.job.plan.PlanTracker;
  import alluxio.master.job.workflow.WorkflowTracker;
  
  import com.beust.jcommander.internal.Lists;
  import org.junit.Assert;
  import org.junit.Before;
  import org.junit.Rule;
  import org.junit.Test;
  import org.junit.rules.ExpectedException;
  
  import java.util.Collections;
  import java.util.List;
  import java.util.Random;
  import java.util.Set;
  
  /**
   * Tests {@link CmdJobTrackerTest}.
   */
  public final class CmdJobTrackerTest {
    private static final int REPEATED_ATTEMPT_COUNT = 5;
    private static final int ONE_ATTEMPT = 1;
    private static final long CAPACITY = 100;
  
    private CmdJobTracker mCmdJobTracker;
    private FileSystem mFs;
  
    private long mLoadJobId;
    private long mMigrateJobId;
    private MigrateCliRunner mMigrateCliRunner;
    private DistLoadCliRunner mDistLoadRunner;
    private PersistRunner mPersistRunner;
    private PlanTracker mPlanTracker;
    private LoadCliConfig mLoad;
    private MigrateCliConfig mMigrate;
    private List<Status> mSearchingCriteria = Lists.newArrayList();
    private WorkflowTracker mWorkflowTracker;
    private JobMaster mMockJobMaster;
  
    private Long mRetentionTime;
  
    @Rule
    public ExpectedException mException = ExpectedException.none();
  
    @Before
    public void before() throws Exception {
      mFs = mock(FileSystem.class);
      mRetentionTime = 1000L;
      FileSystemContext fsCtx = mock(FileSystemContext.class);
  
      mMigrateCliRunner = mock(MigrateCliRunner.class);
      mDistLoadRunner = mock(DistLoadCliRunner.class);
      mPersistRunner = mock(PersistRunner.class);
      mPlanTracker = mock(PlanTracker.class);
      mMockJobMaster = mock(JobMaster.class);
      mWorkflowTracker = new WorkflowTracker(mMockJobMaster);
      mCmdJobTracker = new CmdJobTracker(fsCtx,
              mDistLoadRunner, mMigrateCliRunner, mPersistRunner, mRetentionTime, mPlanTracker);
  
      mLoad = new LoadCliConfig("/path/to/load", 3, 1, Collections.EMPTY_SET,
              Collections.EMPTY_SET, Collections.EMPTY_SET, Collections.EMPTY_SET, true);
      mMigrate = new MigrateCliConfig("/path/from", "/path/to", WriteType.THROUGH, true, 2);
      mLoadJobId = 1;
      mMigrateJobId = 2;
      mSearchingCriteria.clear();
    }
  
    //All tests below are for testing running progresses, job and command statuses.
    @Test
    public void runDistLoadBatchCompleteTest() throws Exception {
      CmdInfo cmdInfo = new CmdInfo(mLoadJobId, OperationType.DIST_LOAD,
              JobSource.CLI, System.currentTimeMillis(), Lists.newArrayList());
      prepareAttemptWithStatus(Status.COMPLETED, cmdInfo, REPEATED_ATTEMPT_COUNT);
  
      prepareDistLoadTest(cmdInfo, mLoad, mLoadJobId);
  
      mCmdJobTracker.run(mLoad, mLoadJobId);
      Status s = mCmdJobTracker.getCmdStatus(mLoadJobId);
      Assert.assertEquals(s, Status.COMPLETED);
    }
  
    @Test
    public void runCleanExpiredJobsTest() throws Exception {
      generateLoadCommandForStatus(Status.CANCELED);
      generateLoadCommandForStatus(Status.RUNNING);
      generateLoadCommandForStatus(Status.FAILED);
      generateLoadCommandForStatus(Status.COMPLETED);
      generateLoadCommandForStatus(Status.CREATED);
      Thread.sleep(70000L);
      // the expired job has been cleaned in mInfoMap
      mSearchingCriteria.clear();
      mSearchingCriteria.add(Status.CANCELED);
      Set<Long> cancelCmdIds = mCmdJobTracker.findCmdIds(mSearchingCriteria);
      Assert.assertEquals(0, cancelCmdIds.size());
      mSearchingCriteria.clear();
      mSearchingCriteria.add(Status.COMPLETED);
      Set<Long> completedCmdIds = mCmdJobTracker.findCmdIds(mSearchingCriteria);
      Assert.assertEquals(0, completedCmdIds.size());
      mSearchingCriteria.clear();
      mSearchingCriteria.add(Status.FAILED);
      Set<Long> failedCmdIds = mCmdJobTracker.findCmdIds(mSearchingCriteria);
      Assert.assertEquals(0, failedCmdIds.size());
      mSearchingCriteria.clear();
      mSearchingCriteria.add(Status.RUNNING);
      Set<Long> runningCmdIds = mCmdJobTracker.findCmdIds(mSearchingCriteria);
      Assert.assertEquals(2, runningCmdIds.size());
    }
  
    @Test
    public void runDistLoadBatchFailTest() throws Exception {
      CmdInfo cmdInfo = new CmdInfo(mLoadJobId, OperationType.DIST_LOAD,
              JobSource.CLI, System.currentTimeMillis(), Lists.newArrayList());
      prepareAttemptWithStatus(Status.FAILED, cmdInfo, ONE_ATTEMPT);
      prepareAttemptWithStatus(Status.COMPLETED, cmdInfo, REPEATED_ATTEMPT_COUNT);
      prepareAttemptWithStatus(Status.CANCELED, cmdInfo, ONE_ATTEMPT);
  
      prepareDistLoadTest(cmdInfo,  mLoad, mLoadJobId);
  
      mCmdJobTracker.run(mLoad, mLoadJobId);
      Status s = mCmdJobTracker.getCmdStatus(mLoadJobId);
      Assert.assertEquals(s, Status.FAILED);
    }
  
    @Test
    public void runDistLoadBatchCancelTest() throws Exception {
      CmdInfo cmdInfo = new CmdInfo(mLoadJobId, OperationType.DIST_LOAD,
              JobSource.CLI, System.currentTimeMillis(), Lists.newArrayList());
      prepareAttemptWithStatus(Status.COMPLETED, cmdInfo, REPEATED_ATTEMPT_COUNT);
      prepareAttemptWithStatus(Status.CANCELED, cmdInfo, ONE_ATTEMPT);
  
      prepareDistLoadTest(cmdInfo,  mLoad, mLoadJobId);
  
      mCmdJobTracker.run(mLoad, mLoadJobId);
      Status s = mCmdJobTracker.getCmdStatus(mLoadJobId);
      Assert.assertEquals(s, Status.CANCELED);
    }
  
    @Test
    public void runDistLoadBatchRunningTest() throws Exception {
      CmdInfo cmdInfo = new CmdInfo(mLoadJobId, OperationType.DIST_LOAD,
              JobSource.CLI, System.currentTimeMillis(), Lists.newArrayList());
      prepareAttemptWithStatus(Status.FAILED, cmdInfo, ONE_ATTEMPT);
      prepareAttemptWithStatus(Status.COMPLETED, cmdInfo, REPEATED_ATTEMPT_COUNT);
      prepareAttemptWithStatus(Status.CANCELED, cmdInfo, ONE_ATTEMPT);
      prepareAttemptWithStatus(Status.RUNNING, cmdInfo, ONE_ATTEMPT);
      prepareAttemptWithStatus(Status.CREATED, cmdInfo, REPEATED_ATTEMPT_COUNT);
  
      prepareDistLoadTest(cmdInfo,  mLoad, mLoadJobId);
  
      mCmdJobTracker.run(mLoad, mLoadJobId);
      Status s = mCmdJobTracker.getCmdStatus(mLoadJobId);
      Assert.assertEquals(s, Status.RUNNING);
    }
  
    @Test
    public void runDistCpBatchCompleteTest() throws Exception {
      CmdInfo cmdInfo = new CmdInfo(mMigrateJobId, OperationType.DIST_CP,
              JobSource.CLI, System.currentTimeMillis(), Lists.newArrayList());
      prepareAttemptWithStatus(Status.COMPLETED, cmdInfo, REPEATED_ATTEMPT_COUNT);
  
      prepareDistCpTest(cmdInfo, mMigrate, mMigrateJobId);
  
      mCmdJobTracker.run(mMigrate, mMigrateJobId);
      Status s = mCmdJobTracker.getCmdStatus(mMigrateJobId);
      Assert.assertEquals(s, Status.COMPLETED);
    }
  
    @Test
    public void runDistCpBatchFailTest() throws Exception {
      CmdInfo cmdInfo = new CmdInfo(mMigrateJobId, OperationType.DIST_CP,
              JobSource.CLI, System.currentTimeMillis(), Lists.newArrayList());
      prepareAttemptWithStatus(Status.FAILED, cmdInfo, ONE_ATTEMPT);
      prepareAttemptWithStatus(Status.COMPLETED, cmdInfo, REPEATED_ATTEMPT_COUNT);
      prepareAttemptWithStatus(Status.CANCELED, cmdInfo, ONE_ATTEMPT);
  
      prepareDistCpTest(cmdInfo, mMigrate, mMigrateJobId);
  
      mCmdJobTracker.run(mMigrate, mMigrateJobId);
      Status s = mCmdJobTracker.getCmdStatus(mMigrateJobId);
      Assert.assertEquals(s, Status.FAILED);
    }
  
    @Test
    public void runDistCpBatchCancelTest() throws Exception {
      CmdInfo cmdInfo = new CmdInfo(mMigrateJobId, OperationType.DIST_CP,
              JobSource.CLI, System.currentTimeMillis(), Lists.newArrayList());
      prepareAttemptWithStatus(Status.COMPLETED, cmdInfo, REPEATED_ATTEMPT_COUNT);
      prepareAttemptWithStatus(Status.CANCELED, cmdInfo, ONE_ATTEMPT);
  
      prepareDistCpTest(cmdInfo, mMigrate, mMigrateJobId);
  
      mCmdJobTracker.run(mMigrate, mMigrateJobId);
      Status s = mCmdJobTracker.getCmdStatus(mMigrateJobId);
      Assert.assertEquals(s, Status.CANCELED);
    }
  
    @Test
    public void runDistCpBatchRunningTest() throws Exception {
      CmdInfo cmdInfo = new CmdInfo(mMigrateJobId, OperationType.DIST_CP,
              JobSource.CLI, System.currentTimeMillis(), Lists.newArrayList());
      prepareAttemptWithStatus(Status.FAILED, cmdInfo, ONE_ATTEMPT);
      prepareAttemptWithStatus(Status.COMPLETED, cmdInfo, REPEATED_ATTEMPT_COUNT);
      prepareAttemptWithStatus(Status.CANCELED, cmdInfo, ONE_ATTEMPT);
      prepareAttemptWithStatus(Status.RUNNING, cmdInfo, ONE_ATTEMPT);
      prepareAttemptWithStatus(Status.CREATED, cmdInfo, REPEATED_ATTEMPT_COUNT);
  
      prepareDistCpTest(cmdInfo, mMigrate, mMigrateJobId);
  
      mCmdJobTracker.run(mMigrate, mMigrateJobId);
      Status s = mCmdJobTracker.getCmdStatus(mMigrateJobId);
      Assert.assertEquals(s, Status.RUNNING);
    }
  
    @Test
    public void testFindCmdIdsForComplete() throws Exception {
      long completedId = generateMigrateCommandForStatus(Status.COMPLETED);
      mSearchingCriteria.add(Status.COMPLETED);
      Set<Long> completedCmdIds = mCmdJobTracker.findCmdIds(mSearchingCriteria);
      Assert.assertEquals(completedCmdIds.size(), 1);
      Assert.assertTrue(completedCmdIds.contains(completedId));
    }
  
    @Test
    public void testFindCmdIdsForFailed() throws Exception {
      long failedId = generateMigrateCommandForStatus(Status.FAILED);
      mSearchingCriteria.add(Status.FAILED);
      Set<Long> failedCmdIds = mCmdJobTracker.findCmdIds(mSearchingCriteria);
      Assert.assertEquals(failedCmdIds.size(), 1);
      Assert.assertTrue(failedCmdIds.contains(failedId));
    }
  
    @Test
    public void testFindCmdIdsForRunning() throws Exception {
      long runningId = generateMigrateCommandForStatus(Status.RUNNING);
      mSearchingCriteria.add(Status.RUNNING);
      Set<Long> runningCmdIds = mCmdJobTracker.findCmdIds(mSearchingCriteria);
      Assert.assertEquals(runningCmdIds.size(), 1);
      Assert.assertTrue(runningCmdIds.contains(runningId));
    }
  
    @Test
    public void testFindCmdIdsForCancel() throws Exception {
      long cancelId = generateMigrateCommandForStatus(Status.CANCELED);
      mSearchingCriteria.add(Status.CANCELED);
      Set<Long> cancelCmdIds = mCmdJobTracker.findCmdIds(mSearchingCriteria);
      Assert.assertEquals(cancelCmdIds.size(), 1);
      Assert.assertTrue(cancelCmdIds.contains(cancelId));
    }
  
    @Test
    public void testFindCmdIdsForMultipleCmds() throws Exception {
      long cancelId = generateLoadCommandForStatus(Status.CANCELED);
      long runningIdA = generateLoadCommandForStatus(Status.RUNNING);
      long runningIdB = generateLoadCommandForStatus(Status.RUNNING);
      long failedId = generateMigrateCommandForStatus(Status.FAILED);
      long completedIdA = generateMigrateCommandForStatus(Status.COMPLETED);
      long completedIB = generateMigrateCommandForStatus(Status.COMPLETED);
      long createdId = generateMigrateCommandForStatus(Status.CREATED);
  
      // test cancel cmd ids.
      mSearchingCriteria.add(Status.CANCELED);
      Set<Long> cancelCmdIds = mCmdJobTracker.findCmdIds(mSearchingCriteria);
      Assert.assertEquals(cancelCmdIds.size(), 1);
      Assert.assertTrue(cancelCmdIds.contains(cancelId));
  
      // test completed cmd ids.
      mSearchingCriteria.clear();
      mSearchingCriteria.add(Status.COMPLETED);
      Set<Long> completedCmdIds = mCmdJobTracker.findCmdIds(mSearchingCriteria);
      Assert.assertEquals(completedCmdIds.size(), 2);
      Assert.assertTrue(completedCmdIds.contains(completedIdA));
      Assert.assertTrue(completedCmdIds.contains(completedIB));
  
      // test failed cmd ids.
      mSearchingCriteria.clear();
      mSearchingCriteria.add(Status.FAILED);
      Set<Long> failCmdIds = mCmdJobTracker.findCmdIds(mSearchingCriteria);
      Assert.assertEquals(failCmdIds.size(), 1);
      Assert.assertTrue(failCmdIds.contains(failedId));
  
      // test running cmd ids.
      mSearchingCriteria.clear();
      mSearchingCriteria.add(Status.RUNNING);
      Set<Long> runningCmdIds = mCmdJobTracker.findCmdIds(mSearchingCriteria);
      Assert.assertEquals(runningCmdIds.size(), 3); // 2 running commands + 1 created command.
      Assert.assertTrue(runningCmdIds.contains(runningIdA));
      Assert.assertTrue(runningCmdIds.contains(runningIdB));
      Assert.assertTrue(runningCmdIds.contains(createdId));
  
      // test running and completed cmd ids.
      mSearchingCriteria.clear();
      mSearchingCriteria.add(Status.COMPLETED);
      mSearchingCriteria.add(Status.RUNNING);
      Set<Long> ids = mCmdJobTracker.findCmdIds(mSearchingCriteria);
      Assert.assertEquals(ids.size(), 5);
      Assert.assertTrue(ids.contains(completedIdA));
      Assert.assertTrue(ids.contains(completedIB));
      Assert.assertTrue(ids.contains(runningIdA));
      Assert.assertTrue(ids.contains(runningIdB));
      Assert.assertTrue(ids.contains(createdId));
    }
  
    // test `getCmdStatusBlock`
    @Test
    public void testGetCmdStatusBlock() throws Exception {
      CmdInfo cmdInfo = new CmdInfo(mMigrateJobId, OperationType.DIST_CP,
              JobSource.CLI, System.currentTimeMillis(), Lists.newArrayList());
      CmdStatusBlock expectedStatusBlock = new CmdStatusBlock(mMigrateJobId, OperationType.DIST_CP);
  
      int completedCount = REPEATED_ATTEMPT_COUNT;
      int runningCount = REPEATED_ATTEMPT_COUNT;
      int offset = completedCount;
      addJobStatusBlockWithStatus(Status.COMPLETED, cmdInfo, completedCount,
              expectedStatusBlock, 0);
      addJobStatusBlockWithStatus(Status.RUNNING, cmdInfo, runningCount,
              expectedStatusBlock, offset);
  
      prepareDistCpTest(cmdInfo, mMigrate, mMigrateJobId);
  
      mCmdJobTracker.run(mMigrate, mMigrateJobId);
      CmdStatusBlock actualStatusBlock = mCmdJobTracker.getCmdStatusBlock(mMigrateJobId);
      Assert.assertEquals(actualStatusBlock.getJobControlId(), expectedStatusBlock.getJobControlId());
      Assert.assertEquals(actualStatusBlock.toProto(), expectedStatusBlock.toProto());
    }
  
   // Below are all help functions.
    private void prepareDistLoadTest(
           CmdInfo cmdInfo, LoadCliConfig loadCliConfig, long loadId) throws Exception {
      AlluxioURI filePath = new AlluxioURI(loadCliConfig.getFilePath());
      int replication = loadCliConfig.getReplication();
      Set<String> workerSet = loadCliConfig.getWorkerSet();
      Set<String> excludedWorkerSet = loadCliConfig.getExcludedWorkerSet();
      Set<String> localityIds = loadCliConfig.getLocalityIds();
      Set<String> excludedLocalityIds = loadCliConfig.getExcludedLocalityIds();
      boolean directCache = loadCliConfig.getDirectCache();
      int batch = loadCliConfig.getBatchSize();
      // Mock the behavior of runDistLoad
      when(mDistLoadRunner.runDistLoad(batch, filePath, replication, workerSet,
              excludedWorkerSet, localityIds, excludedLocalityIds, directCache, loadId))
              .thenReturn(cmdInfo);
    }
  
    private void prepareDistCpTest(
            CmdInfo cmdInfo, MigrateCliConfig migrateCliConfig, long migrateId) throws Exception {
      AlluxioURI src = new AlluxioURI(migrateCliConfig.getSource());
      AlluxioURI dst = new AlluxioURI(migrateCliConfig.getDestination());
      boolean overwt = migrateCliConfig.getOverWrite();
      int batch = migrateCliConfig.getBatchSize();
  
      when(mMigrateCliRunner.runDistCp(src, dst, overwt, batch, migrateId))
              .thenReturn(cmdInfo);
    }
  
    private void prepareAttemptWithStatus(Status status, CmdInfo cmdInfo, int number) {
      for (int i = 0; i < number; i++) {
        CmdRunAttempt attempt = mock(CmdRunAttempt.class);
        long jobId = new Random().nextLong();
        when(attempt.getJobId()).thenReturn(jobId); //sequential job ids.
        when(attempt.checkJobStatus()).thenReturn(status);
        cmdInfo.addCmdRunAttempt(attempt);
      }
    }
  
    // generate command and run the Load command, return job control ID.
    private long generateLoadCommandForStatus(Status status) throws Exception {
      long jobControlId = new Random().nextLong();
      LoadCliConfig config = new LoadCliConfig(
              "/path/to/load", 3, 1, Collections.EMPTY_SET,
              Collections.EMPTY_SET, Collections.EMPTY_SET, Collections.EMPTY_SET, true);
  
      CmdInfo cmdInfo = new CmdInfo(jobControlId, OperationType.DIST_LOAD,
              JobSource.CLI, System.currentTimeMillis(), Lists.newArrayList());
  
      switch (status) {
        case COMPLETED:
          prepareAttemptWithStatus(Status.COMPLETED, cmdInfo, REPEATED_ATTEMPT_COUNT);
          break;
        case FAILED:
          prepareAttemptWithStatus(Status.FAILED, cmdInfo, ONE_ATTEMPT);
          break;
        case RUNNING:
          prepareAttemptWithStatus(Status.RUNNING, cmdInfo, ONE_ATTEMPT);
          break;
        case CANCELED:
          prepareAttemptWithStatus(Status.CANCELED, cmdInfo, ONE_ATTEMPT);
          break;
        case CREATED:
          prepareAttemptWithStatus(Status.CREATED, cmdInfo, REPEATED_ATTEMPT_COUNT);
          break;
        default:
          throw new JobDoesNotExistException("No such job");
      }
  
      prepareDistLoadTest(cmdInfo, config, jobControlId);
      mCmdJobTracker.run(config, jobControlId);
      return jobControlId;
    }
  
    // generate command and run the Migrate command, return job control ID.
    private long generateMigrateCommandForStatus(Status status) throws Exception {
      long jobControlId = new Random().nextLong();
      MigrateCliConfig config = new MigrateCliConfig(
              "/path/from", "/path/to", WriteType.THROUGH, true, 2);
  
      CmdInfo cmdInfo = new CmdInfo(jobControlId, OperationType.DIST_CP,
              JobSource.CLI, System.currentTimeMillis(), Lists.newArrayList());
  
      switch (status) {
        case COMPLETED:
          prepareAttemptWithStatus(Status.COMPLETED, cmdInfo, REPEATED_ATTEMPT_COUNT);
          break;
        case FAILED:
          prepareAttemptWithStatus(Status.FAILED, cmdInfo, ONE_ATTEMPT);
          break;
        case RUNNING:
          prepareAttemptWithStatus(Status.RUNNING, cmdInfo, ONE_ATTEMPT);
          break;
        case CANCELED:
          prepareAttemptWithStatus(Status.CANCELED, cmdInfo, ONE_ATTEMPT);
          break;
        case CREATED:
          prepareAttemptWithStatus(Status.CREATED, cmdInfo, REPEATED_ATTEMPT_COUNT);
          break;
        default:
          throw new JobDoesNotExistException("No such job");
      }
  
      prepareDistCpTest(cmdInfo, config, jobControlId);
      mCmdJobTracker.run(config, jobControlId);
      return jobControlId;
    }
  
    private void addJobStatusBlockWithStatus(
            Status status, CmdInfo cmdInfo, int number, CmdStatusBlock cmdStatusBlock, int offset) {
      for (int i = 0; i < number; i++) {
        CmdRunAttempt attempt = mock(CmdRunAttempt.class);
        String filePath = String.format("filePath-%s", i + offset);
        when(attempt.getJobId()).thenReturn((long) (i + offset)); //sequential job ids.
        when(attempt.checkJobStatus()).thenReturn(status);
        when(attempt.getFilePath()).thenReturn(filePath);
        cmdInfo.addCmdRunAttempt(attempt);
        cmdStatusBlock.addJobStatusBlock(new SimpleJobStatusBlock(
                (long) (i + offset), status, filePath, ""));
      }
    }
  }
  