From f2d6d53c4484a920b763094a29e88491d3d502e7 Mon Sep 17 00:00:00 2001
From: Kondaka <krishkdk@amazon.com>
Date: Sun, 17 Nov 2024 23:00:26 -0800
Subject: [PATCH] Fixed unit tests

Signed-off-by: Kondaka <krishkdk@amazon.com>
---
 .../common/LambdaCommonHandlerTest.java       |  2 +-
 .../lambda/processor/LambdaProcessorTest.java | 19 +++++++++++++++++--
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/data-prepper-plugins/aws-lambda/src/test/java/org/opensearch/dataprepper/plugins/lambda/common/LambdaCommonHandlerTest.java b/data-prepper-plugins/aws-lambda/src/test/java/org/opensearch/dataprepper/plugins/lambda/common/LambdaCommonHandlerTest.java
index 5ff87d8cde..2c7f27654a 100644
--- a/data-prepper-plugins/aws-lambda/src/test/java/org/opensearch/dataprepper/plugins/lambda/common/LambdaCommonHandlerTest.java
+++ b/data-prepper-plugins/aws-lambda/src/test/java/org/opensearch/dataprepper/plugins/lambda/common/LambdaCommonHandlerTest.java
@@ -74,7 +74,7 @@ void testWaitForFutures() {
 
     LambdaCommonHandler.waitForFutures(futureList);
 
-    assertTrue(futureList.isEmpty());
+    assertFalse(futureList.isEmpty());
   }
 
   @Test
diff --git a/data-prepper-plugins/aws-lambda/src/test/java/org/opensearch/dataprepper/plugins/lambda/processor/LambdaProcessorTest.java b/data-prepper-plugins/aws-lambda/src/test/java/org/opensearch/dataprepper/plugins/lambda/processor/LambdaProcessorTest.java
index bcd2c6d4fb..15749b853e 100644
--- a/data-prepper-plugins/aws-lambda/src/test/java/org/opensearch/dataprepper/plugins/lambda/processor/LambdaProcessorTest.java
+++ b/data-prepper-plugins/aws-lambda/src/test/java/org/opensearch/dataprepper/plugins/lambda/processor/LambdaProcessorTest.java
@@ -24,6 +24,9 @@
 import org.opensearch.dataprepper.plugins.lambda.common.config.ClientOptions;
 import static org.opensearch.dataprepper.plugins.lambda.processor.LambdaProcessor.NUMBER_OF_RECORDS_FLUSHED_TO_LAMBDA_FAILED;
 import static org.opensearch.dataprepper.plugins.lambda.processor.LambdaProcessor.NUMBER_OF_RECORDS_FLUSHED_TO_LAMBDA_SUCCESS;
+import static org.opensearch.dataprepper.plugins.lambda.processor.LambdaProcessor.NUMBER_OF_SUCCESSFUL_REQUESTS_TO_LAMBDA;
+import static org.opensearch.dataprepper.plugins.lambda.processor.LambdaProcessor.NUMBER_OF_FAILED_REQUESTS_TO_LAMBDA;
+
 
 import io.micrometer.core.instrument.Counter;
 import io.micrometer.core.instrument.Timer;
@@ -107,8 +110,13 @@ public class LambdaProcessorTest {
     @Mock
     private Counter numberOfRecordsSuccessCounter;
 
+    @Mock
+    private Counter numberOfRequestsSuccessCounter;
+
     @Mock
     private Counter numberOfRecordsFailedCounter;
+    @Mock
+    private Counter numberOfRequestsFailedCounter;
 
     @Mock
     private InvokeResponse invokeResponse;
@@ -131,6 +139,10 @@ public void setUp() throws Exception {
             numberOfRecordsSuccessCounter);
         when(pluginMetrics.counter(eq(NUMBER_OF_RECORDS_FLUSHED_TO_LAMBDA_FAILED))).thenReturn(
             numberOfRecordsFailedCounter);
+        when(pluginMetrics.counter(eq(NUMBER_OF_SUCCESSFUL_REQUESTS_TO_LAMBDA))).thenReturn(
+                numberOfRecordsSuccessCounter);
+        when(pluginMetrics.counter(eq(NUMBER_OF_FAILED_REQUESTS_TO_LAMBDA))).thenReturn(
+                numberOfRecordsFailedCounter);
         when(pluginMetrics.timer(anyString())).thenReturn(lambdaLatencyMetric);
         when(pluginMetrics.gauge(anyString(), any(AtomicLong.class))).thenAnswer(
             invocation -> invocation.getArgument(1));
@@ -200,7 +212,10 @@ private void populatePrivateFields() throws Exception {
         // Use reflection to set the private fields
         setPrivateField(lambdaProcessor, "numberOfRecordsSuccessCounter",
             numberOfRecordsSuccessCounter);
+        setPrivateField(lambdaProcessor, "numberOfRequestsSuccessCounter",
+                numberOfRequestsSuccessCounter);
         setPrivateField(lambdaProcessor, "numberOfRecordsFailedCounter", numberOfRecordsFailedCounter);
+        setPrivateField(lambdaProcessor, "numberOfRequestsFailedCounter", numberOfRequestsFailedCounter);
         setPrivateField(lambdaProcessor, "tagsOnMatchFailure", tagsOnMatchFailure);
         setPrivateField(lambdaProcessor, "lambdaCommonHandler", lambdaCommonHandler);
     }
@@ -274,7 +289,7 @@ public void testDoExecute_WithEmptyResponse() throws Exception {
 
         // Assert
         assertEquals(0, result.size(), "Result should be empty due to empty Lambda response.");
-        //verify(numberOfRecordsSuccessCounter, times(1)).increment(1.0);
+        verify(numberOfRecordsSuccessCounter, times(1)).increment(1.0);
     }
 
     @Test
@@ -292,7 +307,7 @@ public void testDoExecute_WithNullResponse() throws Exception {
 
         // Assert
         assertEquals(0, result.size(), "Result should be empty due to null Lambda response.");
-        //verify(numberOfRecordsSuccessCounter, times(1)).increment(1.0);
+        verify(numberOfRecordsSuccessCounter, times(1)).increment(1.0);
     }
 
     @Test