Add method finishedTraces to MockTracer. (#359)

* Add method finishedTraces to MockTracer. Signed-off-by: qudongfang <qudongfang@gmail.com>
opentracing · Aug 27, 2019 · 2df2a89 · 2df2a89
1 parent 70555ed
commit 2df2a89
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 0 deletions.
diff --git a/opentracing-mock/src/main/java/io/opentracing/mock/MockTracer.java b/opentracing-mock/src/main/java/io/opentracing/mock/MockTracer.java
@@ -25,6 +25,7 @@
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -95,6 +96,28 @@ public synchronized List<MockSpan> finishedSpans() {
         return new ArrayList<>(this.finishedSpans);
     }
 
+    /**
+     * @return all finish()ed Traces(Spans) started by this MockTracer grouped by traceId and spanId in HashMap format.
+     */
+    public synchronized Map<String, Map<String, MockSpan>> finishedTraces() {
+        Map<String, Map<String, MockSpan>> result = new LinkedHashMap<>();
+
+        for (MockSpan span: this.finishedSpans) {
+            String traceId = span.context().toTraceId();
+
+            Map<String, MockSpan> spanId2Span = result.get(traceId);
+            if (null == spanId2Span) {
+                spanId2Span = new LinkedHashMap<>();
+                result.put(traceId, spanId2Span);
+            }
+
+            String spanId = span.context().toSpanId();
+            spanId2Span.put(spanId, span);
+        }
+
+        return result;
+    }
+
     /**
      * Noop method called on {@link Span#finish()}.
      */

diff --git a/opentracing-mock/src/test/java/io/opentracing/mock/MockTracerTest.java b/opentracing-mock/src/test/java/io/opentracing/mock/MockTracerTest.java
@@ -22,6 +22,7 @@
 import io.opentracing.propagation.TextMapAdapter;
 import java.nio.ByteBuffer;
 import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -114,6 +115,30 @@ public void testChildSpan() {
 
     }
 
+    @Test
+    public void testFinishedTraces() {
+        final int TRACE_CNT = 3;
+
+        Map<String, Map<String, MockSpan>> expect = new LinkedHashMap<>();
+
+        try(MockTracer tracer = new MockTracer()) {
+            for (int i = 0; i < TRACE_CNT; i++) {
+                MockSpan parent = tracer.buildSpan("parent").withStartTimestamp(1000).start();
+                MockSpan child = tracer.buildSpan("child").withStartTimestamp(1100).asChildOf(parent).start();
+                child.finish(1900);
+                parent.finish(2000);
+
+                Map<String, MockSpan> spans = new LinkedHashMap<>();
+                spans.put(parent.context().toSpanId(), parent);
+                spans.put(child.context().toSpanId(), child);
+
+                expect.put(parent.context().toTraceId(), spans);
+            }
+
+            Assert.assertTrue(expect.equals(tracer.finishedTraces()));
+        }
+    }
+
     @Test
     public void testStartTimestamp() throws InterruptedException {
         MockTracer tracer = new MockTracer();