From 45ee35b6ecf7de39649e2a7ceed379cc27cd7faf Mon Sep 17 00:00:00 2001
From: Andrea Marziali <andrea.marziali@datadoghq.com>
Date: Wed, 27 Nov 2024 14:20:01 +0100
Subject: [PATCH] fix comparator

---
 .../mule-4/src/test/groovy/mule4/MuleForkedTest.groovy       | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/dd-java-agent/instrumentation/mule-4/src/test/groovy/mule4/MuleForkedTest.groovy b/dd-java-agent/instrumentation/mule-4/src/test/groovy/mule4/MuleForkedTest.groovy
index 799ebae18ffa..9c2e1764568e 100644
--- a/dd-java-agent/instrumentation/mule-4/src/test/groovy/mule4/MuleForkedTest.groovy
+++ b/dd-java-agent/instrumentation/mule-4/src/test/groovy/mule4/MuleForkedTest.groovy
@@ -214,9 +214,10 @@ class MuleForkedTest extends WithHttpServer<MuleTestContainer> {
    */
   private static class TreeComparator implements Comparator<DDSpan> {
     private final Map<DDSpan, Long> levels
+    private final Map<Long, DDSpan> traceMap
 
     TreeComparator(List<DDSpan> trace) {
-      final Map<Long, DDSpan> traceMap =  trace.collectEntries { [(it.spanId): it] }
+      traceMap =  trace.collectEntries { [(it.spanId): it] }
       levels = trace.collectEntries({
         [(it): walkUp(traceMap, it, 0)]
       })
@@ -232,7 +233,7 @@ class MuleForkedTest extends WithHttpServer<MuleTestContainer> {
       if (o1.parentId == o2.parentId) {
         return o1.spanId <=> o2.spanId
       }
-      return o1.parentId <=> o2.parentId
+      return compare(traceMap.get(o1.parentId), traceMap.get(o2.parentId))
     }
 
     def walkUp(Map<Long, DDSpan> traceMap, DDSpan span, int size) {