Add MetricsVerifier to simplify verification of metric values in tests

core/src/test/java/org/apache/druid/java/util/metrics/MetricsVerifier.java

@@ -0,0 +1,98 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.java.util.metrics;
+
+import org.junit.Assert;
+
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Test utility to extract and verify metric values.
+ */
+public interface MetricsVerifier
+{
+  /**
+   * Verifies that no event has been emitted for the given metric.
+   */
+  default void verifyNotEmitted(String metricName)
+  {
+    verifyEmitted(metricName, 0);
+  }
+
+  /**
+   * Verifies that the metric was emitted the expected number of times.
+   */
+  default void verifyEmitted(String metricName, int times)
+  {
+    verifyEmitted(metricName, null, times);
+  }
+
+  /**
+   * Verifies that the metric was emitted for the given dimension filters the
+   * expected number of times.
+   */
+  default void verifyEmitted(String metricName, Map<String, Object> dimensionFilters, int times)
+  {
+    Assert.assertEquals(
+        "Metric was emitted unexpected number of times.",
+        times,
+        getMetricValues(metricName, dimensionFilters).size()
+    );
+  }
+
+  /**
+   * Verifies the value of the specified metric emitted in the previous run.
+   */
+  default void verifyValue(String metricName, Number expectedValue)
+  {
+    verifyValue(metricName, null, expectedValue);
+  }
+
+  /**
+   * Verifies the value of the event corresponding to the specified metric and
+   * dimensionFilters emitted in the previous run.
+   */
+  default void verifyValue(String metricName, Map<String, Object> dimensionFilters, Number expectedValue)
+  {
+    Assert.assertEquals(expectedValue, getValue(metricName, dimensionFilters));
+  }
+
+  /**
+   * Gets the value of the event corresponding to the specified metric and
+   * dimensionFilters.
+   */
+  default Number getValue(String metricName, Map<String, Object> dimensionFilters)
+  {
+    List<Number> values = getMetricValues(metricName, dimensionFilters);
+    Assert.assertEquals(
+        "Metric must have been emitted exactly once for the given dimensions.",
+        1,
+        values.size()
+    );
+    return values.get(0);
+  }
+
+  /**
+   * Gets the metric values for the specified dimension filters.
+   */
+  List<Number> getMetricValues(String metricName, Map<String, Object> dimensionFilters);
+
+}

core/src/test/java/org/apache/druid/java/util/metrics/StubServiceEmitter.java

@@ -24,12 +24,15 @@
 import org.apache.druid.java.util.emitter.service.ServiceMetricEvent;
 
 import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
-public class StubServiceEmitter extends ServiceEmitter
+public class StubServiceEmitter extends ServiceEmitter implements MetricsVerifier
 {
   private final List<Event> events = new ArrayList<>();
-  private final List<ServiceMetricEvent> metricEvents = new ArrayList<>();
+  private final Map<String, List<ServiceMetricEvent>> metricEvents = new HashMap<>();
 
   public StubServiceEmitter(String service, String host)
   {
@@ -40,7 +43,9 @@ public StubServiceEmitter(String service, String host)
   public void emit(Event event)
   {
     if (event instanceof ServiceMetricEvent) {
-      metricEvents.add((ServiceMetricEvent) event);
+      ServiceMetricEvent metricEvent = (ServiceMetricEvent) event;
+      metricEvents.computeIfAbsent(metricEvent.getMetric(), name -> new ArrayList<>())
+                  .add(metricEvent);
     }
     events.add(event);
   }
@@ -53,12 +58,29 @@ public List<Event> getEvents()
     return events;
   }
 
-  /**
-   * Gets all the metric events emitted since the previous {@link #flush()}.
-   */
-  public List<ServiceMetricEvent> getMetricEvents()
+  @Override
+  public List<Number> getMetricValues(
+      String metricName,
+      Map<String, Object> dimensionFilters
+  )
   {
-    return metricEvents;
+    final List<Number> values = new ArrayList<>();
+    final List<ServiceMetricEvent> events =
+        metricEvents.getOrDefault(metricName, Collections.emptyList());
+    final Map<String, Object> filters =
+        dimensionFilters == null ? Collections.emptyMap() : dimensionFilters;
+    for (ServiceMetricEvent event : events) {
+      final Map<String, Object> userDims = event.getUserDims();
+      boolean match = filters.keySet().stream()
+                             .map(d -> filters.get(d).equals(userDims.get(d)))
+                             .reduce((a, b) -> a && b)
+                             .orElse(true);
+      if (match) {
+        values.add(event.getValue());
+      }
+    }
+
+    return values;
   }
 
   @Override

processing/src/test/java/org/apache/druid/query/groupby/GroupByQueryRunnerTest.java

@@ -757,8 +757,7 @@ public void testGroupBy()
         query,
         serviceEmitter
     );
