Java Code Examples for org.apache.helix.task.Workflow#Builder

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

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

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

@Test(dataProvider = "testJobFailureInput")
public void testNormalJobFailure(String comment, List<String> taskStates,
    List<String> expectedTaskEndingStates, String expectedJobEndingStates,
    String expectedWorkflowEndingStates) throws Exception {
  final String JOB_NAME = "test_job";
  final String WORKFLOW_NAME = TestHelper.getTestMethodName() + testNum++;
  System.out.println("Test case comment: " + comment);

  Map<String, Map<String, String>> targetPartitionConfigs =
      createPartitionConfig(taskStates, expectedTaskEndingStates);

  JobConfig.Builder firstJobBuilder =
      new JobConfig.Builder().setWorkflow(WORKFLOW_NAME).setTargetResource(DB_NAME)
          .setTargetPartitionStates(Sets.newHashSet(MasterSlaveSMD.States.MASTER.name()))
          .setCommand(MockTask.TASK_COMMAND)
          .setJobCommandConfigMap(ImmutableMap.of(MockTask.TARGET_PARTITION_CONFIG,
              MockTask.serializeTargetPartitionConfig(targetPartitionConfigs)));

  Workflow.Builder workflowBuilder =
      new Workflow.Builder(WORKFLOW_NAME).addJob(JOB_NAME, firstJobBuilder);

  _driver.start(workflowBuilder.build());
  _driver.pollForJobState(WORKFLOW_NAME, TaskUtil.getNamespacedJobName(WORKFLOW_NAME, JOB_NAME),
      TaskState.valueOf(expectedJobEndingStates));
  _driver.pollForWorkflowState(WORKFLOW_NAME, TaskState.valueOf(expectedWorkflowEndingStates));

  JobContext jobContext =
      _driver.getJobContext(TaskUtil.getNamespacedJobName(WORKFLOW_NAME, JOB_NAME));
  for (int pId : jobContext.getPartitionSet()) {
    Map<String, String> targetPartitionConfig =
        targetPartitionConfigs.get(jobContext.getTargetForPartition(pId));
    Assert
        .assertTrue(Arrays.asList(targetPartitionConfig.get(EXPECTED_ENDING_STATE).split("\\s+"))
            .contains(jobContext.getPartitionState(pId).name()));
  }
}
 

