Fix flaky PipelinesWithAcksIT by awaiting ack callback (#6718)

The acknowledgement callback is invoked asynchronously after sink
processing completes. Tests were asserting the ack result immediately
after confirming sink output, causing intermittent NullPointerException
when the callback had not yet fired.

Wrap ack assertions in await().untilAsserted() to poll until the
callback completes, matching the pattern already used in
three_pipelines_with_all_unrouted_records().

Signed-off-by: Sotaro Hikita <bering1814@gmail.com>

Author: lawofcycles

data-prepper-core/src/integrationTest/java/org/opensearch/dataprepper/integration/PipelinesWithAcksIT.java

@@ -30,9 +30,7 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
-import static org.junit.Assert.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 
 @FixMethodOrder()
 class PipelinesWithAcksIT {
@@ -81,7 +79,10 @@ void simple_pipeline_with_single_record() {
             assertThat(outputRecords, not(empty()));
             assertThat(outputRecords.size(), equalTo(numRecords));
         });
-        assertTrue(inMemorySourceAccessor.getAckReceived());
+        await().atMost(40000, TimeUnit.MILLISECONDS)
+                .untilAsserted(() -> {
+            assertThat(inMemorySourceAccessor.getAckReceived(), equalTo(true));
+        });
 
     }
 
@@ -97,7 +98,10 @@ void simple_pipeline_with_multiple_records() {
             assertThat(outputRecords, not(empty()));
             assertThat(outputRecords.size(), equalTo(numRecords));
         });
-        assertTrue(inMemorySourceAccessor.getAckReceived());
+        await().atMost(40000, TimeUnit.MILLISECONDS)
+                .untilAsserted(() -> {
+            assertThat(inMemorySourceAccessor.getAckReceived(), equalTo(true));
+        });
     }
 
     @Test
@@ -112,7 +116,10 @@ void two_pipelines_with_multiple_records() {
             assertThat(outputRecords, not(empty()));
             assertThat(outputRecords.size(), equalTo(numRecords));
         });
-        assertTrue(inMemorySourceAccessor.getAckReceived());
+        await().atMost(40000, TimeUnit.MILLISECONDS)
+                .untilAsserted(() -> {
+            assertThat(inMemorySourceAccessor.getAckReceived(), equalTo(true));
+        });
     }
 
     @Test
@@ -127,7 +134,10 @@ void three_pipelines_with_multiple_records() {
             assertThat(outputRecords, not(empty()));
             assertThat(outputRecords.size(), equalTo(numRecords));
         });
-        assertTrue(inMemorySourceAccessor.getAckReceived());
+        await().atMost(40000, TimeUnit.MILLISECONDS)
+                .untilAsserted(() -> {
+            assertThat(inMemorySourceAccessor.getAckReceived(), equalTo(true));
+        });
     }
 
     @Test
@@ -139,11 +149,10 @@ void three_pipelines_with_all_unrouted_records() {
         await().atMost(40000, TimeUnit.MILLISECONDS)
                 .untilAsserted(() -> {
                     assertNotNull(inMemorySourceAccessor);
-                    assertNotNull(inMemorySourceAccessor.getAckReceived());
+                    assertThat(inMemorySourceAccessor.getAckReceived(), equalTo(true));
         });
         List<Record<Event>> outputRecords = inMemorySinkAccessor.get(IN_MEMORY_IDENTIFIER);
         assertThat(outputRecords.size(), equalTo(0));
-        assertTrue(inMemorySourceAccessor.getAckReceived());
     }
 
     @Test
@@ -158,7 +167,10 @@ void three_pipelines_with_route_and_multiple_records() {
             assertThat(outputRecords, not(empty()));
             assertThat(outputRecords.size(), lessThanOrEqualTo(numRecords));
         });
-        assertThat(inMemorySourceAccessor.getAckReceived(), equalTo(true));
+        await().atMost(40000, TimeUnit.MILLISECONDS)
+                .untilAsserted(() -> {
+            assertThat(inMemorySourceAccessor.getAckReceived(), equalTo(true));
+        });
     }
 
     @Test
@@ -174,7 +186,10 @@ void three_pipelines_with_default_route_and_multiple_records() {
             assertThat(outputRecords, not(empty()));
             assertThat(outputRecords.size(), equalTo(2*numRecords));
         });
-        assertTrue(inMemorySourceAccessor.getAckReceived());
+        await().atMost(40000, TimeUnit.MILLISECONDS)
+                .untilAsserted(() -> {
+            assertThat(inMemorySourceAccessor.getAckReceived(), equalTo(true));
+        });
     }
 
     @Test
@@ -189,7 +204,10 @@ void two_parallel_pipelines_multiple_records() {
             assertThat(outputRecords, not(empty()));
             assertThat(outputRecords.size(), equalTo(2*numRecords));
         });
-        assertTrue(inMemorySourceAccessor.getAckReceived());
+        await().atMost(40000, TimeUnit.MILLISECONDS)
+                .untilAsserted(() -> {
+            assertThat(inMemorySourceAccessor.getAckReceived(), equalTo(true));
+        });
     }
 
     @Test
@@ -204,7 +222,10 @@ void three_pipelines_multi_sink_multiple_records() {
             assertThat(outputRecords, not(empty()));
             assertThat(outputRecords.size(), equalTo(3*numRecords));
         });
-        assertTrue(inMemorySourceAccessor.getAckReceived());
+        await().atMost(40000, TimeUnit.MILLISECONDS)
+                .untilAsserted(() -> {
+            assertThat(inMemorySourceAccessor.getAckReceived(), equalTo(true));
+        });
     }
 
     @Test
@@ -219,7 +240,10 @@ void one_pipeline_three_sinks_multiple_records() {
             assertThat(outputRecords, not(empty()));
             assertThat(outputRecords.size(), equalTo(3*numRecords));
         });
-        assertTrue(inMemorySourceAccessor.getAckReceived());
+        await().atMost(40000, TimeUnit.MILLISECONDS)
+                .untilAsserted(() -> {
+            assertThat(inMemorySourceAccessor.getAckReceived(), equalTo(true));
+        });
     }
 
     @Test
@@ -250,6 +274,9 @@ void one_pipeline_three_sinks_negative_ack_multiple_records() {
             assertThat(outputRecords, not(empty()));
             assertThat(outputRecords.size(), equalTo(3*numRecords));
         });
-        assertFalse(inMemorySourceAccessor.getAckReceived());
+        await().atMost(40000, TimeUnit.MILLISECONDS)
+                .untilAsserted(() -> {
+            assertThat(inMemorySourceAccessor.getAckReceived(), equalTo(false));
+        });
     }
 }