MET-6679 more tests for ProgressUpdater and RegularDbReader

Author: pWoz

commons/src/main/java/eu/europeana/processing/source/AbstractDbReader.java

@@ -67,11 +67,8 @@ protected AbstractDbReader(
     public void start() {
         LOGGER.info("Starting: {}", getClass().getSimpleName());
         dbConnectionProvider = new DbConnectionProvider(parameterTool);
-        createRepositories();
     }
 
-    protected abstract void createRepositories() ;
-
     @Override
     public InputStatus pollNext(ReaderOutput<R> output) throws Exception {
         LOGGER.debug("Pooling next record");

commons/src/main/java/eu/europeana/processing/source/AbstractEnumerator.java

@@ -90,7 +90,10 @@ protected AbstractEnumerator(SplitEnumeratorContext<P> context, ParameterTool pa
   public void start() {
     LOGGER.info("Starting enumerator");
     dbConnectionProvider = new DbConnectionProvider(parameterTool);
-    progressUpdater = new ProgressUpdater(dbConnectionProvider, parameterTool, emittedRecordCount);
+    progressUpdater = new ProgressUpdater(
+        new TaskInfoRepository(dbConnectionProvider),
+        parameterTool,
+        emittedRecordCount);
     createDbRepositories();
     taskInfoRepo = RetryableMethodExecutor.createRetryProxy(new TaskInfoRepository(dbConnectionProvider));
     validateTaskExists();

commons/src/main/java/eu/europeana/processing/source/DbSourceWithProgressHandling.java

@@ -1,16 +1,19 @@
 package eu.europeana.processing.source;
 
+import eu.europeana.processing.DbConnectionProvider;
 import eu.europeana.processing.model.DataPartition;
 import eu.europeana.processing.model.ExecutionRecord;
+import eu.europeana.processing.repository.ExecutionRecordRepository;
+import eu.europeana.processing.retryable.RetryableMethodExecutor;
 import java.io.Serial;
 import org.apache.flink.api.connector.source.Boundedness;
 import org.apache.flink.api.connector.source.Source;
 import org.apache.flink.api.connector.source.SourceReader;
 import org.apache.flink.api.connector.source.SourceReaderContext;
 import org.apache.flink.api.connector.source.SplitEnumerator;
 import org.apache.flink.api.connector.source.SplitEnumeratorContext;
-import org.apache.flink.util.ParameterTool;
 import org.apache.flink.core.io.SimpleVersionedSerializer;
+import org.apache.flink.util.ParameterTool;
 
 public class DbSourceWithProgressHandling implements Source<ExecutionRecord, DataPartition, DbEnumeratorState> {
 
@@ -42,7 +45,15 @@ public SplitEnumerator<DataPartition, DbEnumeratorState> restoreEnumerator(
 
   @Override
   public SourceReader<ExecutionRecord, DataPartition> createReader(SourceReaderContext readerContext) {
-    return new RegularDbReader(readerContext, parameterTool);
+    return new RegularDbReader(
+        readerContext,
+        parameterTool,
+        //TODO Using retry proxy is maybe not optimal strategy in this case. This source implements asynchronous interface, so
+        // we could do this retries in poolNext() method by returning InputStatus.NOTHING_AVAILABLE, wait a bit and notify
+        // completable future to poll source again. Or simple wait a bit in pollNext() but only once per one retry.
+        // In such cases we would less block checkpointing mechanism, which should work smoothly in case of infrastructure problems
+        // and potential job restarts. And when we do not block we could do more retries or longer pauses.
+        RetryableMethodExecutor.createRetryProxy(new ExecutionRecordRepository(new DbConnectionProvider(parameterTool))));
   }
 
   @Override

commons/src/main/java/eu/europeana/processing/source/ProgressUpdater.java

@@ -17,7 +17,6 @@ public class ProgressUpdater implements Closeable {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(ProgressUpdater.class);
   private final long taskId;
-  private final DbConnectionProvider dbConnectionProvider;
   private final TaskInfoRepository taskInfoRepo;
   private long lastStoredFilesCount;
   private long snapshottedEmittedFilesCount = -1;
@@ -28,13 +27,12 @@ public class ProgressUpdater implements Closeable {
    * @param parameterTool - all the command line parameters of the job
    * @param completedFilesCount - Number of files already completed. It is greater than 0 only if the source is restored from a
    * checkpoint.
-   * @param dbConnectionProvider - db connection provider
+   * @param taskInfoRepository - repository for TaskInfo table
    */
-  public ProgressUpdater(DbConnectionProvider dbConnectionProvider, ParameterTool parameterTool, long completedFilesCount) {
-    this.dbConnectionProvider = dbConnectionProvider;
+  public ProgressUpdater(TaskInfoRepository taskInfoRepository, ParameterTool parameterTool, long completedFilesCount) {
+    taskInfoRepo = taskInfoRepository;
     this.taskId = parameterTool.getLong(JobParamName.TASK_ID);
     lastStoredFilesCount = completedFilesCount;
-    taskInfoRepo = RetryableMethodExecutor.createRetryProxy(new TaskInfoRepository(dbConnectionProvider));
     LOGGER.debug("Created ProgressUpdater");
   }
 
@@ -70,7 +68,6 @@ public void saveProgressInDB() {
   }
 
   public void close() {
-    dbConnectionProvider.close();
     LOGGER.debug("Closed: {}", ProgressUpdater.class.getSimpleName());
   }
 

commons/src/main/java/eu/europeana/processing/source/RegularDbReader.java

@@ -17,28 +17,20 @@
 public class RegularDbReader extends AbstractDbReader<ExecutionRecord> {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(RegularDbReader.class);
-  private ExecutionRecordRepository executionRecordRepository;
+  private final ExecutionRecordRepository executionRecordRepository;
 
   /**
    * Creates RegularDbReader
    *
    * @param context - Flink context
    * @param parameterTool - job parameters
    */
-  public RegularDbReader(SourceReaderContext context, ParameterTool parameterTool) {
+  public RegularDbReader(SourceReaderContext context, ParameterTool parameterTool, ExecutionRecordRepository executionRecordRepository) {
     super(context, parameterTool);
+    this.executionRecordRepository = executionRecordRepository;
     LOGGER.info("Created RegularDbReader");
   }
 
-  protected void createRepositories() {
-    //TODO Using retry proxy is maybe not optimal strategy in this case. This source implements asynchronous interface, so
-    // we could do this retries in poolNext() method by returning InputStatus.NOTHING_AVAILABLE, wait a bit and notify
-    // completable future to poll source again. Or simple wait a bit in pollNext() but only once per one retry.
-    // In such cases we would less block checkpointing mechanism, which should work smoothly in case of infrastructure problems
-    // and potential job restarts. And when we do not block we could do more retries or longer pauses.
-    executionRecordRepository = RetryableMethodExecutor.createRetryProxy(new ExecutionRecordRepository(dbConnectionProvider));
-  }
-
   protected List<ExecutionRecord> fetchRecords() throws IOException {
     return executionRecordRepository.getByDatasetIdAndExecutionIdAndOffsetAndLimit(
         parameterTool.getRequired(JobParamName.DATASET_ID),

commons/src/test/java/eu/europeana/processing/repository/TaskInfoRepositoryTest.java

@@ -10,7 +10,6 @@
 import org.assertj.db.type.AssertDbConnection;
 import org.assertj.db.type.AssertDbConnectionFactory;
 import org.assertj.db.type.Request.Builder;
-import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 
@@ -25,12 +24,6 @@ static void before() {
     prepareRequests();
   }
 
-  @AfterAll
-  static void after() {
-    postgres.stop();
-  }
-
-
   @Test
   void shouldSaveCorrectTaskInfo() {
     TaskInfoRepository taskInfoRepository = prepareRepository();
@@ -42,18 +35,21 @@ void shouldSaveCorrectTaskInfo() {
     assertThat(
         request.parameters(1).build()
     ).hasNumberOfRows(1)
+     .row(0).value(TaskInfoRepository.TASK_ID_COL_NAME).isEqualTo(1)
      .row(0).value(TaskInfoRepository.COMMIT_COUNT_COL_NAME).isEqualTo(12)
      .row(0).value(TaskInfoRepository.WRITE_COUNT_COL_NAME).isEqualTo(14);
 
     assertThat(
         request.parameters(2).build()
     ).hasNumberOfRows(1)
+     .row(0).value(TaskInfoRepository.TASK_ID_COL_NAME).isEqualTo(2)
      .row(0).value(TaskInfoRepository.COMMIT_COUNT_COL_NAME).isEqualTo(0)
      .row(0).value(TaskInfoRepository.WRITE_COUNT_COL_NAME).isEqualTo(10);
 
     assertThat(
         request.parameters(3).build()
     ).hasNumberOfRows(1)
+     .row(0).value(TaskInfoRepository.TASK_ID_COL_NAME).isEqualTo(3)
      .row(0).value(TaskInfoRepository.COMMIT_COUNT_COL_NAME).isEqualTo(20)
      .row(0).value(TaskInfoRepository.WRITE_COUNT_COL_NAME).isEqualTo(30);
   }

commons/src/test/java/eu/europeana/processing/source/ProgressUpdaterTest.java

@@ -0,0 +1,73 @@
+package eu.europeana.processing.source;
+
+import eu.europeana.processing.job.JobParamName;
+import eu.europeana.processing.model.TaskInfo;
+import eu.europeana.processing.repository.TaskInfoRepository;
+import org.apache.flink.util.ParameterTool;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mockito;
+
+class ProgressUpdaterTest {
+
+  @Test
+  void shouldSavePropperProgressInDbForFirstExecution() {
+    TaskInfoRepository taskInfoRepository = Mockito.mock(TaskInfoRepository.class);
+
+    try (ProgressUpdater progressUpdater = new ProgressUpdater(taskInfoRepository, ParameterTool.fromArgs(
+        new String[]{
+            "-" + JobParamName.TASK_ID, "12"
+        }
+    ), 100);) {
+      progressUpdater.saveProgressInDB();
+
+      ArgumentCaptor<TaskInfo> captor = ArgumentCaptor.forClass(TaskInfo.class);
+
+      Mockito.verify(taskInfoRepository, Mockito.times(1)).update(captor.capture());
+      TaskInfo emitted = captor.getValue();
+      Assertions.assertEquals(12, emitted.taskId());
+      Assertions.assertEquals(0, emitted.commitCount());
+      Assertions.assertEquals(-1, emitted.writeCount());
+    }
+  }
+
+  @Test
+  void shouldSavePropperProgressInDbForConsecutiveExecution() {
+    TaskInfoRepository taskInfoRepository = Mockito.mock(TaskInfoRepository.class);
+
+    try (ProgressUpdater progressUpdater = new ProgressUpdater(taskInfoRepository, ParameterTool.fromArgs(
+        new String[]{
+            "-" + JobParamName.TASK_ID, "12"
+        }
+    ), 80);) {
+
+      progressUpdater.snapshotEmittedFilesCount(100);
+      progressUpdater.saveProgressInDB();
+
+      ArgumentCaptor<TaskInfo> captor = ArgumentCaptor.forClass(TaskInfo.class);
+
+      Mockito.verify(taskInfoRepository, Mockito.times(1)).update(captor.capture());
+      TaskInfo emitted = captor.getValue();
+      Assertions.assertEquals(12, emitted.taskId());
+      Assertions.assertEquals(0, emitted.commitCount());
+      Assertions.assertEquals(100, emitted.writeCount());
+    }
+  }
+
+  @Test
+  void shouldNotSavePropperProgressInDb() {
+    TaskInfoRepository taskInfoRepository = Mockito.mock(TaskInfoRepository.class);
+
+    try (ProgressUpdater progressUpdater = new ProgressUpdater(taskInfoRepository, ParameterTool.fromArgs(
+        new String[]{
+            "-" + JobParamName.TASK_ID, "12"
+        }
+    ), 100);) {
+      progressUpdater.snapshotEmittedFilesCount(100);
+      progressUpdater.saveProgressInDB();
+
+      Mockito.verify(taskInfoRepository, Mockito.times(0)).update(Mockito.any());
+    }
+  }
+}