@BeforeClass
 public void setup() throws Exception {
   super.setup();
   createDB(DB_SQL_FILE);
   falconJPAService.init();
   this.dfsCluster = EmbeddedCluster.newCluster("testCluster");
   this.conf = dfsCluster.getConf();
   registerServices();
 }
/** Test cases for JDBCStateStore. */
public class TestJDBCStateStore extends AbstractSchedulerTestBase {
  private static StateStore stateStore = JDBCStateStore.get();
  private static Random randomValGenerator = new Random();
  private static FalconJPAService falconJPAService = FalconJPAService.get();
  private AlarmService mockTimeService;
  private DataAvailabilityService mockDataService;
  private SchedulerService mockSchedulerService;
  private JobCompletionService mockCompletionService;
  private DAGEngine dagEngine;

  @BeforeClass
  public void setup() throws Exception {
    super.setup();
    createDB(DB_SQL_FILE);
    falconJPAService.init();
    this.dfsCluster = EmbeddedCluster.newCluster("testCluster");
    this.conf = dfsCluster.getConf();
    registerServices();
  }

  private void registerServices() throws FalconException {
    mockTimeService = Mockito.mock(AlarmService.class);
    Mockito.when(mockTimeService.getName()).thenReturn("AlarmService");
    Mockito.when(
            mockTimeService.createRequestBuilder(
                Mockito.any(NotificationHandler.class), Mockito.any(ID.class)))
        .thenCallRealMethod();
    mockDataService = Mockito.mock(DataAvailabilityService.class);
    Mockito.when(mockDataService.getName()).thenReturn("DataAvailabilityService");
    Mockito.when(
            mockDataService.createRequestBuilder(
                Mockito.any(NotificationHandler.class), Mockito.any(ID.class)))
        .thenCallRealMethod();
    dagEngine = Mockito.mock(OozieDAGEngine.class);
    Mockito.doNothing().when(dagEngine).resume(Mockito.any(ExecutionInstance.class));
    mockSchedulerService = Mockito.mock(SchedulerService.class);
    Mockito.when(mockSchedulerService.getName()).thenReturn("JobSchedulerService");
    StartupProperties.get().setProperty("dag.engine.impl", MockDAGEngine.class.getName());
    StartupProperties.get()
        .setProperty("execution.service.impl", FalconExecutionService.class.getName());
    dagEngine = Mockito.spy(DAGEngineFactory.getDAGEngine("testCluster"));
    Mockito.when(
            mockSchedulerService.createRequestBuilder(
                Mockito.any(NotificationHandler.class), Mockito.any(ID.class)))
        .thenCallRealMethod();
    mockCompletionService = Mockito.mock(JobCompletionService.class);
    Mockito.when(mockCompletionService.getName()).thenReturn("JobCompletionService");
    Mockito.when(
            mockCompletionService.createRequestBuilder(
                Mockito.any(NotificationHandler.class), Mockito.any(ID.class)))
        .thenCallRealMethod();
    Services.get().register(mockTimeService);
    Services.get().register(mockDataService);
    Services.get().register(mockSchedulerService);
    Services.get().register(mockCompletionService);
  }

  @Test
  public void testInsertRetrieveAndUpdate() throws Exception {
    EntityState entityState = getEntityState(EntityType.PROCESS, "process");
    stateStore.putEntity(entityState);
    EntityID entityID = new EntityID(entityState.getEntity());
    EntityState actualEntityState = stateStore.getEntity(entityID);
    Assert.assertEquals(actualEntityState.getEntity(), entityState.getEntity());
    Assert.assertEquals(actualEntityState.getCurrentState(), entityState.getCurrentState());
    try {
      stateStore.putEntity(entityState);
      Assert.fail("Exception must have been thrown");
    } catch (StateStoreException e) {
      // no op
    }

    entityState.setCurrentState(EntityState.STATE.SCHEDULED);
    stateStore.updateEntity(entityState);
    actualEntityState = stateStore.getEntity(entityID);
    Assert.assertEquals(actualEntityState.getEntity(), entityState.getEntity());
    Assert.assertEquals(actualEntityState.getCurrentState(), entityState.getCurrentState());

    stateStore.deleteEntity(entityID);
    boolean entityExists = stateStore.entityExists(entityID);
    Assert.assertEquals(entityExists, false);

    try {
      stateStore.getEntity(entityID);
      Assert.fail("Exception must have been thrown");
    } catch (StateStoreException e) {
      // no op
    }

    try {
      stateStore.updateEntity(entityState);
      Assert.fail("Exception must have been thrown");
    } catch (StateStoreException e) {
      // no op
    }

    try {
      stateStore.deleteEntity(entityID);
      Assert.fail("Exception must have been thrown");
    } catch (StateStoreException e) {
      // no op
    }
  }

