org.apache.helix.task.Workflow Java Examples

@Test
public void testDifferentTasks() throws Exception {
  // Create a job with two different tasks
  String jobName = TestHelper.getTestMethodName();
  Workflow.Builder workflowBuilder = new Workflow.Builder(jobName);
  List<TaskConfig> taskConfigs = Lists.newArrayListWithCapacity(2);
  TaskConfig taskConfig1 = new TaskConfig("TaskOne", null);
  TaskConfig taskConfig2 = new TaskConfig("TaskTwo", null);
  taskConfigs.add(taskConfig1);
  taskConfigs.add(taskConfig2);
  Map<String, String> jobCommandMap = Maps.newHashMap();
  jobCommandMap.put("Timeout", "1000");
  JobConfig.Builder jobBuilder = new JobConfig.Builder().setCommand("DummyCommand")
      .addTaskConfigs(taskConfigs).setJobCommandConfigMap(jobCommandMap);
  workflowBuilder.addJob(jobName, jobBuilder);
  _driver.start(workflowBuilder.build());

  // Ensure the job completes
  _driver.pollForWorkflowState(jobName, TaskState.COMPLETED);

  // Ensure that each class was invoked
  Assert.assertTrue(_invokedClasses.contains(TaskOne.class.getName()));
  Assert.assertTrue(_invokedClasses.contains(TaskTwo.class.getName()));
}
 

@Test
public void testTaskAssignment() throws InterruptedException {
  _gSetupTool.getClusterManagementTool()
      .enableInstance(CLUSTER_NAME, PARTICIPANT_PREFIX + "_" + (_startPort + 0), false);
  String jobResource = TestHelper.getTestMethodName();
  JobConfig.Builder jobBuilder = new JobConfig.Builder().setCommand(MockTask.TASK_COMMAND)
      .setTargetResource(WorkflowGenerator.DEFAULT_TGT_DB);

  Workflow flow =
      WorkflowGenerator.generateSingleJobWorkflowBuilder(jobResource, jobBuilder).build();
  _driver.start(flow);

  // Wait 1 sec. The task should not be complete since it is not assigned.
  Thread.sleep(1000L);

  // The task is not assigned so the task state should be null in this case.
  Assert.assertNull(
      _driver.getJobContext(TaskUtil.getNamespacedJobName(jobResource)).getPartitionState(0));
}
 

/**
 * This test explicitly allows overlap job assignment.
 * @throws InterruptedException
 */
@Test
// This test does NOT allow multiple jobs being assigned to an instance.
public void testMultipleJobAssignmentOverlapEnabled() throws InterruptedException {
  _runCounts.clear();
  failTask = false;
  String workflowName = TestHelper.getTestMethodName();
  Workflow.Builder workflowBuilder = new Workflow.Builder(workflowName);
  WorkflowConfig.Builder configBuilder = new WorkflowConfig.Builder(workflowName);
  configBuilder.setAllowOverlapJobAssignment(true);
  workflowBuilder.setWorkflowConfig(configBuilder.build());

  for (int i = 0; i < 40; i++) {
    List<TaskConfig> taskConfigs = Lists.newArrayListWithCapacity(1);
    taskConfigs.add(new TaskConfig("TaskOne", new HashMap<>()));
    JobConfig.Builder jobBuilder = new JobConfig.Builder().setCommand("DummyCommand")
        .addTaskConfigs(taskConfigs).setJobCommandConfigMap(_jobCommandMap);
    workflowBuilder.addJob("JOB" + i, jobBuilder);
  }

  _driver.start(workflowBuilder.build());
  _driver.pollForWorkflowState(workflowName, TaskState.COMPLETED);

  Assert.assertEquals(_runCounts.size(), 5);
}
 

@Test
public void testJobAndWorkflowType() throws InterruptedException {
  LOG.info("Start testing job and workflow type");
  String jobName = TestHelper.getTestMethodName();
  JobConfig.Builder jobConfig = JobConfig.Builder.fromMap(WorkflowGenerator.DEFAULT_JOB_CONFIG)
      .setJobCommandConfigMap(WorkflowGenerator.DEFAULT_COMMAND_CONFIG).setJobType(DEFAULT_TYPE);

  Map<String, String> tmp = new HashMap<>();
  tmp.put("WorkflowType", DEFAULT_TYPE);
  Workflow.Builder builder =
      WorkflowGenerator.generateSingleJobWorkflowBuilder(jobName, jobConfig).fromMap(tmp);

  // Start workflow
  _driver.start(builder.build());

  _driver.pollForWorkflowState(jobName, TaskState.COMPLETED);
  String fetchedJobType =
      _driver.getJobConfig(String.format("%s_%s", jobName, jobName)).getJobType();
  String fetchedWorkflowType = _driver.getWorkflowConfig(jobName).getWorkflowType();

  Assert.assertEquals(fetchedJobType, DEFAULT_TYPE);
  Assert.assertEquals(fetchedWorkflowType, DEFAULT_TYPE);
}
 

