1 files changed, 14 insertions, 11 deletions

diff --git a/metrics-proxy/src/test/java/ai/vespa/metricsproxy/metric/model/json/YamasJsonModelTest.java b/metrics-proxy/src/test/java/ai/vespa/metricsproxy/metric/model/json/YamasJsonModelTest.java
index 540445fba5b..a5bc7a0877a 100644
--- a/metrics-proxy/src/test/java/ai/vespa/metricsproxy/metric/model/json/YamasJsonModelTest.java
+++ b/metrics-proxy/src/test/java/ai/vespa/metricsproxy/metric/model/json/YamasJsonModelTest.java
@@ -1,15 +1,17 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.metricsproxy.metric.model.json;
 
+import ai.vespa.metricsproxy.http.yamas.YamasResponse;
 import ai.vespa.metricsproxy.metric.model.MetricsPacket;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import org.junit.Test;
 
 import java.io.BufferedReader;
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Paths;
-import java.util.Collections;
+import java.util.List;
 
 import static ai.vespa.metricsproxy.metric.model.ConsumerId.toConsumerId;
 import static ai.vespa.metricsproxy.metric.model.MetricId.toMetricId;
@@ -32,10 +34,11 @@ public class YamasJsonModelTest {
     public void array_definition_creates_correct_json() throws IOException {
         YamasJsonModel jsonModel = getYamasJsonModel("yamas-array.json");
 
-        YamasArrayJsonModel yamasData = new YamasArrayJsonModel();
-        yamasData.add(jsonModel);
-
-        assertEquals(EXPECTED_JSON, yamasData.serialize());
+        try (ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) {
+            YamasResponse response = new YamasResponse(200, List.of(YamasJsonUtil.toMetricsPacketBuilder(jsonModel).build()));
+            response.render(outputStream);
+            assertEquals(EXPECTED_JSON, outputStream.toString());
+        }
     }
 
     @Test
@@ -48,10 +51,11 @@ public class YamasJsonModelTest {
         assertEquals(5.555555555E9, jsonModel.metrics.get("memory_rss"), 0.1d); //Not using custom double renderer
 
         // Serialize and verify
-        YamasArrayJsonModel yamasArray = new YamasArrayJsonModel();
-        yamasArray.add(jsonModel);
-        String string = yamasArray.serialize();
-        assertEquals(EXPECTED_JSON, string);
+        try (ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) {
+            YamasResponse response = new YamasResponse(200, List.of(YamasJsonUtil.toMetricsPacketBuilder(jsonModel).build()));
+            response.render(outputStream);
+            assertEquals(EXPECTED_JSON, outputStream.toString());
+        }
     }
 
     @Test
@@ -65,8 +69,7 @@ public class YamasJsonModelTest {
         assertEquals(5.555555555E9, metricsPacket.metrics().get(toMetricId("memory_rss")).doubleValue(), 0.1d); //Not using custom double rendrer
 
         // Serialize and verify
-        YamasArrayJsonModel yamasArray = YamasJsonUtil.toYamasArray(Collections.singleton(metricsPacket), true);
-        String string = yamasArray.serialize();
+        String string = YamasJsonUtil.toJson(List.of(metricsPacket), true);
         assertEquals(EXPECTED_JSON, string);
     }