WIP

Author: PerfectSlayer

dd-smoke-tests/src/main/java/datadog/smoketest/trace/SpanMatcher.java

@@ -1,6 +1,5 @@
 package datadog.smoketest.trace;
 
-import static datadog.trace.test.junit.utils.assertions.Matchers.any;
 import static datadog.trace.test.junit.utils.assertions.Matchers.assertValue;
 import static datadog.trace.test.junit.utils.assertions.Matchers.is;
 import static datadog.trace.test.junit.utils.assertions.Matchers.isFalse;
@@ -39,15 +38,20 @@ public final class SpanMatcher {
   private Matcher<CharSequence> resourceNameMatcher;
   private Matcher<String> typeMatcher;
   private Matcher<Boolean> errorMatcher;
-  private Long parentId; // explicit expected parent id (childOf/root); null => not checked by id
-  private int parentSpanIndex = -1; // >= 0 => parent is the span at this index in the sorted trace
-  private boolean childOfPrevious; // parent is the preceding span in the sorted trace
-  private final Map<String, Matcher<String>> metaMatchers = new HashMap<>();
-  private final Map<String, Matcher<Number>> metricMatchers = new HashMap<>();
-  // meta_struct values are nested (map/list/scalar), so matched as Object.
-  private final Map<String, Matcher<?>> metaStructMatchers = new HashMap<>();
+  private Matcher<Long> parentIdMatcher;
+  private int parentSpanIndex;
+  private final Map<String, Matcher<String>> metaMatchers;
+  private final Map<String, Matcher<Number>> metricMatchers;
+  private final Map<String, Matcher<?>> metaStructMatchers;
 
-  private SpanMatcher() {}
+  private static final Matcher<Long> CHILD_OF_PREVIOUS_MATCHER = is(0L);
+
+  private SpanMatcher() {
+    this.parentSpanIndex = -1;
+    this.metaMatchers = new HashMap<>();
+    this.metricMatchers = new HashMap<>();
+    this.metaStructMatchers = new HashMap<>();
+  }
 
   /** Creates a new span matcher; configure it with the fluent methods below. */
   public static SpanMatcher span() {
@@ -97,32 +101,28 @@ public SpanMatcher error(boolean errored) {
 
   /** Matches a root span (no parent). */
   public SpanMatcher root() {
-    this.parentId = 0L;
+    this.parentIdMatcher = is(0L);
     this.parentSpanIndex = -1;
-    this.childOfPrevious = false;
     return this;
   }
 
   /** Matches a span whose parent is the given span id. */
   public SpanMatcher childOf(long parentSpanId) {
-    this.parentId = parentSpanId;
+    this.parentIdMatcher = is(parentSpanId);
     this.parentSpanIndex = -1;
-    this.childOfPrevious = false;
     return this;
   }
 
   /** Matches a span whose parent is the span at {@code parentSpanIndex} in the (sorted) trace. */
   public SpanMatcher childOfIndex(int parentSpanIndex) {
+    this.parentIdMatcher = null;
     this.parentSpanIndex = parentSpanIndex;
-    this.parentId = null;
-    this.childOfPrevious = false;
     return this;
   }
 
   /** Matches a span whose parent is the immediately preceding span in the (sorted) trace. */
   public SpanMatcher childOfPrevious() {
-    this.childOfPrevious = true;
-    this.parentId = null;
+    this.parentIdMatcher = CHILD_OF_PREVIOUS_MATCHER;
     this.parentSpanIndex = -1;
     return this;
   }
@@ -151,16 +151,25 @@ public SpanMatcher metaStruct(String name, Matcher<?> matcher) {
   }
 
   void assertSpan(List<DecodedSpan> trace, int spanIndex) {
-    assertValue(
-        parentIdMatcher(trace, spanIndex),
-        trace.get(spanIndex).getParentId(),
-        "Unexpected parent id");
+    // Apply parent span index
+    if (this.parentSpanIndex >= 0) {
+      this.parentIdMatcher = is(trace.get(this.parentSpanIndex).getSpanId());
+    }
+    // Apply parent id matcher from the previous span
+    else if (this.parentIdMatcher == CHILD_OF_PREVIOUS_MATCHER) {
+      if (spanIndex == 0) {
+        throw new IllegalStateException("Cannot use childOfPrevious() matcher on the first span");
+      }
+      DecodedSpan previousSpan = trace.get(spanIndex - 1);
+      this.parentIdMatcher = is(previousSpan.getSpanId());
+    }
     DecodedSpan span = trace.get(spanIndex);
     assertValue(this.serviceMatcher, span.getService(), "Unexpected service name");
     assertValue(this.operationNameMatcher, span.getName(), "Unexpected operation name");
     assertValue(this.resourceNameMatcher, span.getResource(), "Unexpected resource name");
     assertValue(this.typeMatcher, span.getType(), "Unexpected span type");
     assertValue(this.errorMatcher, span.getError() != 0, "Unexpected error status");
+    assertValue(this.parentIdMatcher, span.getParentId(), "Unexpected parent id");
     assertSpanTags(span.getMeta());
     assertSpanMetrics(span.getMetrics());
     assertSpanMetaStruct(span.getMetaStruct());
@@ -190,17 +199,4 @@ private void assertSpanMetaStruct(Map<String, Object> metaStruct) {
           (Matcher<Object>) entry.getValue(), metaStruct.get(key), "Unexpected meta_struct '" + key + "'");
     }
   }
-
-  private Matcher<Long> parentIdMatcher(List<DecodedSpan> trace, int spanIndex) {
-    if (this.parentSpanIndex >= 0) {
-      return is(trace.get(parentSpanIndex).getSpanId());
-    }
-    if (this.childOfPrevious) {
-      if (spanIndex == 0) {
-        throw new IllegalStateException("childOfPrevious() cannot be used on the first span");
-      }
-      return is(trace.get(spanIndex - 1).getSpanId());
-    }
-    return this.parentId == null ? any() : is(this.parentId);
-  }
 }