  @Test
  public void testGetEntities() throws Exception {
    EntityState entityState1 = getEntityState(EntityType.PROCESS, "process1");
    EntityState entityState2 = getEntityState(EntityType.PROCESS, "process2");
    EntityState entityState3 = getEntityState(EntityType.FEED, "feed1");

    Collection<EntityState> result = stateStore.getAllEntities();
    Assert.assertEquals(result.size(), 0);

    stateStore.putEntity(entityState1);
    stateStore.putEntity(entityState2);
    stateStore.putEntity(entityState3);

    result = stateStore.getAllEntities();
    Assert.assertEquals(result.size(), 3);

    Collection<Entity> entities = stateStore.getEntities(EntityState.STATE.SUBMITTED);
    Assert.assertEquals(entities.size(), 3);
  }

  @Test
  public void testInstanceInsertionAndUpdate() throws Exception {
    storeEntity(EntityType.CLUSTER, "testCluster");
    storeEntity(EntityType.FEED, "clicksFeed");
    storeEntity(EntityType.FEED, "clicksSummary");
    EntityState entityState = getEntityState(EntityType.PROCESS, "process");
    ExecutionInstance executionInstance =
        BeanMapperUtil.getExecutionInstance(
            entityState.getEntity().getEntityType(),
            entityState.getEntity(),
            System.currentTimeMillis(),
            "cluster",
            System.currentTimeMillis());
    InstanceState instanceState = new InstanceState(executionInstance);
    initInstanceState(instanceState);
    stateStore.putExecutionInstance(instanceState);
    InstanceID instanceID = new InstanceID(instanceState.getInstance());
    InstanceState actualInstanceState = stateStore.getExecutionInstance(instanceID);
    Assert.assertEquals(actualInstanceState, instanceState);

    instanceState.setCurrentState(InstanceState.STATE.RUNNING);
    Predicate predicate = new Predicate(Predicate.TYPE.DATA);
    instanceState.getInstance().getAwaitingPredicates().add(predicate);

    stateStore.updateExecutionInstance(instanceState);
    actualInstanceState = stateStore.getExecutionInstance(instanceID);
    Assert.assertEquals(actualInstanceState, instanceState);

    try {
      stateStore.putExecutionInstance(instanceState);
      Assert.fail("Exception must have been thrown");
    } catch (StateStoreException e) {
      // no op
    }

    stateStore.deleteExecutionInstance(instanceID);

    try {
      stateStore.getExecutionInstance(instanceID);
      Assert.fail("Exception must have been thrown");
    } catch (StateStoreException e) {
      // no op
    }

    try {
      stateStore.deleteExecutionInstance(instanceID);
      Assert.fail("Exception must have been thrown");
    } catch (StateStoreException e) {
      // no op
    }

    try {
      stateStore.updateExecutionInstance(instanceState);
      Assert.fail("Exception must have been thrown");
    } catch (StateStoreException e) {
      // no op
    }
  }

