/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.gobblin.runtime;

import java.io.File;
import java.io.IOException;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

import org.apache.commons.io.FileUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.mockito.Mockito;
import static org.mockito.ArgumentMatchers.any;
import org.testng.Assert;
import org.testng.annotations.AfterClass;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;

import com.google.common.base.Optional;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterators;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.typesafe.config.Config;
import com.typesafe.config.ConfigFactory;
import com.typesafe.config.ConfigValueFactory;

import kafka.consumer.ConsumerIterator;
import kafka.message.MessageAndMetadata;
import lombok.Getter;

import org.apache.gobblin.configuration.ConfigurationKeys;
import org.apache.gobblin.configuration.ErrorPatternProfile;
import org.apache.gobblin.configuration.State;
import org.apache.gobblin.kafka.KafkaTestBase;
import org.apache.gobblin.kafka.client.DecodeableKafkaRecord;
import org.apache.gobblin.kafka.client.Kafka09ConsumerClient;
import org.apache.gobblin.metastore.InMemoryErrorPatternStore;
import org.apache.gobblin.metastore.StateStore;
import org.apache.gobblin.metrics.FlowStatus;
import org.apache.gobblin.metrics.GaaSFlowObservabilityEvent;
import org.apache.gobblin.metrics.GaaSJobObservabilityEvent;
import org.apache.gobblin.metrics.GobblinTrackingEvent;
import org.apache.gobblin.metrics.JobStatus;
import org.apache.gobblin.metrics.MetricContext;
import org.apache.gobblin.metrics.event.TimingEvent;
import org.apache.gobblin.metrics.kafka.KafkaAvroEventKeyValueReporter;
import org.apache.gobblin.metrics.kafka.KafkaEventReporter;
import org.apache.gobblin.metrics.kafka.KafkaKeyValueProducerPusher;
import org.apache.gobblin.metrics.kafka.Pusher;
import org.apache.gobblin.runtime.troubleshooter.InMemoryMultiContextIssueRepository;
import org.apache.gobblin.runtime.troubleshooter.IssueTestDataProvider;
import org.apache.gobblin.runtime.troubleshooter.JobIssueEventHandler;
import org.apache.gobblin.runtime.troubleshooter.MultiContextIssueRepository;
import org.apache.gobblin.runtime.troubleshooter.TroubleshooterException;
import org.apache.gobblin.service.ExecutionStatus;
import org.apache.gobblin.service.ServiceConfigKeys;
import org.apache.gobblin.service.modules.orchestration.DagActionStore;
import org.apache.gobblin.service.modules.orchestration.DagManagementStateStore;
import org.apache.gobblin.service.monitoring.GaaSJobObservabilityEventProducer;
import org.apache.gobblin.service.monitoring.JobStatusRetriever;
import org.apache.gobblin.service.monitoring.KafkaAvroJobStatusMonitor;
import org.apache.gobblin.service.monitoring.KafkaJobStatusMonitor;
import org.apache.gobblin.service.monitoring.MockGaaSJobObservabilityEventProducer;
import org.apache.gobblin.service.monitoring.NoopGaaSJobObservabilityEventProducer;
import org.apache.gobblin.util.ConfigUtils;

import static org.apache.gobblin.util.retry.RetryerFactory.RETRY_MULTIPLIER;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;


public class KafkaAvroJobStatusMonitorTest {
  public static final String TOPIC = KafkaAvroJobStatusMonitorTest.class.getSimpleName();

  private KafkaTestBase kafkaTestHelper;
  private final String flowGroup = "myFlowGroup";
  private final String flowName = "myFlowName";
  private final String jobGroup = "myJobGroup";
  private final String jobName = "myJobName";
  private final long flowExecutionId = 1234L;
  private final String stateStoreDir = "/tmp/jobStatusMonitor/statestore";
  private MetricContext context;
  private KafkaAvroEventKeyValueReporter.Builder<?> builder;
  DagActionStore.DagAction enforceJobStartDeadlineDagAction;
  private JobIssueEventHandler eventHandler;
  private ErrorClassifier errorClassifier;