-    Assert.assertEquals(1, serviceEmitter.getEvents().size());
-    Assert.assertEquals(vectorize, serviceEmitter.getEvents().get(0).toMap().getOrDefault("vectorized", null));
+    serviceEmitter.verifyEmitted("query/wait/time", ImmutableMap.of("vectorized", vectorize), 1);
     TestHelper.assertExpectedObjects(expectedResults, results, "groupBy");
   }
 

processing/src/test/java/org/apache/druid/query/scan/ScanQueryRunnerTest.java

@@ -230,8 +230,7 @@ public void testFullOnSelect()
         0,
         3
     );
-    Assert.assertEquals(1, stubServiceEmitter.getEvents().size());
-    Assert.assertEquals(false, stubServiceEmitter.getEvents().get(0).toMap().getOrDefault("vectorized", null));
+    stubServiceEmitter.verifyEmitted("query/wait/time", ImmutableMap.of("vectorized", false), 1);
     verify(expectedResults, populateNullColumnAtLastForQueryableIndexCase(results, "null_column"));
   }
 

processing/src/test/java/org/apache/druid/query/timeseries/TimeseriesQueryRunnerTest.java

@@ -317,11 +317,7 @@ public void testFullOnTimeseries()
       ++count;
     }
 
-    Assert.assertEquals(1, stubServiceEmitter.getEvents().size());
-    Assert.assertEquals(
-        vectorize,
-        stubServiceEmitter.getEvents().get(0).toMap().getOrDefault("vectorized", null)
-    );
+    stubServiceEmitter.verifyEmitted("query/wait/time", ImmutableMap.of("vectorized", vectorize), 1);
     Assert.assertEquals(lastResult.toString(), expectedLast, lastResult.getTimestamp());
   }
 

server/src/test/java/org/apache/druid/server/coordinator/simulate/BalancingStrategiesTest.java

@@ -76,8 +76,8 @@ public void testFreshClusterGetsBalanced()
     runCoordinatorCycle();
     loadQueuedSegments();
     verifyValue(Metric.ASSIGNED_COUNT, 1000L);
-    verifyNoEvent(Metric.MOVED_COUNT);
-    verifyNoEvent(Metric.UNMOVED_COUNT);
+    verifyNotEmitted(Metric.MOVED_COUNT);
+    verifyNotEmitted(Metric.UNMOVED_COUNT);
 
     for (DruidServer historical : historicals) {
       Assert.assertEquals(200, historical.getTotalSegments());
@@ -113,8 +113,8 @@ public void testClusterGetsBalancedWhenServerIsAdded()
     runCoordinatorCycle();
     loadQueuedSegments();
     verifyValue(Metric.ASSIGNED_COUNT, 1000L);
-    verifyNoEvent(Metric.MOVED_COUNT);
-    verifyNoEvent(Metric.UNMOVED_COUNT);
+    verifyNotEmitted(Metric.MOVED_COUNT);
+    verifyNotEmitted(Metric.UNMOVED_COUNT);
 
     // Verify that each server is equally loaded
     for (DruidServer historical : historicals) {
@@ -161,8 +161,8 @@ public void testClusterGetsBalancedWhenServerIsRemoved()
     runCoordinatorCycle();
     loadQueuedSegments();
     verifyValue(Metric.ASSIGNED_COUNT, 1000L);
-    verifyNoEvent(Metric.MOVED_COUNT);
-    verifyNoEvent(Metric.UNMOVED_COUNT);
+    verifyNotEmitted(Metric.MOVED_COUNT);
+    verifyNotEmitted(Metric.UNMOVED_COUNT);
 
     // Verify that each server is equally loaded
     for (DruidServer historical : historicals) {

server/src/test/java/org/apache/druid/server/coordinator/simulate/CoordinatorSimulation.java

@@ -20,7 +20,7 @@
 package org.apache.druid.server.coordinator.simulate;
 
 import org.apache.druid.client.DruidServer;
-import org.apache.druid.java.util.emitter.service.ServiceMetricEvent;
+import org.apache.druid.java.util.metrics.MetricsVerifier;
 import org.apache.druid.server.coordinator.CoordinatorDynamicConfig;
 import org.apache.druid.timeline.DataSegment;
 
@@ -81,9 +81,10 @@ interface CoordinatorState
     DruidServer getInventoryView(String serverName);
 
     /**
-     * Returns the metric events emitted in the previous coordinator run.
+     * Returns a MetricsVerifier which can be used to extract and verify the
+     * metric values emitted in the previous coordinator run.
      */
-    List<ServiceMetricEvent> getMetricEvents();
+    MetricsVerifier getMetricsVerifier();
 
     /**
      * Gets the load percentage of the specified datasource as seen by the coordinator.