  @Test
  public void testBulkInstanceOperations() throws Exception {
    storeEntity(EntityType.CLUSTER, "testCluster");
    storeEntity(EntityType.FEED, "clicksFeed");
    storeEntity(EntityType.FEED, "clicksSummary");
    EntityState entityState = getEntityState(EntityType.PROCESS, "process1");
    ExecutionInstance processExecutionInstance1 =
        BeanMapperUtil.getExecutionInstance(
            entityState.getEntity().getEntityType(),
            entityState.getEntity(),
            System.currentTimeMillis() - 60000,
            "cluster1",
            System.currentTimeMillis() - 60000);
    InstanceState instanceState1 = new InstanceState(processExecutionInstance1);
    instanceState1.setCurrentState(InstanceState.STATE.READY);

    ExecutionInstance processExecutionInstance2 =
        BeanMapperUtil.getExecutionInstance(
            entityState.getEntity().getEntityType(),
            entityState.getEntity(),
            System.currentTimeMillis(),
            "cluster1",
            System.currentTimeMillis());
    InstanceState instanceState2 = new InstanceState(processExecutionInstance2);
    instanceState2.setCurrentState(InstanceState.STATE.RUNNING);

    ExecutionInstance processExecutionInstance3 =
        BeanMapperUtil.getExecutionInstance(
            entityState.getEntity().getEntityType(),
            entityState.getEntity(),
            System.currentTimeMillis(),
            "cluster2",
            System.currentTimeMillis());
    InstanceState instanceState3 = new InstanceState(processExecutionInstance3);
    instanceState3.setCurrentState(InstanceState.STATE.READY);

    stateStore.putExecutionInstance(instanceState1);
    stateStore.putExecutionInstance(instanceState2);
    stateStore.putExecutionInstance(instanceState3);

    Collection<InstanceState> actualInstances =
        stateStore.getAllExecutionInstances(entityState.getEntity(), "cluster1");
    Assert.assertEquals(actualInstances.size(), 2);
    Assert.assertEquals(actualInstances.toArray()[0], instanceState1);
    Assert.assertEquals(actualInstances.toArray()[1], instanceState2);

    actualInstances = stateStore.getAllExecutionInstances(entityState.getEntity(), "cluster2");
    Assert.assertEquals(actualInstances.size(), 1);
    Assert.assertEquals(actualInstances.toArray()[0], instanceState3);

    List<InstanceState.STATE> states = new ArrayList<>();
    states.add(InstanceState.STATE.READY);

    actualInstances = stateStore.getExecutionInstances(entityState.getEntity(), "cluster1", states);
    Assert.assertEquals(actualInstances.size(), 1);
    Assert.assertEquals(actualInstances.toArray()[0], instanceState1);

    EntityClusterID entityClusterID = new EntityClusterID(entityState.getEntity(), "testCluster");
    actualInstances = stateStore.getExecutionInstances(entityClusterID, states);
    Assert.assertEquals(actualInstances.size(), 2);
    Assert.assertEquals(actualInstances.toArray()[0], instanceState1);
    Assert.assertEquals(actualInstances.toArray()[1], instanceState3);

    states.add(InstanceState.STATE.RUNNING);
    actualInstances = stateStore.getExecutionInstances(entityState.getEntity(), "cluster1", states);
    Assert.assertEquals(actualInstances.size(), 2);
    Assert.assertEquals(actualInstances.toArray()[0], instanceState1);
    Assert.assertEquals(actualInstances.toArray()[1], instanceState2);

    InstanceState lastInstanceState =
        stateStore.getLastExecutionInstance(entityState.getEntity(), "cluster1");
    Assert.assertEquals(lastInstanceState, instanceState2);

    InstanceID instanceKey = new InstanceID(instanceState3.getInstance());
    stateStore.deleteExecutionInstance(instanceKey);

    actualInstances = stateStore.getAllExecutionInstances(entityState.getEntity(), "cluster2");
    Assert.assertEquals(actualInstances.size(), 0);

    actualInstances = stateStore.getAllExecutionInstances(entityState.getEntity(), "cluster1");
    Assert.assertEquals(actualInstances.size(), 2);

    stateStore.putExecutionInstance(instanceState3);

    actualInstances = stateStore.getAllExecutionInstances(entityState.getEntity(), "cluster2");
    Assert.assertEquals(actualInstances.size(), 1);

    stateStore.deleteExecutionInstances(entityClusterID.getEntityID());
    actualInstances = stateStore.getAllExecutionInstances(entityState.getEntity(), "cluster1");
    Assert.assertEquals(actualInstances.size(), 0);

    actualInstances = stateStore.getAllExecutionInstances(entityState.getEntity(), "cluster2");
    Assert.assertEquals(actualInstances.size(), 0);
  }