/**
 * This test does NOT allow multiple jobs being assigned to an instance.
 * @throws InterruptedException
 */
@Test
public void testMultipleJobAssignment() throws InterruptedException {
  _runCounts.clear();
  failTask = false;
  String workflowName = TestHelper.getTestMethodName();
  Workflow.Builder workflowBuilder = new Workflow.Builder(workflowName);

  for (int i = 0; i < 20; i++) {
    List<TaskConfig> taskConfigs = Lists.newArrayListWithCapacity(1);
    taskConfigs.add(new TaskConfig("TaskOne", new HashMap<>()));
    JobConfig.Builder jobBuilder = new JobConfig.Builder().setCommand("DummyCommand")
        .addTaskConfigs(taskConfigs).setJobCommandConfigMap(_jobCommandMap);
    workflowBuilder.addJob("JOB" + i, jobBuilder);
  }

  _driver.start(workflowBuilder.build());
  _driver.pollForWorkflowState(workflowName, TaskState.COMPLETED);

  Assert.assertEquals(_runCounts.size(), 5);
}
 

@Test
public void testMultipleTaskAssignment() throws InterruptedException {
  _runCounts.clear();
  failTask = false;
  String workflowName = TestHelper.getTestMethodName();
  Workflow.Builder workflowBuilder = new Workflow.Builder(workflowName);

  List<TaskConfig> taskConfigs = Lists.newArrayListWithCapacity(20);
  for (int i = 0; i < 20; i++) {
    Map<String, String> taskConfigMap = Maps.newHashMap();
    taskConfigs.add(new TaskConfig("TaskOne", taskConfigMap));
  }
  JobConfig.Builder jobBuilder = new JobConfig.Builder().setCommand("DummyCommand")
      .setJobCommandConfigMap(_jobCommandMap).addTaskConfigs(taskConfigs);

  workflowBuilder.addJob("JOB", jobBuilder);
  _driver.start(workflowBuilder.build());
  _driver.pollForWorkflowState(workflowName, TaskState.COMPLETED);

  Assert.assertEquals(_runCounts.size(), 5);
}
 

/**
 * Task type: fixed target
 * Rebalance running task: disabled
 * Story: 1 node is down
 */
@Test
public void testFixedTargetTaskAndDisabledRebalanceAndNodeDown() throws InterruptedException {
  WORKFLOW = TestHelper.getTestMethodName();
  startParticipant(_initialNumNodes);

  JobConfig.Builder jobBuilder =
      new JobConfig.Builder().setWorkflow(WORKFLOW).setTargetResource(DATABASE)
          .setNumConcurrentTasksPerInstance(100).setCommand(MockTask.TASK_COMMAND)
          .setJobCommandConfigMap(ImmutableMap.of(MockTask.JOB_DELAY, "99999999"));

  Workflow.Builder workflowBuilder = new Workflow.Builder(WORKFLOW).addJob(JOB, jobBuilder);
  _driver.start(workflowBuilder.build());
  Assert.assertTrue(checkTasksOnDifferentInstances());
  // Stop a participant and partitions will be moved to the same instance,
  // and tasks rebalanced accordingly
  stopParticipant(_initialNumNodes);
  Assert.assertTrue(checkTasksOnSameInstances());
}
 

/**
 * Task type: generic
 * Rebalance running task: enabled
 * Story: new node added
 * NOTE: This test is disabled because this "load-balancing" would happen at the Task Assigner
 * level. In the legacy assignment strategy (Consistent Hashing) did not take instance's capacity
 * into account. However, the new quota-based scheduling takes capacity into account, and it will
 * generally assign to the most "free" instance, so load-balancing of tasks will happen at the
 * Assigner layer. Deprecating this test.
 */
@Deprecated
@Test(enabled = false)
public void testGenericTaskAndEnabledRebalanceAndNodeAdded() throws InterruptedException {
  WORKFLOW = TestHelper.getTestMethodName();
  JobConfig.Builder jobBuilder = new JobConfig.Builder().setWorkflow(WORKFLOW)
      .setNumberOfTasks(10).setNumConcurrentTasksPerInstance(100)
      .setCommand(MockTask.TASK_COMMAND).setRebalanceRunningTask(true)
      .setJobCommandConfigMap(ImmutableMap.of(MockTask.JOB_DELAY, "99999999")); // task stuck

  Workflow.Builder workflowBuilder = new Workflow.Builder(WORKFLOW).addJob(JOB, jobBuilder);

  _driver.start(workflowBuilder.build());

  // All tasks stuck on the same instance
  Assert.assertTrue(checkTasksOnSameInstances());
  // Add a new instance, and some running tasks will be rebalanced to the new node
  startParticipant(_initialNumNodes);
  Assert.assertTrue(checkTasksOnDifferentInstances());
}
 

