Allow inspecting consumers in MetricsPacket.Builder

2 files changed, 17 insertions, 2 deletions

diff --git a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/MetricsPacket.java b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/MetricsPacket.java
index 8ecf57237ef..8d5a1f50918 100644
--- a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/MetricsPacket.java
+++ b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/MetricsPacket.java
@@ -179,6 +179,10 @@ public class MetricsPacket {
             return this;
         }
 
+        public boolean hasConsumer(ConsumerId id) {
+            return consumers.contains(id);
+        }
+
         public MetricsPacket build() {
             return new MetricsPacket(statusCode, statusMessage, timestamp, service, metrics, dimensions, consumers);
         }
diff --git a/metrics-proxy/src/test/java/ai/vespa/metricsproxy/metric/model/MetricsPacketTest.java b/metrics-proxy/src/test/java/ai/vespa/metricsproxy/metric/model/MetricsPacketTest.java
index 9b37a805245..78c80689299 100644
--- a/metrics-proxy/src/test/java/ai/vespa/metricsproxy/metric/model/MetricsPacketTest.java
+++ b/metrics-proxy/src/test/java/ai/vespa/metricsproxy/metric/model/MetricsPacketTest.java
@@ -14,6 +14,7 @@ import java.util.Map;
 import static ai.vespa.metricsproxy.metric.model.ConsumerId.toConsumerId;
 import static ai.vespa.metricsproxy.metric.model.MetricId.toMetricId;
 import static ai.vespa.metricsproxy.metric.model.ServiceId.toServiceId;
+import static java.util.Collections.singleton;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
@@ -45,13 +46,23 @@ public class MetricsPacketTest {
         MetricsPacket packet = new MetricsPacket.Builder(toServiceId("foo"))
                 .statusCode(0)
                 .statusMessage("")
-                .addConsumers(Collections.singleton(DUPLICATE_CONSUMER))
-                .addConsumers(Collections.singleton(DUPLICATE_CONSUMER))
+                .addConsumers(singleton(DUPLICATE_CONSUMER))
+                .addConsumers(singleton(DUPLICATE_CONSUMER))
                 .build();
         assertEquals(1, packet.consumers().size());
     }
 
     @Test
+    public void builder_allows_inspecting_consumers() {
+        var consumer = toConsumerId("my-consumer");
+        var builder = new MetricsPacket.Builder(toServiceId("foo"))
+                .statusCode(0)
+                .statusMessage("")
+                .addConsumers(singleton(consumer));
+        assertTrue(builder.hasConsumer(consumer));
+    }
+
+    @Test
     public void builder_can_retain_subset_of_metrics() {
         MetricsPacket packet = new MetricsPacket.Builder(toServiceId("foo"))
                 .putMetrics(ImmutableList.of(