  @Test
  public void testGetExecutionInstancesWithRange() throws Exception {
    storeEntity(EntityType.CLUSTER, "testCluster");
    storeEntity(EntityType.FEED, "clicksFeed");
    storeEntity(EntityType.FEED, "clicksSummary");

    long instance1Time = System.currentTimeMillis() - 180000;
    long instance2Time = System.currentTimeMillis();
    EntityState entityState = getEntityState(EntityType.PROCESS, "process1");
    ExecutionInstance processExecutionInstance1 =
        BeanMapperUtil.getExecutionInstance(
            entityState.getEntity().getEntityType(),
            entityState.getEntity(),
            instance1Time,
            "cluster1",
            instance1Time);
    InstanceState instanceState1 = new InstanceState(processExecutionInstance1);
    instanceState1.setCurrentState(InstanceState.STATE.RUNNING);

    ExecutionInstance processExecutionInstance2 =
        BeanMapperUtil.getExecutionInstance(
            entityState.getEntity().getEntityType(),
            entityState.getEntity(),
            instance2Time,
            "cluster1",
            instance2Time);
    InstanceState instanceState2 = new InstanceState(processExecutionInstance2);
    instanceState2.setCurrentState(InstanceState.STATE.RUNNING);

    ExecutionInstance processExecutionInstance3 =
        BeanMapperUtil.getExecutionInstance(
            entityState.getEntity().getEntityType(),
            entityState.getEntity(),
            instance2Time,
            "cluster2",
            instance2Time);
    InstanceState instanceState3 = new InstanceState(processExecutionInstance3);
    instanceState3.setCurrentState(InstanceState.STATE.RUNNING);

    stateStore.putExecutionInstance(instanceState1);
    stateStore.putExecutionInstance(instanceState2);
    stateStore.putExecutionInstance(instanceState3);

    List<InstanceState.STATE> states = new ArrayList<>();
    states.add(InstanceState.STATE.RUNNING);

    Collection<InstanceState> actualInstances =
        stateStore.getExecutionInstances(
            entityState.getEntity(),
            "cluster1",
            states,
            new DateTime(instance1Time),
            new DateTime(instance1Time + 60000));
    Assert.assertEquals(actualInstances.size(), 1);
    Assert.assertEquals(actualInstances.toArray()[0], instanceState1);

    actualInstances =
        stateStore.getExecutionInstances(
            entityState.getEntity(),
            "cluster1",
            states,
            new DateTime(instance2Time),
            new DateTime(instance2Time + 60000));
    Assert.assertEquals(actualInstances.size(), 1);
    Assert.assertEquals(actualInstances.toArray()[0], instanceState2);

    // Ensure we can get instances for a different cluster
    actualInstances =
        stateStore.getExecutionInstances(
            entityState.getEntity(),
            "cluster2",
            states,
            new DateTime(instance2Time),
            new DateTime(instance2Time + 60000));
    Assert.assertEquals(actualInstances.size(), 1);
    Assert.assertEquals(actualInstances.toArray()[0], instanceState3);
  }

  @Test
  public void testGetInstanceFromExternalID() throws Exception {
    storeEntity(EntityType.CLUSTER, "testCluster");
    storeEntity(EntityType.FEED, "clicksFeed");
    storeEntity(EntityType.FEED, "clicksSummary");

    long instance1Time = System.currentTimeMillis() - 180000;
    long instance2Time = System.currentTimeMillis();
    EntityState entityState = getEntityState(EntityType.PROCESS, "processext");
    ExecutionInstance processExecutionInstance1 =
        BeanMapperUtil.getExecutionInstance(
            entityState.getEntity().getEntityType(),
            entityState.getEntity(),
            instance1Time,
            "cluster1",
            instance1Time);
    processExecutionInstance1.setExternalID("external_id_1");
    InstanceState instanceState1 = new InstanceState(processExecutionInstance1);
    instanceState1.setCurrentState(InstanceState.STATE.RUNNING);

    ExecutionInstance processExecutionInstance2 =
        BeanMapperUtil.getExecutionInstance(
            entityState.getEntity().getEntityType(),
            entityState.getEntity(),
            instance2Time,
            "cluster1",
            instance2Time);
    processExecutionInstance2.setExternalID("external_id_2");
    InstanceState instanceState2 = new InstanceState(processExecutionInstance2);
    instanceState2.setCurrentState(InstanceState.STATE.RUNNING);

    stateStore.putExecutionInstance(instanceState1);
    stateStore.putExecutionInstance(instanceState2);

    InstanceState actualInstanceState = stateStore.getExecutionInstance("external_id_1");
    Assert.assertEquals(actualInstanceState.getInstance(), processExecutionInstance1);

    actualInstanceState = stateStore.getExecutionInstance("external_id_2");
    Assert.assertEquals(actualInstanceState.getInstance(), processExecutionInstance2);
  }