/**
 * Task type: generic
 * Rebalance running task: disabled
 * Story: 1 node is down
 */
@Test
public void testGenericTaskAndDisabledRebalanceAndNodeDown() throws InterruptedException {
  WORKFLOW = TestHelper.getTestMethodName();
  startParticipant(_initialNumNodes);

  JobConfig.Builder jobBuilder =
      new JobConfig.Builder().setWorkflow(WORKFLOW).setNumberOfTasks(10)
          // should be enough for
          // consistent hashing to
          // place tasks on
          // different instances
          .setNumConcurrentTasksPerInstance(100).setCommand(MockTask.TASK_COMMAND)
          .setJobCommandConfigMap(ImmutableMap.of(MockTask.JOB_DELAY, "99999999")); // task stuck

  Workflow.Builder workflowBuilder = new Workflow.Builder(WORKFLOW).addJob(JOB, jobBuilder);

  _driver.start(workflowBuilder.build());

  Assert.assertTrue(checkTasksOnDifferentInstances());
  // Stop a participant, tasks rebalanced to the same instance
  stopParticipant(_initialNumNodes);
  Assert.assertTrue(checkTasksOnSameInstances());
}
 

@Test
public void testUnregisteredCommand() throws InterruptedException {
  String workflowName = TestHelper.getTestMethodName();
  Workflow.Builder builder = new Workflow.Builder(workflowName);

  JobConfig.Builder jobBuilder =
      new JobConfig.Builder().setTargetResource(WorkflowGenerator.DEFAULT_TGT_DB)
          .setCommand("OtherCommand").setTimeoutPerTask(10000L).setMaxAttemptsPerTask(2)
          .setJobCommandConfigMap(WorkflowGenerator.DEFAULT_COMMAND_CONFIG);

  builder.addJob("JOB1", jobBuilder);

  _driver.start(builder.build());

  _driver.pollForWorkflowState(workflowName, TaskState.FAILED);
  Assert.assertEquals(_driver.getJobContext(TaskUtil.getNamespacedJobName(workflowName, "JOB1"))
      .getPartitionState(0), TaskPartitionState.ERROR);
  Assert.assertEquals(_driver.getJobContext(TaskUtil.getNamespacedJobName(workflowName, "JOB1"))
      .getPartitionNumAttempts(0), 1);
}
 

@Test
public void testScheduleDelayTaskWithStartTime() throws InterruptedException {
  String workflowName = TestHelper.getTestMethodName();
  Workflow.Builder builder = new Workflow.Builder(workflowName);

  JobConfig.Builder jobBuilder =
      new JobConfig.Builder().setTargetResource(WorkflowGenerator.DEFAULT_TGT_DB)
          .setCommand(MockTask.TASK_COMMAND).setMaxAttemptsPerTask(2)
          .setJobCommandConfigMap(WorkflowGenerator.DEFAULT_COMMAND_CONFIG);

  long currentTime = System.currentTimeMillis();
  builder.addParentChildDependency("Job1", "Job2");
  builder.addJob("Job1", jobBuilder);
  builder.addJob("Job2", jobBuilder.setExecutionStart(currentTime + 5000L));

  _driver.start(builder.build());
  _driver.pollForJobState(workflowName, TaskUtil.getNamespacedJobName(workflowName, "Job2"),
      TaskState.COMPLETED);

  long jobTwoStartTime = _driver.getWorkflowContext(workflowName)
      .getJobStartTime(TaskUtil.getNamespacedJobName(workflowName, "Job2"));

  Assert.assertTrue(jobTwoStartTime - currentTime >= 5000L);
}
 

/**
 * Helper method for creating custom workflows.
 * @param workflowName
 * @param shouldOverlapJobAssign
 * @param quotaType
 * @param numJobs
 * @param numTasks
 * @param taskType
 * @return a workflow per parameters given
 */
private Workflow createWorkflow(String workflowName, boolean shouldOverlapJobAssign,
    String quotaType, int numJobs, int numTasks, String taskType) {

  Workflow.Builder workflowBuilder = new Workflow.Builder(workflowName);
  WorkflowConfig.Builder configBuilder = new WorkflowConfig.Builder(workflowName);
  configBuilder.setAllowOverlapJobAssignment(shouldOverlapJobAssign);
  workflowBuilder.setWorkflowConfig(configBuilder.build());

  for (int jobIndex = 0; jobIndex < numJobs; jobIndex++) {
    String jobName = workflowName + "_" + jobIndex;
    List<TaskConfig> taskConfigs = new ArrayList<>();
    for (int taskIndex = 0; taskIndex < numTasks; taskIndex++) {
      Map<String, String> taskConfigMap = new HashMap<>();
      taskConfigs.add(new TaskConfig(taskType, taskConfigMap));
    }
    JobConfig.Builder jobBuilder = new JobConfig.Builder().setCommand(JOB_COMMAND)
        .setJobCommandConfigMap(_jobCommandMap).addTaskConfigs(taskConfigs).setJobType(quotaType);
    workflowBuilder.addJob(jobName, jobBuilder);
  }
  return workflowBuilder.build();
}
 

