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.workflow;
  
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertNotNull;
  import static org.junit.Assert.assertNull;
  import static org.junit.Assert.fail;
  import static org.mockito.Mockito.any;
  import static org.mockito.Mockito.doAnswer;
  import static org.mockito.Mockito.mock;
  import static org.mockito.Mockito.never;
  import static org.mockito.Mockito.verify;
  import static org.mockito.Mockito.when;
  
  import alluxio.exception.status.ResourceExhaustedException;
  import alluxio.job.JobConfig;
  import alluxio.job.JobServerContext;
  import alluxio.job.SleepJobConfig;
  import alluxio.job.TestPlanConfig;
  import alluxio.job.plan.PlanConfig;
  import alluxio.job.plan.meta.PlanInfo;
  import alluxio.job.wire.Status;
  import alluxio.job.wire.WorkflowInfo;
  import alluxio.job.workflow.composite.CompositeConfig;
  import alluxio.master.job.JobMaster;
  import alluxio.master.job.command.CommandManager;
  import alluxio.master.job.plan.PlanTracker;
  import alluxio.wire.WorkerInfo;
  
  import com.google.common.collect.Lists;
  import org.junit.Before;
  import org.junit.Test;
  
  import java.util.ArrayList;
  
  public class WorkflowTrackerTest {
  
    private static final long CAPACITY = 5;
    private static final long RETENTION_TIME = 0;
    private static final long PURGE_CONUT = -1;
  
    private PlanTracker mPlanTracker;
    private WorkflowTracker mWorkflowTracker;
    private JobMaster mMockJobMaster;
    private long mJobIdCounter;
  
    private ArrayList<WorkerInfo> mWorkers;
    private CommandManager mCommandManager;
    private JobServerContext mMockJobServerContext;
  
    @Before
    public void before() throws Exception {
      mMockJobMaster = mock(JobMaster.class);
      mWorkflowTracker = new WorkflowTracker(mMockJobMaster);
      mPlanTracker = new PlanTracker(CAPACITY, RETENTION_TIME, PURGE_CONUT, mWorkflowTracker);
  
      mJobIdCounter = 100;
      when(mMockJobMaster.getNewJobId()).thenAnswer(invocation -> mJobIdCounter++);
  
      mWorkers = Lists.newArrayList(new WorkerInfo());
      mCommandManager = new CommandManager();
      mMockJobServerContext = mock(JobServerContext.class);
    }
  
    @Test
    public void testEmpty() throws Exception {
      ArrayList<JobConfig> jobs = Lists.newArrayList();
      CompositeConfig config = new CompositeConfig(jobs, true);
      mWorkflowTracker.run(config, 0);
      WorkflowInfo info = mWorkflowTracker.getStatus(0, true);
  
      assertEquals(Status.COMPLETED, info.getStatus());
    }
  
    @Test
    public void testBasic() throws Exception {
  
      ArrayList<JobConfig> jobs = Lists.newArrayList();
  
      TestPlanConfig child1 = new TestPlanConfig("1");
      TestPlanConfig child2 = new TestPlanConfig("2");
      jobs.add(child1);
      jobs.add(child2);
  
      CompositeConfig config = new CompositeConfig(jobs, true);
  
      mWorkflowTracker.run(config, 0);
  
      verify(mMockJobMaster).run(child1, 100);
  
      WorkflowInfo info = mWorkflowTracker.getStatus(0, true);
  
      assertEquals(Status.RUNNING, info.getStatus());
  
      verify(mMockJobMaster, never()).run(child2, 101);
  
      PlanInfo plan100 = new PlanInfo(100, child1, null);
      plan100.setStatus(Status.COMPLETED);
      mWorkflowTracker.onPlanStatusChange(plan100);
  
      verify(mMockJobMaster).run(child2, 101);
  
      assertEquals(Status.RUNNING, mWorkflowTracker.getStatus(0, true).getStatus());
  
      PlanInfo plan101 = new PlanInfo(101, child2, null);
      plan101.setStatus(Status.COMPLETED);
      mWorkflowTracker.onPlanStatusChange(plan101);
  
      assertEquals(Status.COMPLETED, mWorkflowTracker.getStatus(0, true).getStatus());
    }
  
    @Test
    public void testCleanup() throws Exception {
      SleepJobConfig jobConfig = new SleepJobConfig(1);
      mPlanTracker.run(jobConfig, mCommandManager, mMockJobServerContext, mWorkers, 1);
  
      jobConfig = new SleepJobConfig(1);
      mPlanTracker.run(jobConfig, mCommandManager, mMockJobServerContext, mWorkers, 2);
  
      jobConfig = new SleepJobConfig(1);
      mPlanTracker.run(jobConfig, mCommandManager, mMockJobServerContext, mWorkers, 3);
  
      doAnswer(invocation -> {
        PlanConfig config = invocation.getArgument(0, PlanConfig.class);
        long jobId = invocation.getArgument(1, Long.class);
  
        mPlanTracker.run(config, mCommandManager, mMockJobServerContext, mWorkers, jobId);
        return null;
      }).when(mMockJobMaster).run(any(PlanConfig.class), any(Long.class));
  
      ArrayList<JobConfig> jobs = Lists.newArrayList();
  
      SleepJobConfig child1 = new SleepJobConfig(1);
      SleepJobConfig child2 = new SleepJobConfig(2);
      jobs.add(child1);
      jobs.add(child2);
  
      CompositeConfig config = new CompositeConfig(jobs, false);
  
      mWorkflowTracker.run(config, 0);
  
      try {
        mPlanTracker.run(new SleepJobConfig(1), mCommandManager, mMockJobServerContext, mWorkers, 4);
        fail();
      } catch (ResourceExhaustedException e) {
        // Should fail
      }
  
      mPlanTracker.coordinators().stream().filter(coordinator -> coordinator.getJobId() == 100)
          .findFirst().get().setJobAsFailed("TestError", "failed");
  
      mPlanTracker.run(new SleepJobConfig(1), mCommandManager, mMockJobServerContext, mWorkers, 4);
  
      assertNotNull(mWorkflowTracker.getStatus(0, true));
  
      try {
        mPlanTracker.run(new SleepJobConfig(1), mCommandManager, mMockJobServerContext, mWorkers, 5);
        fail();
      } catch (ResourceExhaustedException e) {
        // Should fail
      }
  
      mPlanTracker.coordinators().stream().filter(coordinator -> coordinator.getJobId() == 101)
          .findFirst().get().setJobAsFailed("TestError", "failed");
  
      mPlanTracker.run(new SleepJobConfig(1), mCommandManager, mMockJobServerContext, mWorkers, 5);
  
      assertNull(mWorkflowTracker.getStatus(100, true));
    }
  }
  