private void beginTrace(final ClientTracer tracer, final String sql, final Connection connection)
     throws SQLException {
   final String schema = connection.getSchema();
   tracer.startNewSpan("query");
   tracer.setCurrentClientServiceName(schema == null ? "MySQL" : schema);
   tracer.submitBinaryAnnotation("executed.query", sql);
   tracer.setClientSent();
 }
 private void endTrace(
     final ClientTracer tracer, final int warningCount, final SQLException statementException) {
   try {
     if (warningCount > 0) {
       tracer.submitBinaryAnnotation("warning.count", warningCount);
     }
     if (statementException != null) {
       tracer.submitBinaryAnnotation("error.code", statementException.getErrorCode());
     }
   } finally {
     tracer.setClientReceived();
   }
 }
  @Test
  public void testExecuteHttpCodeNoSuccess() throws Exception {

    final SpanId mockSpanId = mock(SpanId.class);
    when(mockSpanId.getTraceId()).thenReturn(TRACE_ID);
    when(mockSpanId.getSpanId()).thenReturn(SPAN_ID);
    when(mockSpanId.getParentSpanId()).thenReturn(PARENT_SPAN_ID);

    when(mockClientTracer.startNewSpan(PATH)).thenReturn(mockSpanId);
    when(mockClientResponse.getStatus()).thenReturn(SERVER_ERROR_STATUS);

    assertSame(mockClientResponse, interceptor.execute(mockExecutionContext));

    final InOrder inOrder = inOrder(mockClientTracer, mockClientRequest, mockExecutionContext);

    inOrder.verify(mockClientTracer).startNewSpan(PATH);
    inOrder.verify(mockClientRequest).header(BraveHttpHeaders.Sampled.getName(), "true");
    inOrder.verify(mockClientRequest).header(BraveHttpHeaders.TraceId.getName(), TRACE_ID_HEX);
    inOrder.verify(mockClientRequest).header(BraveHttpHeaders.SpanId.getName(), SPAN_ID_HEX);
    inOrder
        .verify(mockClientRequest)
        .header(BraveHttpHeaders.ParentSpanId.getName(), PARENT_SPAN_ID_HEX);
    inOrder.verify(mockClientTracer).setCurrentClientServiceName(CONTEXT_PATH);
    inOrder
        .verify(mockClientTracer)
        .submitBinaryAnnotation(REQUEST_ANNOTATION, HTTP_METHOD + " " + URI);
    inOrder.verify(mockClientTracer).setClientSent();
    inOrder.verify(mockExecutionContext).proceed();
    inOrder
        .verify(mockClientTracer)
        .submitBinaryAnnotation("http.responsecode", SERVER_ERROR_STATUS);
    inOrder.verify(mockClientTracer).submitAnnotation("failure");
    inOrder.verify(mockClientTracer).setClientReceived();
    verifyNoMoreInteractions(mockClientTracer);
  }
  @Test
  public void testExecuteTracingNoParentSpan() throws Exception {

    final SpanId mockSpanId = mock(SpanId.class);
    when(mockSpanId.getTraceId()).thenReturn(TRACE_ID);
    when(mockSpanId.getSpanId()).thenReturn(SPAN_ID);
    when(mockSpanId.getParentSpanId()).thenReturn(null);

    when(mockClientTracer.startNewSpan(PATH)).thenReturn(mockSpanId);
    assertSame(mockClientResponse, interceptor.execute(mockExecutionContext));

    final InOrder inOrder = inOrder(mockClientTracer, mockClientRequest, mockExecutionContext);

    inOrder.verify(mockClientTracer).startNewSpan(PATH);
    inOrder.verify(mockClientRequest).header(BraveHttpHeaders.Sampled.getName(), "true");
    inOrder.verify(mockClientRequest).header(BraveHttpHeaders.TraceId.getName(), TRACE_ID_HEX);
    inOrder.verify(mockClientRequest).header(BraveHttpHeaders.SpanId.getName(), SPAN_ID_HEX);
    inOrder.verify(mockClientTracer).setCurrentClientServiceName(CONTEXT_PATH);
    inOrder
        .verify(mockClientTracer)
        .submitBinaryAnnotation(REQUEST_ANNOTATION, HTTP_METHOD + " " + URI);
    inOrder.verify(mockClientTracer).setClientSent();
    inOrder.verify(mockExecutionContext).proceed();
    inOrder
        .verify(mockClientTracer)
        .submitBinaryAnnotation(HTTP_RESPONSE_CODE_ANNOTATION, OK_STATUS);
    inOrder.verify(mockClientTracer).setClientReceived();
    verifyNoMoreInteractions(mockClientTracer);
  }
  @Test
  public void testTracingFalse()
      throws ClientProtocolException, IOException, UnsatisfiedExpectationException {
    when(clientTracer.startNewSpan(PATH)).thenReturn(null);

    final HttpRequestImpl request = new HttpRequestImpl();
    request
        .method(Method.GET)
        .path(PATH)
        .httpMessageHeader(BraveHttpHeaders.Sampled.getName(), "false");
    final HttpResponseImpl response = new HttpResponseImpl(200, null, null);
    responseProvider.set(request, response);

    final CloseableHttpClient httpclient =
        HttpClients.custom()
            .addInterceptorFirst(new BraveHttpRequestInterceptor(clientTracer))
            .addInterceptorFirst(new BraveHttpResponseInterceptor(clientTracer))
            .build();
    try {
      final HttpGet httpGet = new HttpGet(REQUEST);
      final CloseableHttpResponse httpClientResponse = httpclient.execute(httpGet);
      assertEquals(200, httpClientResponse.getStatusLine().getStatusCode());
      httpClientResponse.close();
      mockServer.verify();

      final InOrder inOrder = inOrder(clientTracer);
      inOrder.verify(clientTracer).startNewSpan(PATH);
      inOrder.verify(clientTracer).submitBinaryAnnotation("request", "GET " + PATH);
      inOrder.verify(clientTracer).setClientSent();
      inOrder.verify(clientTracer).setClientReceived();
      verifyNoMoreInteractions(clientTracer);
    } finally {
      httpclient.close();
    }
  }
  @Test
  public void testExecuteNoTracingSpanNameNotSpecified() throws Exception {

    when(mockClientTracer.startNewSpan(CUSTOM_SPAN_NAME)).thenReturn(null);
    assertSame(mockClientResponse, interceptor.execute(mockExecutionContext));

    final InOrder inOrder = inOrder(mockClientTracer, mockClientRequest, mockExecutionContext);

    inOrder.verify(mockClientTracer).startNewSpan(PATH);
    inOrder.verify(mockClientRequest).header(BraveHttpHeaders.Sampled.getName(), "false");
    inOrder.verify(mockClientTracer).setCurrentClientServiceName(CONTEXT_PATH);
    inOrder
        .verify(mockClientTracer)
        .submitBinaryAnnotation(REQUEST_ANNOTATION, HTTP_METHOD + " " + URI);
    inOrder.verify(mockClientTracer).setClientSent();
    inOrder.verify(mockExecutionContext).proceed();
    inOrder
        .verify(mockClientTracer)
        .submitBinaryAnnotation(HTTP_RESPONSE_CODE_ANNOTATION, OK_STATUS);
    inOrder.verify(mockClientTracer).setClientReceived();
    verifyNoMoreInteractions(mockClientTracer);
  }
  @Test
  public void testExecuteProceedThrowsException() throws Exception {

    final SpanId mockSpanId = mock(SpanId.class);
    when(mockSpanId.getTraceId()).thenReturn(TRACE_ID);
    when(mockSpanId.getSpanId()).thenReturn(SPAN_ID);
    when(mockSpanId.getParentSpanId()).thenReturn(PARENT_SPAN_ID);

    when(mockClientTracer.startNewSpan(PATH)).thenReturn(mockSpanId);
    final IllegalStateException exception = new IllegalStateException("Test exception");
    when(mockExecutionContext.proceed()).thenThrow(exception);

    try {
      interceptor.execute(mockExecutionContext);
      fail("Expected exception.");
    } catch (final IllegalStateException e) {
      assertSame(exception, e);
    }

    final InOrder inOrder = inOrder(mockClientTracer, mockClientRequest, mockExecutionContext);

    inOrder.verify(mockClientTracer).startNewSpan(PATH);
    inOrder.verify(mockClientRequest).header(BraveHttpHeaders.Sampled.getName(), "true");
    inOrder.verify(mockClientRequest).header(BraveHttpHeaders.TraceId.getName(), TRACE_ID_HEX);
    inOrder.verify(mockClientRequest).header(BraveHttpHeaders.SpanId.getName(), SPAN_ID_HEX);
    inOrder
        .verify(mockClientRequest)
        .header(BraveHttpHeaders.ParentSpanId.getName(), Long.toHexString(PARENT_SPAN_ID));
    inOrder.verify(mockClientTracer).setCurrentClientServiceName(CONTEXT_PATH);
    inOrder
        .verify(mockClientTracer)
        .submitBinaryAnnotation(REQUEST_ANNOTATION, HTTP_METHOD + " " + URI);
    inOrder.verify(mockClientTracer).setClientSent();
    inOrder.verify(mockExecutionContext).proceed();
    inOrder.verify(mockClientTracer).submitBinaryAnnotation(HTTP_RESPONSE_CODE_ANNOTATION, 0);
    inOrder.verify(mockClientTracer).submitAnnotation("failure");
    inOrder.verify(mockClientTracer).setClientReceived();
    verifyNoMoreInteractions(mockClientTracer);
  }