/**
 * Tests whether jobs can run successfully without quotaTypes or quota configuration defined in
 * ClusterConfig. This test is to ensure backward-compatibility. This test must go first because
 * we want to make sure there is no quota config set anywhere.
 * @throws InterruptedException
 */
@Test
public void testSchedulingWithoutQuota() throws InterruptedException {
  String workflowName = TestHelper.getTestMethodName();
  Workflow.Builder workflowBuilder = new Workflow.Builder(workflowName);
  WorkflowConfig.Builder configBuilder = new WorkflowConfig.Builder(workflowName);
  configBuilder.setAllowOverlapJobAssignment(true);
  workflowBuilder.setWorkflowConfig(configBuilder.build());

  for (int i = 0; i < 10; i++) {
    List<TaskConfig> taskConfigs = new ArrayList<>();
    taskConfigs.add(new TaskConfig("ShortTask", new HashMap<>()));
    JobConfig.Builder jobConfigBulider = new JobConfig.Builder().setCommand(JOB_COMMAND)
        .addTaskConfigs(taskConfigs).setJobCommandConfigMap(_jobCommandMap);
    workflowBuilder.addJob("JOB" + i, jobConfigBulider);
  }

  _driver.start(workflowBuilder.build());
  _driver.pollForWorkflowState(workflowName, TaskState.COMPLETED);

  for (int i = 0; i < 10; i++) {
    String jobName = workflowName + "_" + "JOB" + i;
    TaskState jobState = _driver.getWorkflowContext(workflowName).getJobState(jobName);
    Assert.assertEquals(jobState, TaskState.COMPLETED);
  }
}
 

@Test
public void stopAndResume() throws Exception {
  Map<String, String> commandConfig = ImmutableMap.of(MockTask.JOB_DELAY, String.valueOf(100));

  JobConfig.Builder jobBuilder = JobConfig.Builder.fromMap(WorkflowGenerator.DEFAULT_JOB_CONFIG);
  jobBuilder.setJobCommandConfigMap(commandConfig);
  Workflow flow =
      WorkflowGenerator.generateSingleJobWorkflowBuilder(JOB_RESOURCE, jobBuilder).build();

  LOG.info("Starting flow " + flow.getName());
  _driver.start(flow);
  _driver.pollForWorkflowState(JOB_RESOURCE, TaskState.IN_PROGRESS);

  LOG.info("Pausing job");
  _driver.stop(JOB_RESOURCE);
  _driver.pollForWorkflowState(JOB_RESOURCE, TaskState.STOPPED);

  LOG.info("Resuming job");
  _driver.resume(JOB_RESOURCE);
  _driver.pollForWorkflowState(JOB_RESOURCE, TaskState.COMPLETED);
}
 

@Test
public void stopAndResumeWorkflow() throws Exception {
  String workflow = "SomeWorkflow";
  Workflow flow = WorkflowGenerator.generateDefaultRepeatedJobWorkflowBuilder(workflow).build();

  LOG.info("Starting flow " + workflow);
  _driver.start(flow);
  _driver.pollForWorkflowState(workflow, TaskState.IN_PROGRESS);

  LOG.info("Pausing workflow");
  _driver.stop(workflow);
  _driver.pollForWorkflowState(workflow, TaskState.STOPPED);

  LOG.info("Resuming workflow");
  _driver.resume(workflow);
  _driver.pollForWorkflowState(workflow, TaskState.COMPLETED);
}
 

@Test
public void testDeplayTimeAndStartTime() throws InterruptedException {
  String workflowName = TestHelper.getTestMethodName();
  Workflow.Builder builder = new Workflow.Builder(workflowName);

  JobConfig.Builder jobBuilder =
      new JobConfig.Builder().setTargetResource(WorkflowGenerator.DEFAULT_TGT_DB)
          .setCommand(MockTask.TASK_COMMAND).setMaxAttemptsPerTask(2)
          .setJobCommandConfigMap(WorkflowGenerator.DEFAULT_COMMAND_CONFIG);

  builder.addParentChildDependency("Job1", "Job2");

  long currentTime = System.currentTimeMillis();
  builder.addJob("Job1", jobBuilder);
  builder
      .addJob("Job2", jobBuilder.setExecutionDelay(2000L).setExecutionStart(currentTime + 5000L));

  _driver.start(builder.build());
  _driver.pollForJobState(workflowName, TaskUtil.getNamespacedJobName(workflowName, "Job2"),
      TaskState.COMPLETED);

  long jobTwoStartTime = _driver.getWorkflowContext(workflowName)
      .getJobStartTime(TaskUtil.getNamespacedJobName(workflowName, "Job2"));

  Assert.assertTrue(jobTwoStartTime - currentTime >= 5000L);
}
 

