diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/annotation/SleuthSpanCreatorAspectFluxTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/annotation/SleuthSpanCreatorAspectFluxTests.java index 02406cbba..88a00b88e 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/annotation/SleuthSpanCreatorAspectFluxTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/annotation/SleuthSpanCreatorAspectFluxTests.java @@ -16,15 +16,27 @@ package org.springframework.cloud.sleuth.annotation; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicReference; +import java.util.stream.Collectors; + import brave.Span; import brave.Tracer; import brave.sampler.Sampler; import org.apache.commons.lang3.StringUtils; -import org.assertj.core.api.BDDAssertions; import org.awaitility.Awaitility; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.util.context.Context; +import zipkin2.Annotation; +import zipkin2.reporter.Reporter; + import org.springframework.beans.factory.annotation.Autowired; import org.springframework.boot.autoconfigure.EnableAutoConfiguration; import org.springframework.boot.test.context.SpringBootTest; @@ -32,17 +44,6 @@ import org.springframework.cloud.sleuth.util.ArrayListSpanReporter; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; import org.springframework.test.context.junit4.SpringJUnit4ClassRunner; -import reactor.core.publisher.Flux; -import reactor.core.publisher.Mono; -import zipkin2.Annotation; -import zipkin2.reporter.Reporter; - -import java.util.ArrayList; -import java.util.Iterator; -import java.util.List; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.atomic.AtomicReference; -import java.util.stream.Collectors; import static org.assertj.core.api.BDDAssertions.then; import static org.springframework.cloud.sleuth.annotation.SleuthSpanCreatorAspectFluxTests.TestBean.TEST_STRING1; @@ -213,7 +214,7 @@ public class SleuthSpanCreatorAspectFluxTests { then(spans.get(0).tags()).containsEntry("customTestTag10", "test"); then(spans.get(0).annotations().stream().map(Annotation::value) .collect(Collectors.toList())).contains("customTest.before", - "customTest.after"); + "customTest.after"); then(spans.get(0).duration()).isNotZero(); then(this.tracer.currentSpan()).isNull(); }); @@ -231,7 +232,7 @@ public class SleuthSpanCreatorAspectFluxTests { then(spans.get(0).tags()).containsEntry("customTestTag10", "test"); then(spans.get(0).annotations().stream().map(Annotation::value) .collect(Collectors.toList())).contains("customTest.before", - "customTest.after"); + "customTest.after"); then(spans.get(0).duration()).isNotZero(); then(this.tracer.currentSpan()).isNull(); }); @@ -257,7 +258,7 @@ public class SleuthSpanCreatorAspectFluxTests { then(spans.get(0).tags()).containsEntry("customTestTag10", "test"); then(spans.get(0).annotations().stream().map(Annotation::value) .collect(Collectors.toList())).contains("customTest.before", - "customTest.after"); + "customTest.after"); then(spans.get(0).duration()).isNotZero(); then(this.tracer.currentSpan()).isNull(); }); @@ -286,7 +287,7 @@ public class SleuthSpanCreatorAspectFluxTests { .containsEntry("customTestTag11", "test"); then(spans.get(0).annotations().stream().map(Annotation::value) .collect(Collectors.toList())).contains("customTest.before", - "customTest.after"); + "customTest.after"); then(spans.get(0).duration()).isNotZero(); then(this.tracer.currentSpan()).isNull(); }); @@ -341,7 +342,7 @@ public class SleuthSpanCreatorAspectFluxTests { then(spans.get(0).tags()).containsEntry("error", "test exception 13"); then(spans.get(0).annotations().stream().map(Annotation::value) .collect(Collectors.toList())).contains("testMethod13.before", - "testMethod13.afterFailure", "testMethod13.after"); + "testMethod13.afterFailure", "testMethod13.after"); then(spans.get(0).duration()).isNotZero(); then(this.tracer.currentSpan()).isNull(); }); @@ -589,13 +590,14 @@ public class SleuthSpanCreatorAspectFluxTests { @Override public Flux newSpanInTraceContext() { - return Flux.defer(() -> Flux.just(id(tracer))); + Long id = id(tracer); + return Flux.defer(() -> Flux.just(id)); } @Override public Flux newSpanInSubscriberContext() { - return Mono.subscriberContext().flatMapMany( - context -> Flux.just(id(tracer))); + return Mono.subscriberContext() + .flatMapMany(context -> Flux.just(id(context, tracer))); } } @@ -607,6 +609,13 @@ public class SleuthSpanCreatorAspectFluxTests { return tracer.currentSpan().context().spanId(); } + protected static Long id(Context context, Tracer tracer) { + if (context.hasKey(Span.class)) { + return context.get(Span.class).context().spanId(); + } + return id(tracer); + } + @Configuration @EnableAutoConfiguration protected static class TestConfiguration { diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/annotation/SleuthSpanCreatorAspectMonoTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/annotation/SleuthSpanCreatorAspectMonoTests.java index e409d99e6..cc333ffbb 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/annotation/SleuthSpanCreatorAspectMonoTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/annotation/SleuthSpanCreatorAspectMonoTests.java @@ -230,7 +230,7 @@ public class SleuthSpanCreatorAspectMonoTests { then(spans.get(0).tags()).containsEntry("customTestTag10", "test"); then(spans.get(0).annotations().stream().map(Annotation::value) .collect(Collectors.toList())).contains("customTest.before", - "customTest.after"); + "customTest.after"); then(spans.get(0).duration()).isNotZero(); then(this.tracer.currentSpan()).isNull(); }); @@ -247,7 +247,7 @@ public class SleuthSpanCreatorAspectMonoTests { then(spans.get(0).tags()).containsEntry("customTestTag10", "test"); then(spans.get(0).annotations().stream().map(Annotation::value) .collect(Collectors.toList())).contains("customTest.before", - "customTest.after"); + "customTest.after"); then(spans.get(0).duration()).isNotZero(); then(this.tracer.currentSpan()).isNull(); }); @@ -275,7 +275,7 @@ public class SleuthSpanCreatorAspectMonoTests { then(spans.get(0).tags()).containsEntry("customTestTag10", "test"); then(spans.get(0).annotations().stream().map(Annotation::value) .collect(Collectors.toList())).contains("customTest.before", - "customTest.after"); + "customTest.after"); then(spans.get(0).duration()).isNotZero(); then(this.tracer.currentSpan()).isNull(); }); @@ -306,7 +306,7 @@ public class SleuthSpanCreatorAspectMonoTests { .containsEntry("customTestTag11", "test"); then(spans.get(0).annotations().stream().map(Annotation::value) .collect(Collectors.toList())).contains("customTest.before", - "customTest.after"); + "customTest.after"); then(spans.get(0).duration()).isNotZero(); then(this.tracer.currentSpan()).isNull(); }); @@ -361,7 +361,7 @@ public class SleuthSpanCreatorAspectMonoTests { then(spans.get(0).tags()).containsEntry("error", "test exception 13"); then(spans.get(0).annotations().stream().map(Annotation::value) .collect(Collectors.toList())).contains("testMethod13.before", - "testMethod13.afterFailure", "testMethod13.after"); + "testMethod13.afterFailure", "testMethod13.after"); then(spans.get(0).duration()).isNotZero(); then(this.tracer.currentSpan()).isNull(); }); @@ -624,8 +624,7 @@ public class SleuthSpanCreatorAspectMonoTests { @Override public Mono newSpanInSubscriberContext() { - return Mono.subscriberContext().flatMap( - context -> Mono.just(id(tracer))); + return Mono.subscriberContext().flatMap(context -> Mono.just(id(tracer))); } } @@ -643,24 +642,17 @@ public class SleuthSpanCreatorAspectMonoTests { @NewSpan(name = "outerSpanInTraceContext") public Mono, Long>> outerNewSpanInTraceContext() { - return Mono - .defer(() -> Mono.just(id(tracer)) - .zipWith( - testBeanInterface.newSpanInTraceContext()) - .map(pair -> Pair.of( - Pair.of(pair.getT1(), id(tracer)), - pair.getT2()))); + return Mono.defer(() -> Mono.just(id(tracer)) + .zipWith(testBeanInterface.newSpanInTraceContext()).map(pair -> Pair + .of(Pair.of(pair.getT1(), id(tracer)), pair.getT2()))); } @NewSpan(name = "outerSpanInSubscriberContext") public Mono, Long>> outerNewSpanInSubscriberContext() { return Mono.subscriberContext() .flatMap(context -> Mono.just(id(tracer)) - .zipWith( - testBeanInterface.newSpanInSubscriberContext()) - .map(pair -> Pair.of( - Pair.of(pair.getT1(), - id(tracer)), + .zipWith(testBeanInterface.newSpanInSubscriberContext()) + .map(pair -> Pair.of(Pair.of(pair.getT1(), id(tracer)), pair.getT2()))); }