taking the coverage to 100%

Signed-off-by: Santhosh Gandhe <[email protected]>
san81 · Oct 23, 2024 · 5a4e966 · 5a4e966
1 parent 647ba10
commit 5a4e966
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/...h/dataprepper/plugins/source/saas_crawler/coordination/scheduler/LeaderSchedulerTest.java b/...h/dataprepper/plugins/source/saas_crawler/coordination/scheduler/LeaderSchedulerTest.java
@@ -19,6 +19,7 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.BDDMockito.given;
+import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoInteractions;
@@ -55,6 +56,7 @@ void testLeaderPartitionsCreation(boolean initializationState) throws Interrupte
         leaderPartition.getProgressState().get().setInitialized(initializationState);
         leaderPartition.getProgressState().get().setLastPollTime(0L);
         given(coordinator.acquireAvailablePartition(LeaderPartition.PARTITION_TYPE)).willReturn(Optional.of(leaderPartition));
+        doThrow(RuntimeException.class).when(coordinator).saveProgressStateForPartition(any(LeaderPartition.class), any(Duration.class));
 
         ExecutorService executorService = Executors.newSingleThreadExecutor();
         executorService.submit(leaderScheduler);

diff --git a/...h/dataprepper/plugins/source/saas_crawler/coordination/scheduler/WorkerSchedulerTest.java b/...h/dataprepper/plugins/source/saas_crawler/coordination/scheduler/WorkerSchedulerTest.java
@@ -80,7 +80,28 @@ void testLeaderPartitionsCreation() throws InterruptedException {
         SaasWorkerProgressState stateObj = (SaasWorkerProgressState)sourcePartition.getProgressState().get();
         verify(crawler, atLeast(1)).executePartition(stateObj, buffer, sourceConfig);
         verify(coordinator, atLeast(1)).completePartition(eq(sourcePartition));
+    }
+
+    @Test
+    void testEmptyProgressState() throws InterruptedException {
+        WorkerScheduler workerScheduler = new WorkerScheduler(buffer, coordinator, sourceConfig, crawler);
+
+        String sourceId =  UUID.randomUUID() + "|" + SaasSourcePartition.PARTITION_TYPE;
+        when(sourcePartitionStoreItem.getPartitionProgressState()).thenReturn(null);
+        when(sourcePartitionStoreItem.getSourceIdentifier()).thenReturn(sourceId);
+        PartitionFactory factory = new PartitionFactory();
+        EnhancedSourcePartition sourcePartition = factory.apply(sourcePartitionStoreItem);
+        given(coordinator.acquireAvailablePartition(SaasSourcePartition.PARTITION_TYPE)).willReturn(Optional.of(sourcePartition));
 
+        ExecutorService executorService = Executors.newSingleThreadExecutor();
+        executorService.submit(workerScheduler);
+
+        Thread.sleep(50);
+        executorService.shutdownNow();
+
+        // Check if crawler was invoked and updated leader lease renewal time
+        verifyNoInteractions(crawler);
+        verify(coordinator, atLeast(1)).completePartition(eq(sourcePartition));
     }
 
     @Test