@Test
public void testWorkflowAndJobTaskUserContentStore() throws InterruptedException {
  String jobName = TestHelper.getTestMethodName();
  Workflow.Builder workflowBuilder = new Workflow.Builder(jobName);
  List<TaskConfig> taskConfigs = Lists.newArrayListWithCapacity(1);
  Map<String, String> taskConfigMap = Maps.newHashMap();
  TaskConfig taskConfig1 = new TaskConfig("ContentStoreTask", taskConfigMap);
  taskConfigs.add(taskConfig1);
  Map<String, String> jobCommandMap = Maps.newHashMap();
  jobCommandMap.put("Timeout", "1000");

  JobConfig.Builder jobBuilder = new JobConfig.Builder().setCommand("DummyCommand")
      .addTaskConfigs(taskConfigs).setWorkflow(jobName)
      .setJobCommandConfigMap(jobCommandMap);
  workflowBuilder.addJob(jobName, jobBuilder);

  _driver.start(workflowBuilder.build());
  _driver.pollForWorkflowState(jobName, TaskState.COMPLETED);
  Assert
      .assertEquals(_driver.getWorkflowContext(jobName).getWorkflowState(), TaskState.COMPLETED);
}
 

public static void submitJobToWorkFlow(JobConfig.Builder jobConfigBuilder,
    String workFlowName,
    String jobName,
    TaskDriver helixTaskDriver,
    HelixManager helixManager,
    long workFlowExpiryTime) throws Exception {

  WorkflowConfig workFlowConfig = new WorkflowConfig.Builder().setExpiry(workFlowExpiryTime, TimeUnit.SECONDS).build();
  // Create a work flow for each job with the name being the queue name
  Workflow workFlow = new Workflow.Builder(workFlowName).setWorkflowConfig(workFlowConfig).addJob(jobName, jobConfigBuilder).build();
  // start the workflow
  helixTaskDriver.start(workFlow);
  log.info("Created a work flow {}", workFlowName);

  waitJobInitialization(helixManager, workFlowName, jobName, Long.MAX_VALUE);
}
 

@Test
public void testThresholdFailure() throws Exception {
  // Create a job with two different tasks
  String jobName = TestHelper.getTestMethodName();
  Workflow.Builder workflowBuilder = new Workflow.Builder(jobName);
  List<TaskConfig> taskConfigs = Lists.newArrayListWithCapacity(2);
  Map<String, String> taskConfigMap = Maps.newHashMap(ImmutableMap.of("fail", "" + true));
  TaskConfig taskConfig1 = new TaskConfig("TaskOne", taskConfigMap);
  TaskConfig taskConfig2 = new TaskConfig("TaskTwo", null);
  taskConfigs.add(taskConfig1);
  taskConfigs.add(taskConfig2);
  Map<String, String> jobConfigMap = Maps.newHashMap();
  jobConfigMap.put("Timeout", "1000");
  JobConfig.Builder jobBuilder = new JobConfig.Builder().setCommand("DummyCommand")
      .setFailureThreshold(1).addTaskConfigs(taskConfigs).setJobCommandConfigMap(jobConfigMap);
  workflowBuilder.addJob(jobName, jobBuilder);
  _driver.start(workflowBuilder.build());

  // Ensure the job completes
  _driver.pollForWorkflowState(jobName, TaskState.IN_PROGRESS);
  _driver.pollForWorkflowState(jobName, TaskState.COMPLETED);

  // Ensure that each class was invoked
  Assert.assertTrue(_invokedClasses.contains(TaskOne.class.getName()));
  Assert.assertTrue(_invokedClasses.contains(TaskTwo.class.getName()));
}
 

