Fix yaml REST tests

Author: rockdaboot

x-pack/plugin/profiling/src/main/java/org/elasticsearch/xpack/profiling/action/GetStackTracesResponseBuilder.java

@@ -137,6 +137,22 @@ public void setTotalSamples(long totalSamples) {
     }
 
     public GetStackTracesResponse build() {
+        // Merge the TraceEvent data into the StackTraces.
+        if (stackTraces != null) {
+            for (Map.Entry<TraceEventID, TraceEvent> entry : stackTraceEvents.entrySet()) {
+                TraceEventID traceEventID = entry.getKey();
+                StackTrace stackTrace = stackTraces.get(traceEventID.stacktraceID());
+                if (stackTrace != null) {
+                    TraceEvent event = entry.getValue();
+                    if (event.subGroups != null) {
+                        stackTrace.subGroups = event.subGroups;
+                    }
+                    stackTrace.count += event.count;
+                    stackTrace.annualCO2Tons += event.annualCO2Tons;
+                    stackTrace.annualCostsUSD += event.annualCostsUSD;
+                }
+            }
+        }
         return new GetStackTracesResponse(stackTraces, stackFrames, executables, stackTraceEvents, totalFrames, samplingRate, totalSamples);
     }
 }

x-pack/plugin/profiling/src/main/java/org/elasticsearch/xpack/profiling/action/StackTrace.java

@@ -27,12 +27,18 @@ final class StackTrace implements ToXContentObject {
     String[] frameIds;
     int[] typeIds;
     SubGroup subGroups;
+    double annualCO2Tons;
+    double annualCostsUSD;
+    long count;
 
     StackTrace(int[] addressOrLines, String[] fileIds, String[] frameIds, int[] typeIds) {
         this.addressOrLines = addressOrLines;
         this.fileIds = fileIds;
         this.frameIds = frameIds;
         this.typeIds = typeIds;
+        annualCO2Tons = 0.0d;
+        annualCostsUSD = 0.0d;
+        count = 0;
     }
 
     private static final int BASE64_FRAME_ID_LENGTH = 32;
@@ -210,12 +216,15 @@ public void forNativeAndKernelFrames(Consumer<String> consumer) {
 
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
-        builder.startObject();
-        builder.field("address_or_lines", this.addressOrLines);
-        builder.field("file_ids", this.fileIds);
-        builder.field("frame_ids", this.frameIds);
-        builder.field("type_ids", this.typeIds);
-        builder.endObject();
+        builder.startObject()
+            .field("address_or_lines", this.addressOrLines)
+            .field("file_ids", this.fileIds)
+            .field("frame_ids", this.frameIds)
+            .field("type_ids", this.typeIds)
+            .field("annual_co2_tons", this.annualCO2Tons)
+            .field("annual_costs_usd", this.annualCostsUSD)
+            .field("count", this.count)
+            .endObject();
         return builder;
     }
 
@@ -230,7 +239,7 @@ public boolean equals(Object o) {
             && Arrays.equals(fileIds, that.fileIds)
             && Arrays.equals(frameIds, that.frameIds)
             && Arrays.equals(typeIds, that.typeIds);
-        // Don't compare metadata like subGroups.
+        // Don't compare metadata like annualized co2, annualized costs, subGroups and count.
     }
 
     // Don't hash metadata like annualized co2, annualized costs, subGroups and count.

x-pack/plugin/profiling/src/test/java/org/elasticsearch/xpack/profiling/action/StackTraceTests.java

@@ -86,6 +86,9 @@ public void testToXContent() throws IOException {
             .array("file_ids", "AAAAAAAAAAUAAAAAAAAB3g")
             .array("frame_ids", "AAAAAAAAAAUAAAAAAAAB3gAAAAAAD67u")
             .array("type_ids", new int[] { 2 })
+            .field("annual_co2_tons", 0.3d)
+            .field("annual_costs_usd", 2.7d)
+            .field("count", 1)
             .endObject();
 
         XContentBuilder actualRequest = XContentFactory.contentBuilder(contentType);
@@ -95,6 +98,9 @@ public void testToXContent() throws IOException {
             new String[] { "AAAAAAAAAAUAAAAAAAAB3gAAAAAAD67u" },
             new int[] { 2 }
         );
+        stackTrace.annualCO2Tons = 0.3d;
+        stackTrace.annualCostsUSD = 2.7d;
+        stackTrace.count = 1;
         stackTrace.toXContent(actualRequest, ToXContent.EMPTY_PARAMS);
 
         assertToXContentEquivalent(BytesReference.bytes(expectedRequest), BytesReference.bytes(actualRequest), contentType);