Prevent message collection from being updated after message count has been received

.github/workflows/ci.yml

@@ -39,7 +39,7 @@ jobs:
       run: ./gradlew clean build -Dbuild.snapshot=false -x test -x integrationTest
 
     - name: Test
-      run: OPENDISTRO_SECURITY_TEST_OPENSSL_OPT=true ./gradlew test integrationTest -i
+      run: OPENDISTRO_SECURITY_TEST_OPENSSL_OPT=true ./gradlew test integrationTest
 
     - name: Coverage
       uses: codecov/codecov-action@v1

src/test/java/org/opensearch/security/auditlog/integration/BasicAuditlogTest.java

@@ -133,7 +133,7 @@ public void testSSLPlainText() throws Exception {
             final RuntimeException ex = Assert.assertThrows(RuntimeException.class,
                 () -> nonSslRestHelper().executeGetRequest("_search", encodeBasicHeader("admin", "admin")));
             Assert.assertEquals("org.apache.hc.core5.http.NoHttpResponseException", ex.getCause().getClass().getName());
-        }, 1);
+        }, 2);
 
         // All of the messages should be the same as the http client is attempting multiple times.
         messages.stream().forEach((message) -> {

src/test/java/org/opensearch/security/auditlog/integration/TestAuditlogImpl.java

@@ -58,30 +58,66 @@ public static synchronized void clear() {
      * Perform an action and then wait until the expected number of messages have been found.
      */
     public static List<AuditMessage> doThenWaitForMessages(final Runnable action, final int expectedCount) {
-        final CountDownLatch latch = new CountDownLatch(expectedCount);
+        final List<AuditMessage> missedMessages = new ArrayList<>();
         final List<AuditMessage> messages = new ArrayList<>();
-        countDownRef.set(latch);
-        messagesRef.set(messages);
-
-        TestAuditlogImpl.sb = new StringBuffer();
-        TestAuditlogImpl.messages = messages; 
+        final CountDownLatch latch = resetAuditStorage(expectedCount, messages);
 
         try {
             action.run();
             final int maxSecondsToWaitForMessages = 1;
-            final boolean foundAll = latch.await(maxSecondsToWaitForMessages, TimeUnit.SECONDS);
-            if (!foundAll) {
+            boolean foundAll = false;
+            foundAll = latch.await(maxSecondsToWaitForMessages, TimeUnit.SECONDS);
+            // After the wait has prevent any new messages from being recieved
+            resetAuditStorage(0, missedMessages);
+            if (!foundAll || messages.size() != expectedCount) {
                 throw new MessagesNotFoundException(expectedCount, messages);
             }
-            if (messages.size() != expectedCount) {
-                throw new RuntimeException("Unexpected number of messages, was expecting " + expectedCount + ", received " + messages.size());
-            }
         } catch (final InterruptedException e) {
             throw new RuntimeException("Unexpected exception", e);
         }
+
+        // Do not check for missed messages if no messages were expected
+        if (expectedCount != 0) {
+            try {
+                Thread.sleep(100);
+                if (missedMessages.size() != 0) {
+                    final String missedMessagesErrorMessage = new StringBuilder()
+                    .append("Audit messages were missed! ")
+                    .append("Found " + (missedMessages.size()) + " messages.")
+                    .append("Messages found during this time: \n\n")
+                    .append(missedMessages.stream()
+                        .map(AuditMessage::toString)
+                        .collect(Collectors.joining("\n")))
+                    .toString();
+
+                    throw new RuntimeException(missedMessagesErrorMessage);
+                }
+            } catch (final Exception e) {
+                throw new RuntimeException("Unexpected exception", e);
+            }
+        }
+
+        // Next usage of this class might be using raw stringbuilder / list so reset before that test might run
+        resetAuditStorage(0, new ArrayList<>());
         return new ArrayList<>(messages);
     }
 
+    /**
+     * Resets all of the mechanics for fresh messages to be captured
+     * 
+     * @param expectedMessageCount The number of messages before the latch is signalled, indicating all messages have been recieved
+     * @param message Where messages will be stored after being recieved
+     */
+    private static CountDownLatch resetAuditStorage(int expectedMessageCount, List<AuditMessage> messages) {
+        final CountDownLatch latch = new CountDownLatch(expectedMessageCount);
+        countDownRef.set(latch);
+        messagesRef.set(messages);
+
+        TestAuditlogImpl.sb = new StringBuffer();
+        TestAuditlogImpl.messages = messages;
+        return latch;
+    }
+
     /**
      * Perform an action and then wait until a single message has been found.
      */