@Test
public void testTaskRetryWithDelay() throws Exception {
  String jobResource = TestHelper.getTestMethodName();
  JobConfig.Builder jobBuilder = JobConfig.Builder.fromMap(WorkflowGenerator.DEFAULT_JOB_CONFIG);
  jobBuilder.setJobCommandConfigMap(WorkflowGenerator.DEFAULT_COMMAND_CONFIG)
      .setMaxAttemptsPerTask(2).setCommand(MockTask.TASK_COMMAND).setWorkflow(jobResource)
      .setFailureThreshold(Integer.MAX_VALUE).setTaskRetryDelay(2000L)
      .setJobCommandConfigMap(ImmutableMap.of(MockTask.FAILURE_COUNT_BEFORE_SUCCESS, "2"));
  Workflow flow =
      WorkflowGenerator.generateSingleJobWorkflowBuilder(jobResource, jobBuilder).build();
  _driver.start(flow);

  // Wait until the job completes.
  _driver.pollForWorkflowState(jobResource, TaskState.COMPLETED);

  long startTime = _driver.getWorkflowContext(jobResource).getStartTime();
  long finishedTime = _driver.getWorkflowContext(jobResource).getFinishTime();

  // It should finish and take more than 2 secs with the retry delay built in
  Assert.assertTrue(finishedTime - startTime >= 2000L);
}
 

@Test
public void testOneTimeScheduled() throws Exception {
  String jobName = TestHelper.getTestMethodName();
  Workflow.Builder workflowBuilder = new Workflow.Builder(jobName);
  List<TaskConfig> taskConfigs = Lists.newArrayListWithCapacity(1);
  Map<String, String> taskConfigMap = Maps.newHashMap();
  TaskConfig taskConfig1 = new TaskConfig("TaskOne", taskConfigMap);
  taskConfigs.add(taskConfig1);
  Map<String, String> jobCommandMap = Maps.newHashMap();
  jobCommandMap.put(MockTask.JOB_DELAY, "1000");

  JobConfig.Builder jobBuilder = new JobConfig.Builder().setCommand("DummyCommand")
      .addTaskConfigs(taskConfigs).setJobCommandConfigMap(jobCommandMap);
  workflowBuilder.addJob(jobName, jobBuilder);

  long inFiveSeconds = System.currentTimeMillis() + (5 * 1000);
  workflowBuilder.setScheduleConfig(ScheduleConfig.oneTimeDelayedStart(new Date(inFiveSeconds)));
  _driver.start(workflowBuilder.build());

  // Ensure the job completes
  _driver.pollForWorkflowState(jobName, TaskState.IN_PROGRESS);
  _driver.pollForWorkflowState(jobName, TaskState.COMPLETED);

  // Ensure that the class was invoked
  Assert.assertTrue(_invokedClasses.contains(TaskOne.class.getName()));

  // Check that the workflow only started after the start time (with a 1 second buffer)
  WorkflowContext workflowCtx = _driver.getWorkflowContext(jobName);
  long startTime = workflowCtx.getStartTime();
  Assert.assertTrue(startTime <= inFiveSeconds);
}
 

/**
 * Tests that stopping a workflow does result in its task ending up in STOPPED state.
 * @throws InterruptedException
 */
@Test(dependsOnMethods = "testStopWorkflow")
public void testStopTask() throws Exception {
  stopTestSetup(1);

  String workflowName = TestHelper.getTestMethodName();
  Workflow.Builder workflowBuilder = new Workflow.Builder(workflowName);
  WorkflowConfig.Builder configBuilder = new WorkflowConfig.Builder(workflowName);
  configBuilder.setAllowOverlapJobAssignment(true);
  workflowBuilder.setWorkflowConfig(configBuilder.build());

  for (int i = 0; i < 1; i++) {
    List<TaskConfig> taskConfigs = new ArrayList<>();
    taskConfigs.add(new TaskConfig("StopTask", new HashMap<>()));
    JobConfig.Builder jobConfigBulider = new JobConfig.Builder().setCommand("Dummy")
        .addTaskConfigs(taskConfigs).setJobCommandConfigMap(new HashMap<>());
    workflowBuilder.addJob("JOB" + i, jobConfigBulider);
  }

  _driver.start(workflowBuilder.build());
  _driver.pollForWorkflowState(workflowName, TaskState.IN_PROGRESS);

  // Stop the workflow
  _driver.stop(workflowName);
  _driver.pollForWorkflowState(workflowName, TaskState.STOPPED);

  Assert.assertEquals(TaskDriver.getWorkflowContext(_manager, workflowName).getWorkflowState(),
      TaskState.STOPPED);

  cleanupParticipants(1);
}
 