  @Test
  public void testCascadingDelete() throws Exception {
    storeEntity(EntityType.CLUSTER, "testCluster");
    storeEntity(EntityType.FEED, "clicksFeed");
    storeEntity(EntityType.FEED, "clicksSummary");
    EntityState entityState = getEntityState(EntityType.PROCESS, "process1");
    stateStore.putEntity(entityState);
    ExecutionInstance processExecutionInstance1 =
        BeanMapperUtil.getExecutionInstance(
            entityState.getEntity().getEntityType(),
            entityState.getEntity(),
            System.currentTimeMillis() - 60000,
            "cluster1",
            System.currentTimeMillis() - 60000);
    InstanceState instanceState1 = new InstanceState(processExecutionInstance1);
    instanceState1.setCurrentState(InstanceState.STATE.READY);

    ExecutionInstance processExecutionInstance2 =
        BeanMapperUtil.getExecutionInstance(
            entityState.getEntity().getEntityType(),
            entityState.getEntity(),
            System.currentTimeMillis(),
            "cluster1",
            System.currentTimeMillis());
    InstanceState instanceState2 = new InstanceState(processExecutionInstance2);
    instanceState2.setCurrentState(InstanceState.STATE.RUNNING);

    stateStore.putExecutionInstance(instanceState1);
    stateStore.putExecutionInstance(instanceState2);

    Collection<InstanceState> instances =
        stateStore.getAllExecutionInstances(entityState.getEntity(), "cluster1");
    Assert.assertEquals(instances.size(), 2);

    stateStore.deleteEntity(new EntityID(entityState.getEntity()));
    deleteEntity(EntityType.PROCESS, "process1");

    instances = stateStore.getAllExecutionInstances(entityState.getEntity(), "cluster1");
    Assert.assertEquals(instances.size(), 0);
  }

  @Test
  public void testGetExecutionSummaryWithRange() throws Exception {
    storeEntity(EntityType.CLUSTER, "testCluster");
    storeEntity(EntityType.FEED, "clicksFeed");
    storeEntity(EntityType.FEED, "clicksSummary");

    long instance1Time = System.currentTimeMillis() - 180000;
    long instance2Time = System.currentTimeMillis();
    EntityState entityState = getEntityState(EntityType.PROCESS, "clicksProcess");
    ExecutionInstance processExecutionInstance1 =
        BeanMapperUtil.getExecutionInstance(
            entityState.getEntity().getEntityType(),
            entityState.getEntity(),
            instance1Time,
            "cluster1",
            instance1Time);
    InstanceState instanceState1 = new InstanceState(processExecutionInstance1);
    instanceState1.setCurrentState(InstanceState.STATE.RUNNING);

    ExecutionInstance processExecutionInstance2 =
        BeanMapperUtil.getExecutionInstance(
            entityState.getEntity().getEntityType(),
            entityState.getEntity(),
            instance2Time,
            "cluster1",
            instance2Time);
    InstanceState instanceState2 = new InstanceState(processExecutionInstance2);
    instanceState2.setCurrentState(InstanceState.STATE.SUCCEEDED);

    stateStore.putExecutionInstance(instanceState1);
    stateStore.putExecutionInstance(instanceState2);

    Map<InstanceState.STATE, Long> summary =
        stateStore.getExecutionInstanceSummary(
            entityState.getEntity(),
            "cluster1",
            new DateTime(instance1Time),
            new DateTime(instance1Time + 60000));
    Assert.assertEquals(summary.size(), 1);
    Assert.assertEquals(summary.get(InstanceState.STATE.RUNNING).longValue(), 1L);

    summary =
        stateStore.getExecutionInstanceSummary(
            entityState.getEntity(),
            "cluster1",
            new DateTime(instance2Time),
            new DateTime(instance2Time + 60000));
    Assert.assertEquals(summary.size(), 1);
    Assert.assertEquals(summary.get(InstanceState.STATE.SUCCEEDED).longValue(), 1L);
  }

  private void initInstanceState(InstanceState instanceState) {
    instanceState.setCurrentState(InstanceState.STATE.READY);
    instanceState.getInstance().setExternalID(RandomStringUtils.randomNumeric(6));
    instanceState.getInstance().setInstanceSequence(randomValGenerator.nextInt());
    instanceState.getInstance().setActualStart(new DateTime(System.currentTimeMillis()));
    instanceState.getInstance().setActualEnd(new DateTime(System.currentTimeMillis()));
    List<Predicate> predicates = new ArrayList<>();
    Predicate predicate = new Predicate(Predicate.TYPE.JOB_COMPLETION);
    predicates.add(predicate);
    instanceState.getInstance().setAwaitingPredicates(predicates);
  }

  private EntityState getEntityState(EntityType entityType, String name) throws Exception {
    storeEntity(entityType, name);
    Entity entity = getStore().get(entityType, name);
    Assert.assertNotNull(entity);
    return new EntityState(entity);
  }

  @AfterTest
  public void cleanUpTables() throws StateStoreException {
    try {
      stateStore.deleteEntities();
    } catch (Exception e) {
      // ignore
    }
  }

  @AfterClass
  public void cleanup() throws IOException {
    super.cleanup();
  }
}