/*
 * Copyright 2013-2021 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.springframework.cloud.sleuth.test;

import java.util.List;
import java.util.Queue;

import org.springframework.cloud.sleuth.Span;
import org.springframework.cloud.sleuth.exporter.FinishedSpan;

public interface TestSpanHandler extends Iterable<FinishedSpan> {

	List<FinishedSpan> reportedSpans();

	FinishedSpan takeLocalSpan();

	void clear();

	FinishedSpan takeRemoteSpan(Span.Kind kind);

	FinishedSpan takeRemoteSpanWithError(Span.Kind kind);

	FinishedSpan get(int index);

	void assertAllSpansWereFinishedOrAbandoned(Queue<Span> createdSpans);

}

@@ -1,8 +1,9 @@
+%0A
 /*%0A * Co
@@ -711,16 +711,52 @@
 l.Queue;
+%0Aimport java.util.stream.Collectors;
 %0A%0Aimport
@@ -859,16 +859,72 @@
 dSpan;%0A%0A
+import static org.assertj.core.api.BDDAssertions.then;%0A%0A
 public i
@@ -1199,24 +1199,32 @@
 t index);%0A%0A%09
+default 
 void assertA
@@ -1279,13 +1279,1484 @@
 edSpans)
-;
+ %7B%0A%09%09List%3CFinishedSpan%3E finishedSpans = reportedSpans();%0A%09%09then(finishedSpans).as(%22There should be that many finished spans as many created ones%22)%0A%09%09%09%09.hasSize(createdSpans.size());%0A%09%09// finished -%3E a,b,c ; created -%3E b,c,d =%3E matchedFinished = b,c%0A%09%09List%3CFinishedSpan%3E matchedFinishedSpans = finishedSpans.stream()%0A%09%09%09%09.filter(f -%3E createdSpans.stream().anyMatch(cs -%3E f.getSpanId().equals(cs.context().spanId())))%0A%09%09%09%09.collect(Collectors.toList());%0A%09%09// finished -%3E a,b,c ; created -%3E b,c,d =%3E matchedCreated = b,c%0A%09%09List%3CSpan%3E matchedCreatedSpans = createdSpans.stream()%0A%09%09%09%09.filter(cs -%3E finishedSpans.stream().anyMatch(f -%3E cs.context().spanId().equals(f.getSpanId())))%0A%09%09%09%09.collect(Collectors.toList());%0A%09%09// finished -%3E a,b,c ; created -%3E b,c,d =%3E missingFinished = a%0A%09%09List%3CFinishedSpan%3E missingFinishedSpans = finishedSpans.stream()%0A%09%09%09%09.filter(f -%3E matchedFinishedSpans.stream().noneMatch(m -%3E m.getSpanId().equals(f.getSpanId())))%0A%09%09%09%09.collect(Collectors.toList());%0A%09%09// finished -%3E a,b,c ; created -%3E b,c,d =%3E missingCreated = d%0A%09%09List%3CSpan%3E missingCreatedSpans = createdSpans.stream().filter(%0A%09%09%09%09f -%3E matchedCreatedSpans.stream().noneMatch(m -%3E m.context().spanId().equals(f.context().spanId())))%0A%09%09%09%09.collect(Collectors.toList());%0A%09%09if (!missingFinishedSpans.isEmpty() %7C%7C !missingCreatedSpans.isEmpty()) %7B%0A%09%09%09throw new AssertionError(%22There were unmatched created spans %22 + missingCreatedSpans%0A%09%09%09%09%09+ %22 and/or finished span %22 + missingFinishedSpans);%0A%09%09%7D%0A%09%7D
 %0A%0A%7D%0A

Update TestSpanHandler.java