@Test(dependsOnMethods = "testJobFailsWithMissingTarget")
public void testJobContinueUponParentJobFailure() throws Exception {
  String workflowName = TestHelper.getTestMethodName();

  // Create a workflow
  LOG.info("Starting job-queue: " + workflowName);
  Workflow.Builder builder = new Workflow.Builder(workflowName).setWorkflowConfig(
      new WorkflowConfig.Builder(workflowName).setFailureThreshold(10).build());
  // Create and add jobs
  List<String> currentJobNames = new ArrayList<String>();
  for (int i = 0; i < _numDbs; i++) {
    JobConfig.Builder jobConfig =
        new JobConfig.Builder().setCommand(MockTask.TASK_COMMAND).setTargetResource(_testDbs.get(i))
            .setTargetPartitionStates(Sets.newHashSet("SLAVE")).setIgnoreDependentJobFailure(true);
    String jobName = "job" + _testDbs.get(i);
    builder.addJob(jobName, jobConfig);
    if (i > 0) {
      builder.addParentChildDependency("job" + _testDbs.get(i - 1), "job" + _testDbs.get(i));
    }
    currentJobNames.add(jobName);
  }

  _driver.start(builder.build());

  String namedSpaceJob1 = String.format("%s_%s", workflowName, currentJobNames.get(1));
  _driver.pollForJobState(workflowName, namedSpaceJob1, TaskState.FAILED);
  String lastJob =
      String.format("%s_%s", workflowName, currentJobNames.get(currentJobNames.size() - 1));
  _driver.pollForJobState(workflowName, lastJob, TaskState.COMPLETED);

  _driver.delete(workflowName);
}
 

@Test(dependsOnMethods = "testJobContinueUponParentJobFailure")
public void testJobFailsWithMissingTargetInRunning() throws Exception {
  String workflowName = TestHelper.getTestMethodName();

  // Create a workflow
  LOG.info("Starting job-queue: " + workflowName);
  Workflow.Builder builder = new Workflow.Builder(workflowName);
  // Create and add jobs
  List<String> currentJobNames = new ArrayList<String>();
  for (int i = 0; i < _numDbs; i++) {
    JobConfig.Builder jobConfig =
        new JobConfig.Builder().setCommand(MockTask.TASK_COMMAND).setTargetResource(_testDbs.get(i))
            .setTargetPartitionStates(Sets.newHashSet("SLAVE"));
    String jobName = "job" + _testDbs.get(i);
    builder.addJob(jobName, jobConfig);
    if (i > 0) {
      builder.addParentChildDependency("job" + _testDbs.get(i - 1), "job" + _testDbs.get(i));
    }      currentJobNames.add(jobName);
  }

  _driver.start(builder.build());
  _gSetupTool.dropResourceFromCluster(CLUSTER_NAME, _testDbs.get(0));

  String namedSpaceJob1 = String.format("%s_%s", workflowName, currentJobNames.get(0));
  _driver.pollForJobState(workflowName, namedSpaceJob1, TaskState.FAILED);
  _driver.pollForWorkflowState(workflowName, TaskState.FAILED);

  _driver.delete(workflowName);
}
 

@Test
public void testJobFailsWithMissingTarget() throws Exception {
  String workflowName = TestHelper.getTestMethodName();

  // Create a workflow
  LOG.info("Starting job-queue: " + workflowName);
  Workflow.Builder builder = new Workflow.Builder(workflowName);
  // Create and Enqueue jobs
  List<String> currentJobNames = new ArrayList<String>();
  for (int i = 0; i < _numDbs; i++) {
    JobConfig.Builder jobConfig =
        new JobConfig.Builder().setCommand(MockTask.TASK_COMMAND).setTargetResource(
            _testDbs.get(i))
            .setTargetPartitionStates(Sets.newHashSet("SLAVE"));
    String jobName = "job" + _testDbs.get(i);
    builder.addJob(jobName, jobConfig);
    if (i > 0) {
      builder.addParentChildDependency("job" + _testDbs.get(i - 1), "job" + _testDbs.get(i));
    }
    currentJobNames.add(jobName);
  }

  _gSetupTool.dropResourceFromCluster(CLUSTER_NAME, _testDbs.get(1));
  _driver.start(builder.build());

  String namedSpaceJob = String.format("%s_%s", workflowName, currentJobNames.get(1));
  _driver.pollForJobState(workflowName, namedSpaceJob, TaskState.FAILED);
  _driver.pollForWorkflowState(workflowName, TaskState.FAILED);

  _driver.delete(workflowName);
}
 

