Fix testApmIntegration histogram assertions (#115906)

Author: ldematte

test/external-modules/apm-integration/src/javaRestTest/java/org/elasticsearch/test/apmintegration/MetricsApmIT.java

@@ -30,7 +30,10 @@
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Predicate;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
+import static java.util.Map.entry;
 import static org.hamcrest.Matchers.closeTo;
 import static org.hamcrest.Matchers.equalTo;
 
@@ -57,27 +60,21 @@ protected String getTestRestCluster() {
 
     @SuppressWarnings("unchecked")
     public void testApmIntegration() throws Exception {
-        Map<String, Predicate<Map<String, Object>>> sampleAssertions = new HashMap<>(
+        Map<String, Predicate<Map<String, Object>>> valueAssertions = new HashMap<>(
             Map.ofEntries(
                 assertion("es.test.long_counter.total", m -> (Double) m.get("value"), closeTo(1.0, 0.001)),
                 assertion("es.test.double_counter.total", m -> (Double) m.get("value"), closeTo(1.0, 0.001)),
                 assertion("es.test.async_double_counter.total", m -> (Double) m.get("value"), closeTo(1.0, 0.001)),
                 assertion("es.test.async_long_counter.total", m -> (Integer) m.get("value"), equalTo(1)),
                 assertion("es.test.double_gauge.current", m -> (Double) m.get("value"), closeTo(1.0, 0.001)),
-                assertion("es.test.long_gauge.current", m -> (Integer) m.get("value"), equalTo(1)),
-                assertion(
-                    "es.test.double_histogram.histogram",
-                    m -> ((Collection<Integer>) m.get("counts")).stream().mapToInt(Integer::intValue).sum(),
-                    equalTo(2)
-                ),
-                assertion(
-                    "es.test.long_histogram.histogram",
-                    m -> ((Collection<Integer>) m.get("counts")).stream().mapToInt(Integer::intValue).sum(),
-                    equalTo(2)
-                )
+                assertion("es.test.long_gauge.current", m -> (Integer) m.get("value"), equalTo(1))
             )
         );
 
+        Map<String, Integer> histogramAssertions = new HashMap<>(
+            Map.ofEntries(entry("es.test.double_histogram.histogram", 2), entry("es.test.long_histogram.histogram", 2))
+        );
+
         CountDownLatch finished = new CountDownLatch(1);
 
         // a consumer that will remove the assertions from a map once it matched
@@ -89,15 +86,36 @@ public void testApmIntegration() throws Exception {
                 var metricset = (Map<String, Object>) apmMessage.get("metricset");
                 var samples = (Map<String, Object>) metricset.get("samples");
 
-                samples.entrySet().forEach(sampleEntry -> {
-                    var assertion = sampleAssertions.get(sampleEntry.getKey());// sample name
-                    if (assertion != null && assertion.test((Map<String, Object>) sampleEntry.getValue())) {// sample object
-                        sampleAssertions.remove(sampleEntry.getKey());
+                samples.forEach((key, value) -> {
+                    var valueAssertion = valueAssertions.get(key);// sample name
+                    if (valueAssertion != null) {
+                        logger.info("Matched {}:{}", key, value);
+                        var sampleObject = (Map<String, Object>) value;
+                        if (valueAssertion.test(sampleObject)) {// sample object
+                            logger.info("{} assertion PASSED", key);
+                            valueAssertions.remove(key);
+                        } else {
+                            logger.error("{} assertion FAILED", key);
+                        }
+                    }
+                    var histogramAssertion = histogramAssertions.get(key);
+                    if (histogramAssertion != null) {
+                        logger.info("Matched {}:{}", key, value);
+                        var samplesObject = (Map<String, Object>) value;
+                        var counts = ((Collection<Integer>) samplesObject.get("counts")).stream().mapToInt(Integer::intValue).sum();
+                        var remaining = histogramAssertion - counts;
+                        if (remaining == 0) {
+                            logger.info("{} assertion PASSED", key);
+                            histogramAssertions.remove(key);
+                        } else {
+                            logger.info("{} assertion PENDING: {} remaining", key, remaining);
+                            histogramAssertions.put(key, remaining);
+                        }
                     }
                 });
             }
 
-            if (sampleAssertions.isEmpty()) {
+            if (valueAssertions.isEmpty()) {
                 finished.countDown();
             }
         };
@@ -106,18 +124,18 @@ public void testApmIntegration() throws Exception {
 
         client().performRequest(new Request("GET", "/_use_apm_metrics"));
 
-        assertTrue(
-            "Timeout when waiting for assertions to complete. Remaining assertions to match: " + sampleAssertions,
-            finished.await(30, TimeUnit.SECONDS)
-        );
+        var completed = finished.await(30, TimeUnit.SECONDS);
+        var remainingAssertions = Stream.concat(valueAssertions.keySet().stream(), histogramAssertions.keySet().stream())
+            .collect(Collectors.joining());
+        assertTrue("Timeout when waiting for assertions to complete. Remaining assertions to match: " + remainingAssertions, completed);
     }
 
     private <T> Map.Entry<String, Predicate<Map<String, Object>>> assertion(
         String sampleKeyName,
         Function<Map<String, Object>, T> accessor,
         Matcher<T> expected
     ) {
-        return Map.entry(sampleKeyName, new Predicate<>() {
+        return entry(sampleKeyName, new Predicate<>() {
             @Override
             public boolean test(Map<String, Object> sampleObject) {
                 return expected.matches(accessor.apply(sampleObject));