@Test
public void testJobSubmitGenericWorkflows() throws InterruptedException {
  String workflowName = TestHelper.getTestMethodName();
  JobConfig.Builder jobBuilder =
      new JobConfig.Builder().setTargetResource(WorkflowGenerator.DEFAULT_TGT_DB)
          .setCommand(MockTask.TASK_COMMAND).setMaxAttemptsPerTask(2);
  Workflow.Builder builder = new Workflow.Builder(workflowName);
  for (int i = 0; i < 5; i++) {
    builder.addJob("JOB" + i, jobBuilder);
  }

  /**
   * Dependency visualization
   *               JOB0
   *
   *             /   |    \
   *
   *         JOB1 <-JOB2   JOB4
   *
   *                 |     /
   *
   *                JOB3
   */

  builder.addParentChildDependency("JOB0", "JOB1");
  builder.addParentChildDependency("JOB0", "JOB2");
  builder.addParentChildDependency("JOB0", "JOB4");
  builder.addParentChildDependency("JOB1", "JOB2");
  builder.addParentChildDependency("JOB2", "JOB3");
  builder.addParentChildDependency("JOB4", "JOB3");
  _driver.start(builder.build());

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

@Test
public void testWorkflowRunningTime() throws InterruptedException {
  String workflowName = TestHelper.getTestMethodName();
  _jobBuilder.setWorkflow(workflowName);
  Workflow.Builder workflowBuilder = new Workflow.Builder(workflowName)
      .setWorkflowConfig(new WorkflowConfig.Builder(workflowName).setTimeout(1000).build())
      .addJob(JOB_NAME, _jobBuilder);
  _driver.start(workflowBuilder.build());

  _driver.pollForWorkflowState(workflowName, 10000L, TaskState.TIMED_OUT);
}
 

@Test
public void testEnableDisablePartitions() throws InterruptedException {
  _admin.enablePartition(false, CLUSTER_NAME, (PARTICIPANT_PREFIX + "_" + _startPort),
      WorkflowGenerator.DEFAULT_TGT_DB, Arrays.asList(new String[] { "TestDB_0", "TestDB_2" }));
  _admin.enablePartition(false, CLUSTER_NAME, (PARTICIPANT_PREFIX + "_" + (_startPort + 1)),
      WorkflowGenerator.DEFAULT_TGT_DB, Arrays.asList(new String[] { "TestDB_0", "TestDB_2" }));

  IdealState idealState =
      _admin.getResourceIdealState(CLUSTER_NAME, WorkflowGenerator.DEFAULT_TGT_DB);
  List<String> preferenceList =
      Arrays.asList(new String[] { "localhost_12919", "localhost_12918" });
  for (String partitionName : idealState.getPartitionSet()) {
    idealState.setPreferenceList(partitionName, preferenceList);
  }
  idealState.setRebalanceMode(IdealState.RebalanceMode.SEMI_AUTO);
  _admin.setResourceIdealState(CLUSTER_NAME, WorkflowGenerator.DEFAULT_TGT_DB, idealState);

  Assert.assertTrue(_clusterVerifier.verifyByPolling());

  String workflowName = TestHelper.getTestMethodName();
  Workflow.Builder builder = new Workflow.Builder(workflowName);
  JobConfig.Builder jobBuilder =
      new JobConfig.Builder().setWorkflow(workflowName).setCommand(MockTask.TASK_COMMAND)
          .setTargetResource(WorkflowGenerator.DEFAULT_TGT_DB)
          .setTargetPartitionStates(Collections.singleton("SLAVE"));
  builder.addJob("JOB", jobBuilder);
  _driver.start(builder.build());
  Thread.sleep(2000L);
  JobContext jobContext =
      _driver.getJobContext(TaskUtil.getNamespacedJobName(workflowName, "JOB"));
  int n = idealState.getNumPartitions();
  for ( int i = 0; i < n; i++) {
    String targetPartition = jobContext.getTargetForPartition(i);
    if (targetPartition.equals("TestDB_0") || targetPartition.equals("TestDB_2")) {
      Assert.assertEquals(jobContext.getPartitionState(i), null);
    } else {
      Assert.assertEquals(jobContext.getPartitionState(i), TaskPartitionState.COMPLETED);
    }
  }
}
 

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

@Test public void testJobPoll() throws InterruptedException {
  String jobResource = TestHelper.getTestMethodName();
  Workflow.Builder builder =
      WorkflowGenerator.generateDefaultSingleJobWorkflowBuilder(jobResource);
  _driver.start(builder.build());

  TaskState polledState = _driver
      .pollForJobState(jobResource, String.format("%s_%s", jobResource, jobResource), 4000L,
          TaskState.COMPLETED, TaskState.FAILED);
  Assert.assertEquals(TaskState.COMPLETED, polledState);
}
 

protected Map<String, Workflow> createWorkflows(String cluster, int numWorkflows) {
  Map<String, Workflow> workflows = new HashMap<>();
  HelixPropertyStore<ZNRecord> propertyStore =
      new ZkHelixPropertyStore<>((ZkBaseDataAccessor<ZNRecord>) _baseAccessor,
          PropertyPathBuilder.propertyStore(cluster), null);

  for (int i = 0; i < numWorkflows; i++) {
    Workflow.Builder workflow = new Workflow.Builder(WORKFLOW_PREFIX + i);
    int j = 0;
    for (JobConfig.Builder job : createJobs(cluster, WORKFLOW_PREFIX + i, 3)) {
      workflow.addJob(JOB_PREFIX + j++, job);
    }
    workflows.put(WORKFLOW_PREFIX + i, workflow.build());
    WorkflowContext workflowContext = TaskTestUtil
        .buildWorkflowContext(WORKFLOW_PREFIX + i, TaskState.IN_PROGRESS,
            System.currentTimeMillis(), TaskState.COMPLETED, TaskState.COMPLETED,
            TaskState.IN_PROGRESS);
    _baseAccessor.set(String.format("/%s/%s%s/%s/%s", cluster, PropertyType.PROPERTYSTORE.name(),
        TaskConstants.REBALANCER_CONTEXT_ROOT, WORKFLOW_PREFIX + i, TaskConstants.CONTEXT_NODE),
        workflowContext.getRecord(), AccessOption.PERSISTENT);
    _configAccessor.setResourceConfig(cluster, WORKFLOW_PREFIX + i, workflow.getWorkflowConfig());

    // Add workflow user content
    propertyStore.create(Joiner.on("/")
            .join(TaskConstants.REBALANCER_CONTEXT_ROOT, WORKFLOW_PREFIX + i,
                TaskUtil.USER_CONTENT_NODE), new ZNRecord(TaskUtil.USER_CONTENT_NODE),
        AccessOption.PERSISTENT);
  }
  return workflows;
}
 

@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 testDeleteCompletedWorkflowForcefully() throws Exception {
  // Create a simple workflow and wait for its completion. Then delete the IdealState,
  // WorkflowContext and WorkflowConfig using ForceDelete.
  String workflowName = TestHelper.getTestMethodName();
  Workflow.Builder builder = createCustomWorkflow(workflowName, SHORT_EXECUTION_TIME, "0");
  _driver.start(builder.build());

  // Wait until workflow is created and completed.
  _driver.pollForWorkflowState(workflowName, TaskState.COMPLETED);

  // Check that WorkflowConfig, WorkflowContext, and IdealState are indeed created for this
  // workflow
  Assert.assertNotNull(_driver.getWorkflowConfig(workflowName));
  Assert.assertNotNull(_driver.getWorkflowContext(workflowName));
  Assert.assertNotNull(_admin.getResourceIdealState(CLUSTER_NAME, workflowName));

  // Stop the Controller
  _controller.syncStop();

  _driver.delete(workflowName, true);

  // Check that WorkflowConfig, WorkflowContext, and IdealState are indeed deleted for this
  // workflow.
  boolean isWorkflowDeleted = TestHelper.verify(() -> {
    WorkflowConfig wfcfg = _driver.getWorkflowConfig(workflowName);
    WorkflowContext wfctx = _driver.getWorkflowContext(workflowName);
    IdealState is = _admin.getResourceIdealState(CLUSTER_NAME, workflowName);
    return (wfcfg == null && wfctx == null && is == null);
  }, 60 * 1000);
  Assert.assertTrue(isWorkflowDeleted);
}
 

@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 upon Participant reconnect, the Controller correctly sends the current
 * state of the task partition to DROPPED. This is to avoid a resource leak in case of a
 * Participant disconnect/reconnect and an ensuing reset() on all of the partitions on that
 * Participant.
 */
@Test
public void testDropOnParticipantReset() throws InterruptedException {
  // Create a workflow with some long-running jobs in progress
  String workflowName = TestHelper.getTestMethodName();
  String jobName = "JOB";
  Workflow.Builder builder = new Workflow.Builder(workflowName);
  List<TaskConfig> taskConfigs = new ArrayList<>();
  for (int j = 0; j < 2; j++) { // 2 tasks to ensure that they execute
    String taskID = jobName + "_TASK_" + j;
    TaskConfig.Builder taskConfigBuilder = new TaskConfig.Builder();
    taskConfigBuilder.setTaskId(taskID).setCommand(MockTask.TASK_COMMAND)
        .addConfig(MockTask.JOB_DELAY, "3000");
    taskConfigs.add(taskConfigBuilder.build());
  }
  JobConfig.Builder jobBuilder = new JobConfig.Builder().setCommand(MockTask.TASK_COMMAND)
      .setMaxAttemptsPerTask(10).setJobCommandConfigMap(WorkflowGenerator.DEFAULT_COMMAND_CONFIG)
      .addTaskConfigs(taskConfigs).setIgnoreDependentJobFailure(true)
      // 1 task at a time
      .setNumConcurrentTasksPerInstance(1);
  builder.addJob(jobName, jobBuilder);

  // Modify maxConcurrentTask for the instance so that it only accepts 1 task at most
  InstanceConfig instanceConfig = _gSetupTool.getClusterManagementTool()
      .getInstanceConfig(CLUSTER_NAME, _participants[0].getInstanceName());
  instanceConfig.setMaxConcurrentTask(1);
  _gSetupTool.getClusterManagementTool().setInstanceConfig(CLUSTER_NAME,
      _participants[0].getInstanceName(), instanceConfig);

  // Start the workflow
  _driver.start(builder.build());
  _driver.pollForJobState(workflowName, workflowName + "_" + jobName, TaskState.IN_PROGRESS);
  Thread.sleep(1500L); // Wait for the Participant to process the message
  // Stop and start the participant to mimic a connection issue
  _participants[0].syncStop();
  // Upon starting the participant, the first task partition should be dropped and assigned anew
  // on the instance. Then the rest of the tasks will execute and the workflow will complete
  startParticipant(0);

  TaskState workflowState = _driver.pollForWorkflowState(workflowName, TaskState.COMPLETED);
  TaskState jobState =
      _driver.pollForJobState(workflowName, workflowName + "_" + jobName, TaskState.COMPLETED);
  Assert.assertEquals(workflowState, TaskState.COMPLETED);
  Assert.assertEquals(jobState, TaskState.COMPLETED);
}
 

@Test(dependsOnMethods = "testDeleteRunningWorkflowForcefully")
public void testDeleteStoppedWorkflowForcefully() throws Exception {
  // Create a simple workflow. Stop the workflow and wait until it's fully stopped. Then delete
  // the IdealState, WorkflowContext and WorkflowConfig using ForceDelete.

  // Start the Controller
  String controllerName = CONTROLLER_PREFIX + "_0";
  _controller = new ClusterControllerManager(ZK_ADDR, CLUSTER_NAME, controllerName);
  _controller.syncStart();

  String workflowName = TestHelper.getTestMethodName();
  Workflow.Builder builder = createCustomWorkflow(workflowName, MEDIUM_EXECUTION_TIME, "0");
  _driver.start(builder.build());

  // Wait until workflow is created and running.
  _driver.pollForWorkflowState(workflowName, TaskState.IN_PROGRESS);

  _driver.stop(workflowName);

  // Wait until workflow is stopped.
  _driver.pollForWorkflowState(workflowName, TaskState.STOPPED);

  // Wait until workflow is stopped. Also, jobs should be either stopped or not created (null).
  boolean areJobsStopped = TestHelper.verify(() -> {
    WorkflowContext wCtx1 = _driver.getWorkflowContext(workflowName);
    TaskState job0 = wCtx1.getJobState(TaskUtil.getNamespacedJobName(workflowName, "JOB0"));
    TaskState job1 = wCtx1.getJobState(TaskUtil.getNamespacedJobName(workflowName, "JOB1"));
    TaskState job2 = wCtx1.getJobState(TaskUtil.getNamespacedJobName(workflowName, "JOB2"));
    return ((job0 == null || job0 == TaskState.STOPPED)
        && (job1 == null || job1 == TaskState.STOPPED)
        && (job2 == null || job2 == TaskState.STOPPED));
  }, 60 * 1000);
  Assert.assertTrue(areJobsStopped);

  // Check that WorkflowConfig, WorkflowContext, and IdealState are indeed created for this
  // workflow.
  Assert.assertNotNull(_driver.getWorkflowConfig(workflowName));
  Assert.assertNotNull(_driver.getWorkflowContext(workflowName));
  Assert.assertNotNull(_admin.getResourceIdealState(CLUSTER_NAME, workflowName));

  // Stop the Controller
  _controller.syncStop();

  _driver.delete(workflowName, true);

  // Check that WorkflowConfig, WorkflowContext, and IdealState are indeed deleted for this
  // workflow.
  boolean isWorkflowDeleted = TestHelper.verify(() -> {
    WorkflowConfig wfcfg = _driver.getWorkflowConfig(workflowName);
    WorkflowContext wfctx = _driver.getWorkflowContext(workflowName);
    IdealState is = _admin.getResourceIdealState(CLUSTER_NAME, workflowName);
    return (wfcfg == null && wfctx == null && is == null);
  }, 60 * 1000);
  Assert.assertTrue(isWorkflowDeleted);
}
 

@Test(dependsOnMethods = "testDeleteCompletedWorkflowForcefully")
public void testDeleteRunningWorkflowForcefully() throws Exception {
  // Create a simple workflow and wait until it reaches the Running state. Then delete the
  // IdealState, WorkflowContext and WorkflowConfig using ForceDelete.

  // Start the Controller
  String controllerName = CONTROLLER_PREFIX + "_0";
  _controller = new ClusterControllerManager(ZK_ADDR, CLUSTER_NAME, controllerName);
  _controller.syncStart();

  String workflowName = TestHelper.getTestMethodName();
  Workflow.Builder builder = createCustomWorkflow(workflowName, LONG_EXECUTION_TIME, "0");
  _driver.start(builder.build());

  // Wait until workflow is created and running.
  _driver.pollForWorkflowState(workflowName, TaskState.IN_PROGRESS);

  // Check the status of JOB0 to make sure it is running.
  _driver.pollForJobState(workflowName, TaskUtil.getNamespacedJobName(workflowName, "JOB0"),
      TaskState.IN_PROGRESS);

  // Check that WorkflowConfig, WorkflowContext, and IdealState are indeed created for this
  // workflow
  Assert.assertNotNull(_driver.getWorkflowConfig(workflowName));
  Assert.assertNotNull(_driver.getWorkflowContext(workflowName));
  Assert.assertNotNull(_admin.getResourceIdealState(CLUSTER_NAME, workflowName));

  // Stop the Controller
  _controller.syncStop();

  _driver.delete(workflowName, true);

  // Check that WorkflowConfig, WorkflowContext, and IdealState are indeed deleted for this
  // workflow
  boolean isWorkflowDeleted = TestHelper.verify(() -> {
    WorkflowConfig wfcfg = _driver.getWorkflowConfig(workflowName);
    WorkflowContext wfctx = _driver.getWorkflowContext(workflowName);
    IdealState is = _admin.getResourceIdealState(CLUSTER_NAME, workflowName);
    return (wfcfg == null && wfctx == null && is == null);
  }, 60 * 1000);
  Assert.assertTrue(isWorkflowDeleted);
}
 

@Test
public void testWorkflowJobFail() throws Exception {
  String workflowName = TestHelper.getTestMethodName();
  String job1 = JOB_NAME + "1";
  String job2 = JOB_NAME + "2";
  String job3 = JOB_NAME + "3";
  String job4 = JOB_NAME + "4";
  long workflowExpiry = 10000;
  long timeout = 10000;

  JobConfig.Builder jobBuilder = createJobConfigBuilder(workflowName, false, 1);
  JobConfig.Builder failedJobBuilder = createJobConfigBuilder(workflowName, true, 1);

  Workflow.Builder workflowBuilder = new Workflow.Builder(workflowName)
      .setWorkflowConfig(new WorkflowConfig.Builder(workflowName).setWorkFlowType(WORKFLOW_TYPE)
          .setTimeout(timeout).setParallelJobs(4).setFailureThreshold(1).build())
      .addJob(job1, jobBuilder).addJob(job2, jobBuilder).addJob(job3, failedJobBuilder)
      .addJob(job4, jobBuilder).addParentChildDependency(job1, job2)
      .addParentChildDependency(job1, job3).addParentChildDependency(job2, job4)
      .addParentChildDependency(job3, job4).setExpiry(workflowExpiry);

  _driver.start(workflowBuilder.build());

  _driver.pollForWorkflowState(workflowName, 10000L, TaskState.FAILED);

  // Timeout is longer than fail time, so the failover should occur earlier
  WorkflowContext context = _driver.getWorkflowContext(workflowName);
  long finishTime = context.getFinishTime();
  Assert.assertTrue(context.getFinishTime() - context.getStartTime() < timeout);

  // job1 will complete
  _driver.pollForJobState(workflowName, getJobNameToPoll(workflowName, job1), 10000L,
      TaskState.COMPLETED);

  // Possible race between 2 and 3 so it's likely for job2 to stay in either COMPLETED or ABORTED
  _driver.pollForJobState(workflowName, getJobNameToPoll(workflowName, job2), 10000L,
      TaskState.COMPLETED, TaskState.ABORTED);

  // job3 meant to fail
  _driver.pollForJobState(workflowName, getJobNameToPoll(workflowName, job3), 10000L,
      TaskState.FAILED);

  // because job4 has dependency over job3, it will fail as well
  _driver.pollForJobState(workflowName, getJobNameToPoll(workflowName, job4), 10000L,
      TaskState.FAILED);

  // Check MBean is updated
  ObjectName objectName = getWorkflowMBeanObjectName(workflowName);
  Assert.assertEquals((long) beanServer.getAttribute(objectName, "FailedWorkflowCount"), 1);

  // For a failed workflow, after timing out, it will be purged
  verifyWorkflowCleanup(workflowName, getJobNameToPoll(workflowName, job1),
      getJobNameToPoll(workflowName, job2), getJobNameToPoll(workflowName, job3),
      getJobNameToPoll(workflowName, job4));

  long cleanUpTime = System.currentTimeMillis();
  Assert.assertTrue(cleanUpTime - finishTime >= workflowExpiry);
}