Skip to content

Commit 7f89976

Browse files
jrhee17jrhee17
committed
fix test
1 parent 037585c commit 7f89976

File tree

1 file changed

+56
-33
lines changed

1 file changed

+56
-33
lines changed

‎brave/brave6/src/test/java/com/linecorp/armeria/it/brave/BraveIntegrationTest.java

Lines changed: 56 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import static com.linecorp.armeria.common.SessionProtocol.H1C;
2323
import static org.assertj.core.api.Assertions.assertThat;
2424
import static org.assertj.core.api.Assertions.assertThatThrownBy;
25+
import static org.awaitility.Awaitility.await;
2526

2627
import java.time.Duration;
2728
import java.util.ArrayList;
@@ -40,6 +41,7 @@
4041

4142
import org.apache.thrift.async.AsyncMethodCallback;
4243
import org.apache.thrift.transport.TTransportException;
44+
import org.junit.jupiter.api.AfterAll;
4345
import org.junit.jupiter.api.AfterEach;
4446
import org.junit.jupiter.api.Test;
4547
import org.junit.jupiter.api.extension.RegisterExtension;
@@ -57,6 +59,7 @@
5759
import com.linecorp.armeria.client.WebClient;
5860
import com.linecorp.armeria.client.brave.BraveClient;
5961
import com.linecorp.armeria.client.thrift.ThriftClients;
62+
import com.linecorp.armeria.common.AggregatedHttpResponse;
6063
import com.linecorp.armeria.common.HttpRequest;
6164
import com.linecorp.armeria.common.HttpResponse;
6265
import com.linecorp.armeria.common.HttpStatus;
@@ -97,6 +100,7 @@
97100
class BraveIntegrationTest {
98101

99102
private static final String CLIENT_TYPE_HEADER = "x-client-type";
103+
private static final String TIMEOUT_HEADER = "x-timeout";
100104
private static final SpanHandlerImpl spanHandler = new SpanHandlerImpl();
101105

102106
@RegisterExtension
@@ -130,7 +134,8 @@ protected void configure(ServerBuilder sb) throws Exception {
130134
sb.service("/qux", tHttpDecorate("service/qux", (name, resultHandler) ->
131135
resultHandler.onComplete("Hello, " + name + '!')));
132136

133-
sb.service("/pool", httpDecorate("service/pool", new AbstractHttpService() {
137+
final Tracing servicePoolTracing = newTracing("service/pool");
138+
sb.service("/pool", httpDecorate(servicePoolTracing, new AbstractHttpService() {
134139
@Override
135140
protected HttpResponse doGet(ServiceRequestContext ctx, HttpRequest req)
136141
throws Exception {
@@ -145,7 +150,7 @@ protected HttpResponse doGet(ServiceRequestContext ctx, HttpRequest req)
145150
countDownLatch.countDown();
146151
countDownLatch.await();
147152
}
148-
final Span span = Tracing.currentTracer().nextSpan().start();
153+
final Span span = servicePoolTracing.tracer().nextSpan().start();
149154
try (SpanInScope unused =
150155
Tracing.currentTracer().withSpanInScope(span)) {
151156
if (i == 1) {
@@ -166,8 +171,8 @@ protected HttpResponse doGet(ServiceRequestContext ctx, HttpRequest req)
166171
result -> allAsList(IntStream.range(1, 3).mapToObj(
167172
i -> executorService.submit(
168173
RequestContext.current().makeContextAware(() -> {
169-
final ScopedSpan span = Tracing.currentTracer()
170-
.startScopedSpan("aloha");
174+
final ScopedSpan span =
175+
servicePoolTracing.tracer().startScopedSpan("aloha");
171176
try {
172177
return null;
173178
} finally {
@@ -185,19 +190,30 @@ protected HttpResponse doGet(ServiceRequestContext ctx, HttpRequest req)
185190
}
186191
}));
187192

188-
sb.service("/timeout", tHttpDecorate("service/timeout",
189-
// This service never calls the handler and will timeout.
190-
(name, resultHandler) -> {}));
193+
sb.service("/timeout",
194+
tHttpDecorate("service/timeout",
195+
// This service never calls the handler and will timeout.
196+
(name, resultHandler) -> {
197+
final ServiceRequestContext ctx = ServiceRequestContext.current();
198+
if (ctx.request().headers().contains(TIMEOUT_HEADER)) {
199+
ctx.timeoutNow();
200+
}
201+
}));
191202

192203
sb.service("/http", (req, ctx) -> HttpResponse.of(HttpStatus.OK));
193204
}
194205
};
195206

196207
@AfterEach
197-
void shouldHaveNoExtraSpans() {
208+
void afterEach() {
198209
assertThat(spanHandler.spans).isEmpty();
199210
}
200211

212+
@AfterAll
213+
static void afterAll() throws Exception {
214+
Tracing.current().close();
215+
}
216+
201217
private static HttpService tHttpDecorate(String name, AsyncIface asyncIface) {
202218
final THttpService service =
203219
THttpService.builder()
@@ -213,16 +229,16 @@ private static HttpService tHttpDecorate(String name, AsyncIface asyncIface) {
213229
};
214230
}
215231

216-
private static HttpService httpDecorate(String name, HttpService service) {
217-
return BraveService.newDecorator(newTracing(name)).apply(service);
232+
private static HttpService httpDecorate(Tracing tracing, HttpService service) {
233+
return BraveService.newDecorator(tracing).apply(service);
218234
}
219235

220236
private static TestService.AsyncIface newClient(String path) {
221237
final ServiceRequestContext ctx = ServiceRequestContext.current();
222238
final String braveServiceType = ctx.request().headers().get(CLIENT_TYPE_HEADER);
223239
return ThriftClients.builder(server.httpUri())
224240
.path(path)
225-
.addHeader(CLIENT_TYPE_HEADER, "http")
241+
.addHeader(CLIENT_TYPE_HEADER, braveServiceType)
226242
.decorator(BraveClient.newDecorator(newTracing("client" + path)))
227243
.build(TestService.AsyncIface.class);
228244
}
@@ -434,7 +450,9 @@ void testServiceInitiatedTrace(String type) throws Exception {
434450

435451
@Test
436452
void testSpanInThreadPoolHasSameTraceId() throws Exception {
437-
server.webClient().get("pool").aggregate().get();
453+
final AggregatedHttpResponse res = server.blockingWebClient().get("pool");
454+
assertThat(res.contentUtf8()).isEqualTo("Lee");
455+
await().untilAsserted(() -> assertThat(spanHandler.spans).hasSize(5));
438456
final MutableSpan[] spans = spanHandler.take(5);
439457
assertThat(Arrays.stream(spans).map(MutableSpan::traceId).collect(toImmutableSet())).hasSize(1);
440458
assertThat(Arrays.stream(spans).map(MutableSpan::parentId)
@@ -445,24 +463,29 @@ void testSpanInThreadPoolHasSameTraceId() throws Exception {
445463
@ParameterizedTest
446464
@ValueSource(strings = {"http", "rpc"})
447465
void testServerTimesOut(String type) throws Exception {
448-
final TestService.Iface timeoutClient =
449-
ThriftClients.builder(server.httpUri())
450-
.path("/timeout")
451-
.addHeader(CLIENT_TYPE_HEADER, type)
452-
.decorator(BraveClient.newDecorator(newTracing("client/timeout")))
453-
.build(TestService.Iface.class);
454-
assertThatThrownBy(() -> timeoutClient.hello("name"))
455-
.isInstanceOf(TTransportException.class)
456-
.hasCauseInstanceOf(InvalidResponseHeadersException.class);
457-
final MutableSpan[] spans = spanHandler.take(2);
458-
459-
final MutableSpan serverSpan = findSpan(spans, "service/timeout");
460-
final MutableSpan clientSpan = findSpan(spans, "client/timeout");
461-
462-
// Server timed out meaning it did still send a timeout response to the client and we have all
463-
// annotations.
464-
assertThat(serverSpan.annotations()).hasSize(2);
465-
assertThat(clientSpan.annotations()).hasSize(2);
466+
try (ClientFactory cf = ClientFactory.builder().build()) {
467+
final TestService.Iface timeoutClient =
468+
ThriftClients.builder(server.httpUri())
469+
.path("/timeout")
470+
.factory(cf)
471+
.addHeader(CLIENT_TYPE_HEADER, type)
472+
.addHeader(TIMEOUT_HEADER, true)
473+
.decorator(BraveClient.newDecorator(newTracing("client/timeout")))
474+
.build(TestService.Iface.class);
475+
assertThatThrownBy(() -> timeoutClient.hello("name"))
476+
.isInstanceOf(TTransportException.class)
477+
.hasCauseInstanceOf(InvalidResponseHeadersException.class);
478+
final MutableSpan[] spans = spanHandler.take(2);
479+
480+
final MutableSpan serverSpan = findSpan(spans, "service/timeout");
481+
final MutableSpan clientSpan = findSpan(spans, "client/timeout");
482+
483+
// Server timed out meaning it did still send a timeout response to the client and we have all
484+
// annotations. A separate client factory is used to guarantee that client span annotations
485+
// always contain connection related extra annotations.
486+
assertThat(serverSpan.annotations()).hasSize(2);
487+
assertThat(clientSpan.annotations()).hasSize(6);
488+
}
466489
}
467490

468491
@ParameterizedTest
@@ -473,7 +496,7 @@ void testHttp2ClientTimesOut(String type) throws Exception {
473496
.path("/timeout")
474497
.addHeader(CLIENT_TYPE_HEADER, type)
475498
.decorator(BraveClient.newDecorator(newTracing("client/timeout")))
476-
.responseTimeout(Duration.ofSeconds(3))
499+
.responseTimeout(Duration.ofSeconds(1))
477500
.build(TestService.Iface.class);
478501
testClientTimesOut(timeoutClientClientTimesOut);
479502
}
@@ -582,8 +605,8 @@ public void onError(Exception exception) {
582605
}
583606
}
584607

585-
private static class SpanHandlerImpl extends SpanHandler {
586-
private final BlockingQueue<MutableSpan> spans = new LinkedBlockingQueue<>();
608+
static final class SpanHandlerImpl extends SpanHandler {
609+
final BlockingQueue<MutableSpan> spans = new LinkedBlockingQueue<>();
587610

588611
@Override
589612
public boolean end(TraceContext context, MutableSpan span, Cause cause) {

0 commit comments

Comments
 (0)