@Test
public void timeouts() throws Exception {
  final String jobResource = "timeouts";

  JobConfig.Builder jobBuilder = JobConfig.Builder.fromMap(WorkflowGenerator.DEFAULT_JOB_CONFIG);
  jobBuilder.setJobCommandConfigMap(WorkflowGenerator.DEFAULT_COMMAND_CONFIG)
      .setMaxAttemptsPerTask(2).setTimeoutPerTask(1); // This timeout needs to be very short

  Workflow flow =
      WorkflowGenerator.generateSingleJobWorkflowBuilder(jobResource, jobBuilder).build();
  _driver.start(flow);

  // Wait until the job reports failure.
  _driver.pollForWorkflowState(jobResource, TaskState.FAILED);

  // Check that all partitions timed out up to maxAttempts
  JobContext ctx = _driver.getJobContext(TaskUtil.getNamespacedJobName(jobResource));
  int maxAttempts = 0;
  boolean sawTimedoutTask = false;
  for (int i = 0; i < _numPartitions; i++) {
    TaskPartitionState state = ctx.getPartitionState(i);
    if (state != null) {
      if (state == TaskPartitionState.TIMED_OUT) {
        sawTimedoutTask = true;
      }
      // At least one task timed out, other might be aborted due to job failure.
      Assert.assertTrue(
          state == TaskPartitionState.TIMED_OUT || state == TaskPartitionState.TASK_ABORTED);
      maxAttempts = Math.max(maxAttempts, ctx.getPartitionNumAttempts(i));
    }
  }

  Assert.assertTrue(sawTimedoutTask);
  // 2 or 3 both okay only for tests - TODO: Fix this later
  Assert.assertTrue(maxAttempts == 2 || maxAttempts == 3);
}
 

@Test
public void testRepeatedWorkflow() throws Exception {
  String workflowName = "SomeWorkflow";
  Workflow flow =
      WorkflowGenerator.generateDefaultRepeatedJobWorkflowBuilder(workflowName).build();
  new TaskDriver(_manager).start(flow);

  // Wait until the workflow completes
  _driver.pollForWorkflowState(workflowName, TaskState.COMPLETED);

  // Assert completion for all tasks within two minutes
  for (String task : flow.getJobConfigs().keySet()) {
    _driver.pollForJobState(workflowName, task, TaskState.COMPLETED);
  }
}
 

@Test
public void testJobScheduleAfterResourceDisabled() throws InterruptedException {
  String workflowName = TestHelper.getTestMethodName();
  _gSetupTool.getClusterManagementTool()
      .enableResource(CLUSTER_NAME, WorkflowGenerator.DEFAULT_TGT_DB, false);
  Workflow.Builder workflow = new Workflow.Builder(workflowName);
  workflow.addJob(_jobName, _jobCfg);
  _driver.start(workflow.build());

  _driver.pollForWorkflowState(workflowName, TaskState.FAILED);
}
 

@Test
public void testTaskNumAttemptsWithDelay() throws Exception {
  int maxAttempts = Integer.MAX_VALUE; // Allow the count to increase infinitely
  // Use a delay that's long enough for multiple rounds of pipeline
  long retryDelay = 4000L;

  String workflowName = TestHelper.getTestMethodName();
  JobConfig.Builder jobBuilder = JobConfig.Builder.fromMap(WorkflowGenerator.DEFAULT_JOB_CONFIG);
  jobBuilder.setJobCommandConfigMap(WorkflowGenerator.DEFAULT_COMMAND_CONFIG)
      .setMaxAttemptsPerTask(maxAttempts).setCommand(MockTask.TASK_COMMAND)
      .setWorkflow(workflowName).setFailureThreshold(Integer.MAX_VALUE)
      .setTaskRetryDelay(retryDelay).setJobCommandConfigMap(
          ImmutableMap.of(MockTask.FAILURE_COUNT_BEFORE_SUCCESS, String.valueOf(maxAttempts)));
  Workflow flow =
      WorkflowGenerator.generateSingleJobWorkflowBuilder(workflowName, jobBuilder).build();
  _driver.start(flow);

  // Wait until the job is running
  _driver.pollForWorkflowState(workflowName, TaskState.IN_PROGRESS);

  JobContext jobContext =
      _driver.getJobContext(TaskUtil.getNamespacedJobName(workflowName, workflowName));
  int expectedAttempts = jobContext.getPartitionNumAttempts(0);

  // Check 3 times to see if maxAttempts match up
  for (int i = 0; i < 3; i++) {
    // Add a small delay to make sure the check falls in the middle of the scheduling timeline
    Thread.sleep(retryDelay + 1000L);
    JobContext ctx =
        _driver.getJobContext(TaskUtil.getNamespacedJobName(workflowName, workflowName));
    expectedAttempts++;
    Assert.assertEquals(ctx.getPartitionNumAttempts(0), expectedAttempts);
  }
}
 

@Test
public void testJobScheduleBeforeResourceDisabled() throws InterruptedException {
  String workflowName = TestHelper.getTestMethodName();
  Workflow.Builder workflow = new Workflow.Builder(workflowName);
  _jobCfg.setJobCommandConfigMap(ImmutableMap.of(MockTask.JOB_DELAY, "1000000"));
  workflow.addJob(_jobName, _jobCfg);
  _driver.start(workflow.build());
  Thread.sleep(1000);
  _gSetupTool.getClusterManagementTool()
      .enableResource(CLUSTER_NAME, WorkflowGenerator.DEFAULT_TGT_DB, false);
  _driver.pollForWorkflowState(workflowName, TaskState.FAILED);
}