add failure counter unit tests

Signed-off-by: Haidong <[email protected]>
wanghd89 · Nov 14, 2023 · 0339a72 · 0339a72
1 parent c600cf4
commit 0339a72
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/...org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBSnapshotWorker.java b/...org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBSnapshotWorker.java
@@ -158,7 +158,7 @@ private void startProcessPartition(SourcePartition<MongoDBSnapshotProgressState>
                     }
                     successItemsCounter.increment();
                     successRecords += 1;
-                } catch (JsonProcessingException e) {
+                } catch (Exception e) {
                     LOG.error("failed to add record to buffer with error {}", e.getMessage());
                     failureItemsCounter.increment();
                     failedRecords += 1;

diff --git a/...opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBSnapshotWorkerTest.java b/...opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBSnapshotWorkerTest.java
@@ -172,6 +172,18 @@ public void test_shouldGiveUpPartitionIfExceptionOccurred() throws InterruptedEx
         verify(sourceCoordinator, times(1)).giveUpPartitions();
     }
 
+    @Test
+    public void test_shouldCountFailureIfBufferFailed() throws Exception {
+        doThrow(new RuntimeException("")).when(buffer).write(any(), anyInt());
+        this.mockDependencyAndProcessPartition("test.collection|0|1|java.lang.Integer", false);
+        final ArgumentCaptor<MongoDBSnapshotProgressState> progressStateCapture = ArgumentCaptor.forClass(MongoDBSnapshotProgressState.class);
+        verify(sourceCoordinator, times(1)).saveProgressStateForPartition(anyString(), progressStateCapture.capture());
+        List<MongoDBSnapshotProgressState> progressStates = progressStateCapture.getAllValues();
+        assertThat(progressStates.get(0).getTotal(), is(2L));
+        assertThat(progressStates.get(0).getSuccess(), is(0L));
+        assertThat(progressStates.get(0).getFailed(), is(2L));
+    }
+
     @Test
     public void test_shouldThreadSleepIfExceptionOccurred() throws InterruptedException {
         doThrow(new RuntimeException("")).when(sourceCoordinator).getNextPartition(mongoDBPartitionCreationSupplier);