  @BeforeClass
  public void setUp() throws Exception {
    cleanUpDir();
    kafkaTestHelper = new KafkaTestBase();
    kafkaTestHelper.startServers();
    kafkaTestHelper.provisionTopic(TOPIC);

    // Create KeyValueProducerPusher instance.
    Pusher<Pair<byte[], byte[]>> pusher = new KafkaKeyValueProducerPusher<>("localhost:dummy", TOPIC,
        Optional.of(ConfigFactory.parseMap(ImmutableMap.of(
            ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:" + this.kafkaTestHelper.getKafkaServerPort()))));

    //Create an event reporter instance
    context = MetricContext.builder("context").build();
    builder = KafkaAvroEventKeyValueReporter.Factory.forContext(context);
    builder = builder.withKafkaPusher(pusher).withKeys(Lists.newArrayList(TimingEvent.FlowEventConstants.FLOW_NAME_FIELD,
        TimingEvent.FlowEventConstants.FLOW_GROUP_FIELD, TimingEvent.FlowEventConstants.FLOW_EXECUTION_ID_FIELD));

    this.enforceJobStartDeadlineDagAction = new DagActionStore.DagAction(this.flowGroup, this.flowName, this.flowExecutionId, this.jobName,
        DagActionStore.DagActionType.ENFORCE_JOB_START_DEADLINE);
  }

  @BeforeMethod
  public void setUpMethod() {
    eventHandler = Mockito.mock(JobIssueEventHandler.class);
    errorClassifier = Mockito.mock(ErrorClassifier.class);
  }

  @Test
  public void testProcessMessageForSuccessfulFlow() throws IOException, ReflectiveOperationException {
    DagManagementStateStore dagManagementStateStore = mock(DagManagementStateStore.class);
    KafkaEventReporter kafkaReporter = builder.build("localhost:0000", "topic1");

    //Submit GobblinTrackingEvents to Kafka
    ImmutableList.of(
        createFlowCompiledEvent(),
        createJobOrchestratedEvent(1, 2),
        createJobStartEvent(),
        createJobSucceededEvent(),
        createDummyEvent(), // note position
        createJobStartEvent()
    ).forEach(event -> {
      context.submitEvent(event);
      kafkaReporter.report();
    });

    try {
      Thread.sleep(1000);
    } catch(InterruptedException ex) {
      Thread.currentThread().interrupt();
    }
    MockKafkaAvroJobStatusMonitor jobStatusMonitor = createMockKafkaAvroJobStatusMonitor(new AtomicBoolean(false), ConfigFactory.empty(),
        new NoopGaaSJobObservabilityEventProducer(), dagManagementStateStore);
    jobStatusMonitor.buildMetricsContextAndMetrics();

    Iterator<DecodeableKafkaRecord<byte[], byte[]>> recordIterator = Iterators.transform(
        this.kafkaTestHelper.getIteratorForTopic(TOPIC),
        this::convertMessageAndMetadataToDecodableKafkaRecord);

    State state = getNextJobStatusState(jobStatusMonitor, recordIterator, "NA", "NA");
    Assert.assertEquals(state.getProp(JobStatusRetriever.EVENT_NAME_FIELD), ExecutionStatus.COMPILED.name());
    Mockito.verify(dagManagementStateStore, Mockito.never()).addJobDagAction(any(), any(), anyLong(), any(),
        eq(DagActionStore.DagActionType.REEVALUATE));
    Mockito.verify(dagManagementStateStore, Mockito.never()).deleteDagAction(eq(this.enforceJobStartDeadlineDagAction));

    state = getNextJobStatusState(jobStatusMonitor, recordIterator, this.jobGroup, this.jobName);
    Assert.assertEquals(state.getProp(JobStatusRetriever.EVENT_NAME_FIELD), ExecutionStatus.ORCHESTRATED.name());
    Mockito.verify(dagManagementStateStore, Mockito.never()).addJobDagAction(any(), any(), anyLong(), any(),
        eq(DagActionStore.DagActionType.REEVALUATE));
    Mockito.verify(dagManagementStateStore, Mockito.never()).deleteDagAction(eq(this.enforceJobStartDeadlineDagAction));

    state = getNextJobStatusState(jobStatusMonitor, recordIterator, this.jobGroup, this.jobName);
    Assert.assertEquals(state.getProp(JobStatusRetriever.EVENT_NAME_FIELD), ExecutionStatus.RUNNING.name());
    Mockito.verify(dagManagementStateStore, Mockito.never()).addJobDagAction(any(), any(), anyLong(), any(),
        eq(DagActionStore.DagActionType.REEVALUATE));
    Mockito.verify(dagManagementStateStore, Mockito.times(1)).deleteDagAction(eq(this.enforceJobStartDeadlineDagAction));

    state = getNextJobStatusState(jobStatusMonitor, recordIterator, this.jobGroup, this.jobName);
    Assert.assertEquals(state.getProp(JobStatusRetriever.EVENT_NAME_FIELD), ExecutionStatus.COMPLETE.name());
    Mockito.verify(dagManagementStateStore, Mockito.times(1)).addJobDagAction(any(), any(),
        anyLong(), any(), eq(DagActionStore.DagActionType.REEVALUATE));
    Mockito.verify(dagManagementStateStore, Mockito.times(1)).deleteDagAction(eq(this.enforceJobStartDeadlineDagAction));

    // (per above, is a 'dummy' event)
    Assert.assertNull(jobStatusMonitor.parseJobStatus(
        jobStatusMonitor.deserializeEvent(recordIterator.next())));

    // Check that state didn't get set to running since it was already complete
    state = getNextJobStatusState(jobStatusMonitor, recordIterator, this.jobGroup, this.jobName);
    Assert.assertEquals(state.getProp(JobStatusRetriever.EVENT_NAME_FIELD), ExecutionStatus.COMPLETE.name());

    jobStatusMonitor.shutDown();
  }

  /**
   * Functional test for ErrorClassifier integration with KafkaJobStatusMonitor:
   * Tests that when a job fails, the monitor properly invokes ErrorClassifier to classify issues
   * and logs the final classified error through the JobIssueEventHandler.
   */
  @Test (dependsOnMethods = "testObservabilityEventFlowFailed")
  public void testErrorClassifierFinalIssueSystemInfra()
      throws IOException, ReflectiveOperationException, TroubleshooterException {
    DagManagementStateStore dagManagementStateStore = mock(DagManagementStateStore.class);
    Config config = ConfigFactory.empty().withValue(ServiceConfigKeys.ERROR_CLASSIFICATION_ENABLED_KEY, ConfigValueFactory.fromAnyRef(Boolean.TRUE));
    List<ErrorPatternProfile> errorPatternProfiles = IssueTestDataProvider.getSortedPatterns();

    InMemoryErrorPatternStore store = new InMemoryErrorPatternStore(config);
    store.upsertCategory(IssueTestDataProvider.TEST_CATEGORIES);
    store.upsertPatterns(errorPatternProfiles);
    store.setDefaultCategory(IssueTestDataProvider.TEST_DEFAULT_ERROR_CATEGORY);

    this.errorClassifier =  new ErrorClassifier(store, config);
    KafkaEventReporter kafkaReporter = builder.build("localhost:0000", "topic2");

    //Submit GobblinTrackingEvents to Kafka
    ImmutableList.of(
        createJobFailedEvent()
    ).forEach(event -> {
      context.submitEvent(event);
      kafkaReporter.report();
    });

    //Set issues for error classification
    Mockito.doReturn(IssueTestDataProvider.testSystemInfraCategoryIssues()).when(eventHandler).getErrorListForClassification(Mockito.any());

    try {
      Thread.sleep(1000L);
    } catch(InterruptedException ex) {
      Thread.currentThread().interrupt();
    }

    MockKafkaAvroJobStatusMonitor jobStatusMonitor = createMockKafkaAvroJobStatusMonitor(new AtomicBoolean(false),
        config, new NoopGaaSJobObservabilityEventProducer(), dagManagementStateStore);
    jobStatusMonitor.buildMetricsContextAndMetrics();

    ConsumerIterator<byte[], byte[]> iterator = this.kafkaTestHelper.getIteratorForTopic(TOPIC);
    MessageAndMetadata<byte[], byte[]> messageAndMetadata = iterator.next();
    jobStatusMonitor.processMessage(convertMessageAndMetadataToDecodableKafkaRecord(messageAndMetadata));

    // Verify that the JobIssueEventHandler was called with the final classified error
    Mockito.verify(eventHandler, Mockito.times(1))
        .logFinalError(Mockito.argThat(issue ->
                issue.getSummary().contains("SYSTEM INFRA")),
            Mockito.eq(this.flowName),
            Mockito.eq(this.flowGroup),
            Mockito.eq(String.valueOf(this.flowExecutionId)),
            Mockito.eq(this.jobName));

    jobStatusMonitor.shutDown();
  }

  /**
   * Test for verifying that error classifier module is not called when it is disabled.
   */
  @Test (dependsOnMethods = "testErrorClassifierFinalIssueSystemInfra")
  public void testDisabledErrorClassification()
      throws IOException, ReflectiveOperationException, TroubleshooterException {
    DagManagementStateStore dagManagementStateStore = mock(DagManagementStateStore.class);
    Config config = ConfigFactory.empty().withValue(ServiceConfigKeys.ERROR_CLASSIFICATION_ENABLED_KEY, ConfigValueFactory.fromAnyRef(Boolean.FALSE));
    List<ErrorPatternProfile> errorPatternProfiles = IssueTestDataProvider.getSortedPatterns();

    InMemoryErrorPatternStore store = new InMemoryErrorPatternStore(config);
    store.upsertCategory(IssueTestDataProvider.TEST_CATEGORIES);
    store.upsertPatterns(errorPatternProfiles);
    store.setDefaultCategory(IssueTestDataProvider.TEST_DEFAULT_ERROR_CATEGORY);

    this.errorClassifier =  new ErrorClassifier(store, config);
    KafkaEventReporter kafkaReporter = builder.build("localhost:0000", "topic2");

    //Submit GobblinTrackingEvents to Kafka
    ImmutableList.of(
        createJobFailedEvent()
    ).forEach(event -> {
      context.submitEvent(event);
      kafkaReporter.report();
    });

    //Set issues for error classification
    Mockito.doReturn(IssueTestDataProvider.testSystemInfraCategoryIssues()).when(eventHandler).getErrorListForClassification(Mockito.any());

    try {
      Thread.sleep(1000L);
    } catch(InterruptedException ex) {
      Thread.currentThread().interrupt();
    }

    MockKafkaAvroJobStatusMonitor jobStatusMonitor = createMockKafkaAvroJobStatusMonitor(new AtomicBoolean(false),
        config, new NoopGaaSJobObservabilityEventProducer(), dagManagementStateStore);
    jobStatusMonitor.buildMetricsContextAndMetrics();

    ConsumerIterator<byte[], byte[]> iterator = this.kafkaTestHelper.getIteratorForTopic(TOPIC);
    MessageAndMetadata<byte[], byte[]> messageAndMetadata = iterator.next();
    jobStatusMonitor.processMessage(convertMessageAndMetadataToDecodableKafkaRecord(messageAndMetadata));

    // Verify that the JobIssueEventHandler was not called for error classification
    Mockito.verify(eventHandler, Mockito.times(0))
        .logFinalError(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any());
    Mockito.verify(eventHandler, Mockito.times(0))
        .getErrorListForClassification(Mockito.any());

    jobStatusMonitor.shutDown();
  }

  @Test (dependsOnMethods = "testProcessMessageForSuccessfulFlow")
  public void testProcessMessageForFailedFlow() throws IOException, ReflectiveOperationException {
    DagManagementStateStore dagManagementStateStore = mock(DagManagementStateStore.class);
    KafkaEventReporter kafkaReporter = builder.build("localhost:0000", "topic2");

    //Submit GobblinTrackingEvents to Kafka
    ImmutableList.of(
        createFlowCompiledEvent(),
        createJobOrchestratedEvent(1, 2),
        createJobStartEvent(),
        createJobFailedEvent(),
        // Mimic retrying - job orchestration
        // set maximum attempt to 2, and current attempt to 2
        createJobOrchestratedEvent(2, 2),
        // Mimic retrying - job start (current attempt = 2)
        createJobStartEvent(),
        // Mimic retrying - job failed again (current attempt = 2)
        createJobFailedEvent()
    ).forEach(event -> {
      context.submitEvent(event);
      kafkaReporter.report();
    });

    try {
      Thread.sleep(1000L);
    } catch(InterruptedException ex) {
      Thread.currentThread().interrupt();
    }

    MockKafkaAvroJobStatusMonitor jobStatusMonitor = createMockKafkaAvroJobStatusMonitor(new AtomicBoolean(false),
        ConfigFactory.empty(), new NoopGaaSJobObservabilityEventProducer(), dagManagementStateStore);
    jobStatusMonitor.buildMetricsContextAndMetrics();

    ConsumerIterator<byte[], byte[]> iterator = this.kafkaTestHelper.getIteratorForTopic(TOPIC);

    MessageAndMetadata<byte[], byte[]> messageAndMetadata = iterator.next();
    // Verify undecodeable message is skipped
    byte[] undecodeableMessage = Arrays.copyOf(messageAndMetadata.message(),
        messageAndMetadata.message().length - 1);
    ConsumerRecord<byte[], byte[]> undecodeableRecord = new ConsumerRecord<>(TOPIC, messageAndMetadata.partition(),
        messageAndMetadata.offset(), messageAndMetadata.key(), undecodeableMessage);
    Assert.assertEquals(jobStatusMonitor.getMessageParseFailures().getCount(), 0L);
    jobStatusMonitor.processMessage(new Kafka09ConsumerClient.Kafka09ConsumerRecord<>(undecodeableRecord));
    Assert.assertEquals(jobStatusMonitor.getMessageParseFailures().getCount(), 1L);
    // Re-test when properly encoded, as expected for a normal event
    jobStatusMonitor.processMessage(convertMessageAndMetadataToDecodableKafkaRecord(messageAndMetadata));

    StateStore<State> stateStore = jobStatusMonitor.getStateStore();
    String storeName = KafkaJobStatusMonitor.jobStatusStoreName(flowGroup, flowName);
    String tableName = KafkaJobStatusMonitor.jobStatusTableName(this.flowExecutionId, "NA", "NA");
    List<State> stateList  = stateStore.getAll(storeName, tableName);
    Assert.assertEquals(stateList.size(), 1);
    State state = stateList.get(0);
    Assert.assertEquals(state.getProp(JobStatusRetriever.EVENT_NAME_FIELD), ExecutionStatus.COMPILED.name());

    Iterator<DecodeableKafkaRecord<byte[], byte[]>> recordIterator = Iterators.transform(
        iterator,
        this::convertMessageAndMetadataToDecodableKafkaRecord);

    state = getNextJobStatusState(jobStatusMonitor, recordIterator, this.jobGroup, this.jobName);
    Assert.assertEquals(state.getProp(JobStatusRetriever.EVENT_NAME_FIELD), ExecutionStatus.ORCHESTRATED.name());
    Mockito.verify(dagManagementStateStore, Mockito.never()).addJobDagAction(any(), any(), anyLong(), any(),
        eq(DagActionStore.DagActionType.REEVALUATE));
    Mockito.verify(dagManagementStateStore, Mockito.never()).deleteDagAction(eq(this.enforceJobStartDeadlineDagAction));

    state = getNextJobStatusState(jobStatusMonitor, recordIterator, this.jobGroup, this.jobName);
    Assert.assertEquals(state.getProp(JobStatusRetriever.EVENT_NAME_FIELD), ExecutionStatus.RUNNING.name());
    Mockito.verify(dagManagementStateStore, Mockito.never()).addJobDagAction(any(), any(), anyLong(), any(),
        eq(DagActionStore.DagActionType.REEVALUATE));
    Mockito.verify(dagManagementStateStore, Mockito.times(1)).deleteDagAction(eq(this.enforceJobStartDeadlineDagAction));

    state = getNextJobStatusState(jobStatusMonitor, recordIterator, this.jobGroup, this.jobName);

    //Because the maximum attempt is set to 2, so the state is set to PENDING_RETRY after the first failure
    Assert.assertEquals(state.getProp(JobStatusRetriever.EVENT_NAME_FIELD), ExecutionStatus.PENDING_RETRY.name());
    Assert.assertEquals(state.getProp(TimingEvent.FlowEventConstants.SHOULD_RETRY_FIELD), Boolean.toString(true));
    Mockito.verify(dagManagementStateStore, Mockito.times(1)).addJobDagAction(any(), any(), anyLong(), any(),
        eq(DagActionStore.DagActionType.REEVALUATE));
    Mockito.verify(dagManagementStateStore, Mockito.times(1)).deleteDagAction(eq(this.enforceJobStartDeadlineDagAction));

    state = getNextJobStatusState(jobStatusMonitor, recordIterator, this.jobGroup, this.jobName);
    //Job orchestrated for retrying
    Assert.assertEquals(state.getProp(JobStatusRetriever.EVENT_NAME_FIELD), ExecutionStatus.ORCHESTRATED.name());
    Mockito.verify(dagManagementStateStore, Mockito.times(1)).addJobDagAction(any(), any(), anyLong(), any(),
        eq(DagActionStore.DagActionType.REEVALUATE));
    Mockito.verify(dagManagementStateStore, Mockito.times(1)).deleteDagAction(eq(this.enforceJobStartDeadlineDagAction));

    state = getNextJobStatusState(jobStatusMonitor, recordIterator, this.jobGroup, this.jobName);
    //Because the maximum attempt is set to 2, so the state is set to PENDING_RETRY after the first failure
    Assert.assertEquals(state.getProp(JobStatusRetriever.EVENT_NAME_FIELD), ExecutionStatus.RUNNING.name());
    Mockito.verify(dagManagementStateStore, Mockito.times(1)).addJobDagAction(any(), any(), anyLong(), any(),
        eq(DagActionStore.DagActionType.REEVALUATE));
    Mockito.verify(dagManagementStateStore, Mockito.times(2)).deleteDagAction(eq(this.enforceJobStartDeadlineDagAction));

    state = getNextJobStatusState(jobStatusMonitor, recordIterator, this.jobGroup, this.jobName);
    //Because the maximum attempt is set to 2, so the state is set to Failed after trying twice
    Assert.assertEquals(state.getProp(JobStatusRetriever.EVENT_NAME_FIELD), ExecutionStatus.FAILED.name());
    Assert.assertEquals(state.getProp(TimingEvent.FlowEventConstants.SHOULD_RETRY_FIELD), Boolean.toString(false));
    Mockito.verify(dagManagementStateStore, Mockito.times(2)).addJobDagAction(any(), any(),
        anyLong(), any(), eq(DagActionStore.DagActionType.REEVALUATE));
    Mockito.verify(dagManagementStateStore, Mockito.times(2)).deleteDagAction(eq(this.enforceJobStartDeadlineDagAction));

    jobStatusMonitor.shutDown();
  }

  @Test (dependsOnMethods = "testProcessMessageForFailedFlow")
  public void testProcessMessageForSkippedFlow() throws IOException, ReflectiveOperationException {
    DagManagementStateStore dagManagementStateStore = mock(DagManagementStateStore.class);
    KafkaEventReporter kafkaReporter = builder.build("localhost:0000", "topic2");

    //Submit GobblinTrackingEvents to Kafka
    ImmutableList.of(
        createFlowCompiledEvent(),
        createJobOrchestratedEvent(1, 2),
        createJobSkippedEvent()
    ).forEach(event -> {
      context.submitEvent(event);
      kafkaReporter.report();
    });

    try {
      Thread.sleep(1000);
    } catch(InterruptedException ex) {
      Thread.currentThread().interrupt();
    }

    MockKafkaAvroJobStatusMonitor jobStatusMonitor = createMockKafkaAvroJobStatusMonitor(new AtomicBoolean(false), ConfigFactory.empty(),
        new NoopGaaSJobObservabilityEventProducer(), dagManagementStateStore);
    jobStatusMonitor.buildMetricsContextAndMetrics();

    ConsumerIterator<byte[], byte[]> iterator = this.kafkaTestHelper.getIteratorForTopic(TOPIC);

    MessageAndMetadata<byte[], byte[]> messageAndMetadata = iterator.next();
    // Verify undecodeable message is skipped
    byte[] undecodeableMessage = Arrays.copyOf(messageAndMetadata.message(),
        messageAndMetadata.message().length - 1);
    ConsumerRecord<byte[], byte[]> undecodeableRecord = new ConsumerRecord<>(TOPIC, messageAndMetadata.partition(),
        messageAndMetadata.offset(), messageAndMetadata.key(), undecodeableMessage);
    Assert.assertEquals(jobStatusMonitor.getMessageParseFailures().getCount(), 0L);
    jobStatusMonitor.processMessage(new Kafka09ConsumerClient.Kafka09ConsumerRecord<>(undecodeableRecord));
    Assert.assertEquals(jobStatusMonitor.getMessageParseFailures().getCount(), 1L);
    // Re-test when properly encoded, as expected for a normal event
    jobStatusMonitor.processMessage(convertMessageAndMetadataToDecodableKafkaRecord(messageAndMetadata));

    StateStore<State> stateStore = jobStatusMonitor.getStateStore();
    String storeName = KafkaJobStatusMonitor.jobStatusStoreName(flowGroup, flowName);
    String tableName = KafkaJobStatusMonitor.jobStatusTableName(this.flowExecutionId, "NA", "NA");
    List<State> stateList  = stateStore.getAll(storeName, tableName);
    Assert.assertEquals(stateList.size(), 1);
    State state = stateList.get(0);
    Assert.assertEquals(state.getProp(JobStatusRetriever.EVENT_NAME_FIELD), ExecutionStatus.COMPILED.name());

    Iterator<DecodeableKafkaRecord<byte[], byte[]>> recordIterator = Iterators.transform(
        iterator,
        this::convertMessageAndMetadataToDecodableKafkaRecord);

    state = getNextJobStatusState(jobStatusMonitor, recordIterator, this.jobGroup, this.jobName);
    Assert.assertEquals(state.getProp(JobStatusRetriever.EVENT_NAME_FIELD), ExecutionStatus.ORCHESTRATED.name());
    Mockito.verify(dagManagementStateStore, Mockito.never()).addJobDagAction(any(), any(), anyLong(), any(),
        eq(DagActionStore.DagActionType.REEVALUATE));
    Mockito.verify(dagManagementStateStore, Mockito.never()).deleteDagAction(eq(this.enforceJobStartDeadlineDagAction));

    state = getNextJobStatusState(jobStatusMonitor, recordIterator, this.jobGroup, this.jobName);
    Assert.assertEquals(state.getProp(JobStatusRetriever.EVENT_NAME_FIELD), ExecutionStatus.CANCELLED.name());
    Mockito.verify(dagManagementStateStore, Mockito.times(1)).addJobDagAction(any(), any(),
        anyLong(), any(), eq(DagActionStore.DagActionType.REEVALUATE));
    Mockito.verify(dagManagementStateStore, Mockito.times(1)).deleteDagAction(eq(this.enforceJobStartDeadlineDagAction));
    jobStatusMonitor.shutDown();
  }

  @Test (dependsOnMethods = "testProcessMessageForSkippedFlow")
  public void testProcessingRetriedForApparentlyTransientErrors() throws IOException, ReflectiveOperationException {
    DagManagementStateStore dagManagementStateStore = mock(DagManagementStateStore.class);
    KafkaEventReporter kafkaReporter = builder.build("localhost:0000", "topic3");

    //Submit GobblinTrackingEvents to Kafka
    ImmutableList.of(
        createFlowCompiledEvent(),
        createJobOrchestratedEvent(1, 2)
    ).forEach(event -> {
      context.submitEvent(event);
      kafkaReporter.report();
    });

    try {
      Thread.sleep(1000L);
    } catch(InterruptedException ex) {
      Thread.currentThread().interrupt();
    }
    int minNumFakeExceptionsExpected = 10;
    AtomicBoolean shouldThrowFakeExceptionInParseJobStatusToggle = new AtomicBoolean(false);
    Config conf = ConfigFactory.empty().withValue(
        KafkaJobStatusMonitor.JOB_STATUS_MONITOR_PREFIX + "." + RETRY_MULTIPLIER, ConfigValueFactory.fromAnyRef(TimeUnit.MILLISECONDS.toMillis(1L)));
    MockKafkaAvroJobStatusMonitor jobStatusMonitor = createMockKafkaAvroJobStatusMonitor(shouldThrowFakeExceptionInParseJobStatusToggle, conf,
        new NoopGaaSJobObservabilityEventProducer(), dagManagementStateStore);

    jobStatusMonitor.buildMetricsContextAndMetrics();
    Iterator<DecodeableKafkaRecord<byte[], byte[]>> recordIterator = Iterators.transform(
        this.kafkaTestHelper.getIteratorForTopic(TOPIC),
        this::convertMessageAndMetadataToDecodableKafkaRecord);

    State state = getNextJobStatusState(jobStatusMonitor, recordIterator, "NA", "NA");;
    Assert.assertEquals(state.getProp(JobStatusRetriever.EVENT_NAME_FIELD), ExecutionStatus.COMPILED.name());

    shouldThrowFakeExceptionInParseJobStatusToggle.set(true);
    // since `processMessage` hereafter effectively hangs, launch eventual re-toggling before calling again
    ScheduledExecutorService toggleManagementExecutor = Executors.newScheduledThreadPool(2);
    toggleManagementExecutor.scheduleAtFixedRate(() -> {
      if (jobStatusMonitor.getNumFakeExceptionsFromParseJobStatus().intValue() > minNumFakeExceptionsExpected) { // curtail faking: simulate resolution
        shouldThrowFakeExceptionInParseJobStatusToggle.set(false);
      }
    }, 2, 2, TimeUnit.SECONDS);
    Thread mainThread = Thread.currentThread();
    // guardrail against excessive retries (befitting this unit test):
    toggleManagementExecutor.scheduleAtFixedRate(mainThread::interrupt, 20, 5, TimeUnit.SECONDS);

    state = getNextJobStatusState(jobStatusMonitor, recordIterator, this.jobGroup, this.jobName);

    Assert.assertTrue(jobStatusMonitor.getNumFakeExceptionsFromParseJobStatus().intValue() > minNumFakeExceptionsExpected,
        String.format("processMessage returned with only %d (faked) exceptions",
            jobStatusMonitor.getNumFakeExceptionsFromParseJobStatus().intValue()));

    Assert.assertEquals(state.getProp(JobStatusRetriever.EVENT_NAME_FIELD), ExecutionStatus.ORCHESTRATED.name());
    Mockito.verify(dagManagementStateStore, Mockito.never()).addJobDagAction(any(), any(), anyLong(), any(),
        eq(DagActionStore.DagActionType.REEVALUATE));
    Mockito.verify(dagManagementStateStore, Mockito.never()).deleteDagAction(eq(this.enforceJobStartDeadlineDagAction));

    toggleManagementExecutor.shutdownNow();
    jobStatusMonitor.shutDown();
  }

  @Test (dependsOnMethods = "testProcessingRetriedForApparentlyTransientErrors")
  public void testProcessMessageForCancelledAndKilledEvent() throws IOException, ReflectiveOperationException {
    DagManagementStateStore dagManagementStateStore = mock(DagManagementStateStore.class);
    KafkaEventReporter kafkaReporter = builder.build("localhost:0000", "topic4");

    //Submit GobblinTrackingEvents to Kafka
    ImmutableList.of(
        createFlowCompiledEvent(),
        createJobOrchestratedEvent(1, 4),
        createJobSLAKilledEvent(),
        createJobOrchestratedEvent(2, 4),
        createJobStartSLAKilledEvent(),
        // Verify that kill event will not retry
        createJobOrchestratedEvent(3, 4),
        createJobCancelledEvent()
    ).forEach(event -> {
      context.submitEvent(event);
      kafkaReporter.report();
    });

    try {
      Thread.sleep(1000);
    } catch(InterruptedException ex) {
      Thread.currentThread().interrupt();
    }

    MockKafkaAvroJobStatusMonitor jobStatusMonitor = createMockKafkaAvroJobStatusMonitor(new AtomicBoolean(false),
        ConfigFactory.empty(), new NoopGaaSJobObservabilityEventProducer(), dagManagementStateStore);
    jobStatusMonitor.buildMetricsContextAndMetrics();
    Iterator<DecodeableKafkaRecord<byte[], byte[]>> recordIterator = Iterators.transform(
        this.kafkaTestHelper.getIteratorForTopic(TOPIC),
        this::convertMessageAndMetadataToDecodableKafkaRecord);

    State state = getNextJobStatusState(jobStatusMonitor, recordIterator, "NA", "NA");
    Assert.assertEquals(state.getProp(JobStatusRetriever.EVENT_NAME_FIELD), ExecutionStatus.COMPILED.name());
    Mockito.verify(dagManagementStateStore, Mockito.never()).addJobDagAction(any(), any(), anyLong(), any(),
        eq(DagActionStore.DagActionType.REEVALUATE));
    Mockito.verify(dagManagementStateStore, Mockito.never()).deleteDagAction(eq(this.enforceJobStartDeadlineDagAction));

    state = getNextJobStatusState(jobStatusMonitor, recordIterator, this.jobGroup, this.jobName);
    Assert.assertEquals(state.getProp(JobStatusRetriever.EVENT_NAME_FIELD), ExecutionStatus.ORCHESTRATED.name());
    Mockito.verify(dagManagementStateStore, Mockito.never()).addJobDagAction(any(), any(), anyLong(), any(),
        eq(DagActionStore.DagActionType.REEVALUATE));
    Mockito.verify(dagManagementStateStore, Mockito.never()).deleteDagAction(eq(this.enforceJobStartDeadlineDagAction));

    state = getNextJobStatusState(jobStatusMonitor, recordIterator, this.jobGroup, this.jobName);
    Assert.assertEquals(state.getProp(JobStatusRetriever.EVENT_NAME_FIELD), ExecutionStatus.PENDING_RETRY.name());
    Assert.assertEquals(state.getProp(TimingEvent.FlowEventConstants.SHOULD_RETRY_FIELD), Boolean.toString(true));
    Mockito.verify(dagManagementStateStore, Mockito.times(1)).addJobDagAction(any(), any(), anyLong(), any(),
        eq(DagActionStore.DagActionType.REEVALUATE));
    // this is a cancelled job, so its JobStartDeadlineDagAction should get removed
    Mockito.verify(dagManagementStateStore, Mockito.times(1)).deleteDagAction(eq(this.enforceJobStartDeadlineDagAction));

    state = getNextJobStatusState(jobStatusMonitor, recordIterator, this.jobGroup, this.jobName);
    //Job orchestrated for retrying
