package azkaban.execapp;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.fail;

import azkaban.event.Event;
import azkaban.executor.ExecutableFlow;
import azkaban.executor.ExecutableFlowBase;
import azkaban.executor.ExecutableNode;
import azkaban.executor.InteractiveTestJob;
import azkaban.executor.Status;
import java.util.function.Function;
import org.junit.Assert;

public class FlowRunnerTestBase {

  protected FlowRunner runner;
  protected EventCollectorListener eventCollector;

  public static boolean isStarted(final Status status) {
    if (status == Status.QUEUED) {
      return false;
    } else if (!Status.isStatusFinished(status) && !Status.isStatusRunning(status)) {
      return false;
    }
    return true;
  }

  public void assertThreadShutDown() {
    waitFlowRunner(
        runner -> Status.isStatusFinished(runner.getExecutableFlow().getStatus())
        && !runner.isRunnerThreadAlive());
  }

  public void assertThreadRunning() {
    waitFlowRunner(
        runner -> Status.isStatusRunning(runner.getExecutableFlow().getStatus())
        && runner.isRunnerThreadAlive());
  }

  public void waitFlowRunner(final Function<FlowRunner, Boolean> statusCheck) {
    for (int i = 0; i < 100; i++) {
      if (statusCheck.apply(this.runner)) {
        return;
      }
      synchronized (EventCollectorListener.handleEvent) {
        try {
          EventCollectorListener.handleEvent.wait(10L);
        } catch (final InterruptedException e) {
        }
      }
    }
    Assert.fail("Flow didn't reach expected status");
  }

  public void waitJobStatuses(final Function<Status, Boolean> statusCheck,
      final String... jobs) {
    for (int i = 0; i < 100; i++) {
      if (checkJobStatuses(statusCheck, jobs)) {
        return;
      }
      synchronized (EventCollectorListener.handleEvent) {
        try {
          EventCollectorListener.handleEvent.wait(10L);
        } catch (final InterruptedException e) {
        }
      }
    }
    Assert.fail("Jobs didn't reach expected statuses");
  }

  public void waitJobsStarted(final FlowRunner runner, final String... jobs) {
    waitJobStatuses(FlowRunnerTest::isStarted, jobs);
  }

  protected void waitEventFired(final String nestedId, final Status status)
      throws InterruptedException {
    for (int i = 0; i < 100; i++) {
      for (final Event event : this.eventCollector.getEventList()) {
        if (event.getData().getStatus() == status && event.getData().getNestedId()
            .equals(nestedId)) {
          return;
        }
      }
      synchronized (EventCollectorListener.handleEvent) {
        EventCollectorListener.handleEvent.wait(10L);
      }
    }
    fail("Event wasn't fired with [" + nestedId + "], " + status);
  }

  public boolean checkJobStatuses(final Function<Status, Boolean> statusCheck,
      final String[] jobs) {
    final ExecutableFlow exFlow = this.runner.getExecutableFlow();
    for (final String name : jobs) {
      final ExecutableNode node = exFlow.getExecutableNodePath(name);
      assertNotNull(name + " wasn't found", node);
      if (!statusCheck.apply(node.getStatus())) {
        return false;
      }
    }
    return true;
  }

  protected void assertFlowStatus(final Status status) {
    final ExecutableFlow flow = this.runner.getExecutableFlow();
    for (int i = 0; i < 100; i++) {
      if (flow.getStatus() == status) {
        break;
      }
      synchronized (EventCollectorListener.handleEvent) {
        try {
          EventCollectorListener.handleEvent.wait(10L);
        } catch (final InterruptedException e) {
        }
      }
    }
    printStatuses(status, flow);
    assertEquals(status, flow.getStatus());
  }

  protected void assertStatus(final String name, final Status status) {
    final ExecutableFlow exFlow = this.runner.getExecutableFlow();
    final ExecutableNode node = exFlow.getExecutableNodePath(name);
    assertNotNull(name + " wasn't found", node);
    for (int i = 0; i < 100; i++) {
      if (node.getStatus() == status) {
        break;
      }
      synchronized (EventCollectorListener.handleEvent) {
        try {
          EventCollectorListener.handleEvent.wait(10L);
        } catch (final InterruptedException e) {
        }
      }
    }
    printStatuses(status, node);
    assertEquals("Wrong status for [" + name + "]", status, node.getStatus());
  }

  protected void printStatuses(final Status status, final ExecutableNode node) {
    if (status != node.getStatus()) {
      printTestJobs();
      printFlowJobs(this.runner.getExecutableFlow());
    }
  }

  private void printTestJobs() {
    for (final String testJob : InteractiveTestJob.testJobs.keySet()) {
      final ExecutableNode testNode = this.runner.getExecutableFlow()
          .getExecutableNodePath(testJob);
      System.err.println("testJob: " + testNode.getNestedId() + " " + testNode.getStatus());
    }
  }

  private void printFlowJobs(final ExecutableFlowBase flow) {
    System.err.println("ExecutableFlow: " + flow.getNestedId() + " " + flow.getStatus());
    for (final ExecutableNode node : flow.getExecutableNodes()) {
      if (node instanceof ExecutableFlowBase) {
        printFlowJobs((ExecutableFlowBase) node);
      } else {
        System.err.println("ExecutableNode: " + node.getNestedId() + " " + node.getStatus());
      }
    }
  }

  protected void succeedJobs(final String... jobs) {
    waitJobsStarted(this.runner, jobs);
    for (final String name : jobs) {
      InteractiveTestJob.getTestJob(name).succeedJob();
    }
  }

}

@@ -677,19 +677,20 @@
    %7D
- else if (!
+%0A    return 
 Stat
@@ -717,20 +717,19 @@
 status) 
-&& !
+%7C%7C 
 Status.i
@@ -754,53 +754,8 @@
 tus)
-) %7B%0A      return false;%0A    %7D%0A    return true
 ;%0A  
@@ -3328,77 +3328,35 @@
     
-for (int i = 0; i %3C 100; i++) %7B%0A      if (flow.get
+waitFor
 Status(
-) ==
+flow,
  status)
  %7B%0A 
@@ -3355,234 +3355,9 @@
 tus)
- %7B%0A        break;%0A      %7D%0A      synchronized (EventCollectorListener.handleEvent) %7B%0A        try %7B%0A          EventCollectorListener.handleEvent.wait(10L);%0A        %7D catch (final InterruptedException e) %7B%0A        %7D%0A      %7D%0A    %7D
+;
 %0A   
@@ -3687,24 +3687,241 @@
 nd%22, node);%0A
+    waitForStatus(node, status);%0A    printStatuses(status, node);%0A    assertEquals(%22Wrong status for %5B%22 + name + %22%5D%22, status, node.getStatus());%0A  %7D%0A%0A  private void waitForStatus(ExecutableNode node, Status status) %7B%0A
     for (int
@@ -4176,32 +4176,47 @@
 dException e) %7B%0A
+          i--;%0A
         %7D%0A      
@@ -4221,126 +4221,14 @@
   %7D%0A
+
     %7D%0A
-    printStatuses(status, node);%0A    assertEquals(%22Wrong status for %5B%22 + name + %22%5D%22, status, node.getStatus());%0A
   %7D%0A

Test cleanup: FlowRunnerTestBase (#1154)