Skip to content

Commit

Permalink
fix test
Browse files Browse the repository at this point in the history
  • Loading branch information
@jrhee17
jrhee17 committed Feb 20, 2025
1 parent 037585c commit dd69a72
Showing 1 changed file with 54 additions and 28 deletions.
82 changes: 54 additions & 28 deletions brave/brave6/src/test/java/com/linecorp/armeria/it/brave/BraveIntegrationTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,16 +30,19 @@
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.BlockingDeque;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.stream.IntStream;

import org.apache.thrift.async.AsyncMethodCallback;
import org.apache.thrift.transport.TTransportException;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
Expand Down Expand Up @@ -98,6 +101,8 @@ class BraveIntegrationTest {

private static final String CLIENT_TYPE_HEADER = "x-client-type";
private static final SpanHandlerImpl spanHandler = new SpanHandlerImpl();
private static final BlockingDeque<Tracing> perTestTracings = new LinkedBlockingDeque<>();
private static final BlockingDeque<Tracing> perClassTracings = new LinkedBlockingDeque<>();

@RegisterExtension
static ServerExtension server = new ServerExtension(true) {
Expand Down Expand Up @@ -194,8 +199,15 @@ protected HttpResponse doGet(ServiceRequestContext ctx, HttpRequest req)
};

@AfterEach
void shouldHaveNoExtraSpans() {
void afterEach() {
assertThat(spanHandler.spans).isEmpty();
perTestTracings.forEach(Tracing::close);
perTestTracings.clear();
}

@AfterAll
static void afterAll() {
perClassTracings.forEach(Tracing::close);
}

private static HttpService tHttpDecorate(String name, AsyncIface asyncIface) {
Expand All @@ -207,38 +219,48 @@ private static HttpService tHttpDecorate(String name, AsyncIface asyncIface) {
return (ctx, req) -> {
final String braveServiceType = ctx.request().headers().get(CLIENT_TYPE_HEADER);
if ("http".equals(braveServiceType)) {
return BraveService.newDecorator(newTracing(name)).apply(service).serve(ctx, req);
return BraveService.newDecorator(newTracing(name, false)).apply(service).serve(ctx, req);
}
return service.serve(ctx, req);
};
}

private static HttpService httpDecorate(String name, HttpService service) {
return BraveService.newDecorator(newTracing(name)).apply(service);
return BraveService.newDecorator(newTracing(name, false)).apply(service);
}

private static TestService.AsyncIface newClient(String path) {
final ServiceRequestContext ctx = ServiceRequestContext.current();
final String braveServiceType = ctx.request().headers().get(CLIENT_TYPE_HEADER);
return ThriftClients.builder(server.httpUri())
.path(path)
.addHeader(CLIENT_TYPE_HEADER, "http")
.addHeader(CLIENT_TYPE_HEADER, braveServiceType)
.decorator(BraveClient.newDecorator(newTracing("client" + path)))
.build(TestService.AsyncIface.class);
}

private static Tracing newTracing(String name) {
return newTracing(name, true);
}

private static Tracing newTracing(String name, boolean perTest) {
final CurrentTraceContext currentTraceContext =
RequestContextCurrentTraceContext.builder()
.nonRequestThread("nonrequest-")
.addScopeDecorator(StrictScopeDecorator.create())
.build();
return Tracing.newBuilder()
.currentTraceContext(currentTraceContext)
.localServiceName(name)
.addSpanHandler(spanHandler)
.sampler(Sampler.ALWAYS_SAMPLE)
.build();
final Tracing tracing = Tracing.newBuilder()
.currentTraceContext(currentTraceContext)
.localServiceName(name)
.addSpanHandler(spanHandler)
.sampler(Sampler.ALWAYS_SAMPLE)
.build();
if (perTest) {
perTestTracings.add(tracing);
} else {
perClassTracings.add(tracing);
}
return tracing;
}

@Test
Expand Down Expand Up @@ -445,24 +467,28 @@ void testSpanInThreadPoolHasSameTraceId() throws Exception {
@ParameterizedTest
@ValueSource(strings = {"http", "rpc"})
void testServerTimesOut(String type) throws Exception {
final TestService.Iface timeoutClient =
ThriftClients.builder(server.httpUri())
.path("/timeout")
.addHeader(CLIENT_TYPE_HEADER, type)
.decorator(BraveClient.newDecorator(newTracing("client/timeout")))
.build(TestService.Iface.class);
assertThatThrownBy(() -> timeoutClient.hello("name"))
.isInstanceOf(TTransportException.class)
.hasCauseInstanceOf(InvalidResponseHeadersException.class);
final MutableSpan[] spans = spanHandler.take(2);

final MutableSpan serverSpan = findSpan(spans, "service/timeout");
final MutableSpan clientSpan = findSpan(spans, "client/timeout");

// Server timed out meaning it did still send a timeout response to the client and we have all
// annotations.
assertThat(serverSpan.annotations()).hasSize(2);
assertThat(clientSpan.annotations()).hasSize(2);
try (ClientFactory cf = ClientFactory.builder().build()) {
final TestService.Iface timeoutClient =
ThriftClients.builder(server.httpUri())
.path("/timeout")
.factory(cf)
.addHeader(CLIENT_TYPE_HEADER, type)
.decorator(BraveClient.newDecorator(newTracing("client/timeout")))
.build(TestService.Iface.class);
assertThatThrownBy(() -> timeoutClient.hello("name"))
.isInstanceOf(TTransportException.class)
.hasCauseInstanceOf(InvalidResponseHeadersException.class);
final MutableSpan[] spans = spanHandler.take(2);

final MutableSpan serverSpan = findSpan(spans, "service/timeout");
final MutableSpan clientSpan = findSpan(spans, "client/timeout");

// Server timed out meaning it did still send a timeout response to the client and we have all
// annotations. A separate client factory is used to guarantee that client span annotations
// always contain connection related extra annotations.
assertThat(serverSpan.annotations()).hasSize(2);
assertThat(clientSpan.annotations()).hasSize(6);
}
}

@ParameterizedTest
Expand Down

0 comments on commit dd69a72

Please sign in to comment.