com.amazonaws.xray.AWSXRay Java Examples

@Test
public void testRaceConditionOnRecorderInitialization() {
    AWSXRay.setGlobalRecorder(null);
    // TracingHandler will not have the initialized recorder
    AWSLambda lambda = AWSLambdaClientBuilder
        .standard()
        .withRequestHandlers(new TracingHandler())
        .withRegion(Regions.US_EAST_1)
        .withCredentials(new AWSStaticCredentialsProvider(new BasicAWSCredentials("fake", "fake")))
        .build();

    mockHttpClient(lambda, "null");

    // Now init the global recorder
    AWSXRayRecorder recorder = AWSXRayRecorderBuilder.defaultRecorder();
    recorder.setContextMissingStrategy(new LogErrorContextMissingStrategy());
    AWSXRay.setGlobalRecorder(recorder);

    // Test logic
    InvokeRequest request = new InvokeRequest();
    request.setFunctionName("testFunctionName");
    lambda.invoke(request);
}
 

@Test
public void testCaptureRuntimeExceptionWithoutSegment() throws Exception {
    ContextMissingStrategy oldStrategy = AWSXRay.getGlobalRecorder().getContextMissingStrategy();
    AWSXRay.getGlobalRecorder().setContextMissingStrategy(new IgnoreErrorContextMissingStrategy());
    try {
        RuntimeException exception = new RuntimeException("foo");
        when(delegate.execute(SQL)).thenThrow(exception);
        try {
            statement.execute(SQL);
            fail("Expected exception is not thrown");
        } catch (RuntimeException th) {
            assertEquals(exception, th);
        }
    } finally {
        AWSXRay.getGlobalRecorder().setContextMissingStrategy(oldStrategy);
    }
}
 

@Test
public void testAsyncServletRequestWithCompletedAsync() throws IOException, ServletException {
    AWSXRayServletFilter servletFilter = new AWSXRayServletFilter("test");

    HttpServletRequest request = Mockito.mock(HttpServletRequest.class);
    Mockito.when(request.getRequestURL()).thenReturn(new StringBuffer("test_url"));
    Mockito.when(request.getMethod()).thenReturn("TEST_METHOD");
    Mockito.when(request.isAsyncStarted()).thenReturn(true);
    Mockito.when(request.getAsyncContext()).thenThrow(IllegalStateException.class);

    HttpServletResponse response = Mockito.mock(HttpServletResponse.class);

    FilterChain chain = mockChain(request, response);

    AsyncEvent event = Mockito.mock(AsyncEvent.class);
    Mockito.when(event.getSuppliedRequest()).thenReturn(request);
    Mockito.when(event.getSuppliedResponse()).thenReturn(response);

    servletFilter.doFilter(request, response, chain);
    Assert.assertNull(AWSXRay.getTraceEntity());

    Mockito.verify(AWSXRay.getGlobalRecorder().getEmitter(), Mockito.times(1)).sendSegment(Mockito.any());
}
 

@Before
public void setup() throws Exception {
    statement = TracingStatement.decorateStatement(delegate);
    preparedStatement = TracingStatement.decoratePreparedStatement(preparedDelegate, SQL);
    callableStatement = TracingStatement.decorateCallableStatement(callableDelegate, SQL);
    when(delegate.getConnection()).thenReturn(connection);
    when(preparedDelegate.getConnection()).thenReturn(connection);
    when(callableDelegate.getConnection()).thenReturn(connection);
    when(connection.getMetaData()).thenReturn(metaData);
    when(connection.getCatalog()).thenReturn(CATALOG);
    when(metaData.getURL()).thenReturn(URL);
    when(metaData.getUserName()).thenReturn(USER);
    when(metaData.getDriverVersion()).thenReturn(DRIVER_VERSION);
    when(metaData.getDatabaseProductName()).thenReturn(DB_TYPE);
    when(metaData.getDatabaseProductVersion()).thenReturn(DB_VERSION);
    expectedSqlParams = new HashMap<>();
    expectedSqlParams.put("url", URL);
    expectedSqlParams.put("user", USER);
    expectedSqlParams.put("driver_version", DRIVER_VERSION);
    expectedSqlParams.put("database_type", DB_TYPE);
    expectedSqlParams.put("database_version", DB_VERSION);
    AWSXRay.beginSegment("foo");
}
 

@Test
public void testGlobalMapSegmentContext() {
    Segment test = AWSXRay.beginSegment("test");


    List<Integer> list = new ArrayList<>();
    for (int i = 0; i < 100; i++) {
        list.add(i);
    }

    list.parallelStream().forEach(e -> {
        AWSXRay.setTraceEntity(test);
        AWSXRay.createSubsegment("parallelPrint", (subsegment) -> {
        });
    });

    Assert.assertEquals(100, test.getTotalSize().intValue());

    AWSXRay.endSegment();
}
 

@Test
public void testMultipleSegmentListeners() {
    SecondSegmentListener secondSegmentListener = new SecondSegmentListener();
    AWSXRay.getGlobalRecorder().addSegmentListener(secondSegmentListener);
    Segment test = AWSXRay.beginSegment("test");
    String beginAnnotation = test.getAnnotations().get("beginTest").toString();


    Assert.assertEquals(1, secondSegmentListener.getTestVal());

    Assert.assertEquals("isPresent", beginAnnotation);

    AWSXRay.endSegment();
    String endAnnotation = test.getAnnotations().get("endTest").toString();

    Assert.assertEquals("isPresent", endAnnotation);
    Assert.assertEquals(1, secondSegmentListener.getTestVal2());
}
 

@Test
public void testStreamSomeChildrenRemovedFromParent() {
    TraceID traceId = new TraceID();
    DefaultStreamingStrategy defaultStreamingStrategy = new DefaultStreamingStrategy(1);

    Segment bigSegment = new SegmentImpl(AWSXRay.getGlobalRecorder(), "big", traceId);
    bigSegment.setStartTime(1.0);

    for (int i = 0; i < 5; i++) {
        Subsegment subsegment = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "child" + i, bigSegment);
        subsegment.setStartTime(1.0);
        bigSegment.addSubsegment(subsegment);
        subsegment.end();
    }
    Assert.assertTrue(defaultStreamingStrategy.requiresStreaming(bigSegment));
    defaultStreamingStrategy.streamSome(bigSegment, AWSXRay.getGlobalRecorder().getEmitter());
    Assert.assertTrue(bigSegment.getTotalSize().intValue() == 0);
}
 

@Test
public void testStreamSomeChildrenNotRemovedFromParent() {
    TraceID traceId = new TraceID();
    DefaultStreamingStrategy defaultStreamingStrategy = new DefaultStreamingStrategy(1);

    Segment bigSegment = new SegmentImpl(AWSXRay.getGlobalRecorder(), "big", traceId);
    bigSegment.setStartTime(1.0);

    for (int i = 0; i < 5; i++) {
        Subsegment subsegment = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "child" + i, bigSegment);
        subsegment.setStartTime(1.0);
        bigSegment.addSubsegment(subsegment);
    }
    Assert.assertTrue(defaultStreamingStrategy.requiresStreaming(bigSegment));
    defaultStreamingStrategy.streamSome(bigSegment, AWSXRay.getGlobalRecorder().getEmitter());
    Assert.assertTrue(bigSegment.getTotalSize().intValue() == 5);
}
 

@Test
public void testLeakedSubsegmentsAreCleanedBetweenInvocations() {

    LambdaSegmentContext lsc = new LambdaSegmentContext();

    PowerMockito.stub(PowerMockito.method(LambdaSegmentContext.class, "getTraceHeaderFromEnvironment"))
                .toReturn(TraceHeader.fromString(TRACE_HEADER));
    Subsegment firstInvocation = lsc.beginSubsegment(AWSXRay.getGlobalRecorder(), "test");
    Assert.assertNotNull(AWSXRay.getTraceEntity());

    PowerMockito.stub(PowerMockito.method(LambdaSegmentContext.class, "getTraceHeaderFromEnvironment"))
                .toReturn(TraceHeader.fromString(TRACE_HEADER_2));
    Subsegment secondInvocation = lsc.beginSubsegment(AWSXRay.getGlobalRecorder(), "test");
    Assert.assertNotNull(AWSXRay.getTraceEntity());

    Assert.assertTrue(FacadeSegment.class.isInstance(firstInvocation.getParent()));
    Assert.assertTrue(FacadeSegment.class.isInstance(secondInvocation.getParent()));
}
 

@Before
public void setupAWSXRay() {
    Emitter blankEmitter = Mockito.mock(Emitter.class);
    Mockito.doReturn(true).when(blankEmitter).sendSegment(Mockito.any());
    Mockito.doReturn(true).when(blankEmitter).sendSubsegment(Mockito.any());
    CustomSegmentListener segmentListener = new CustomSegmentListener();

    AWSXRay.setGlobalRecorder(AWSXRayRecorderBuilder.standard()
                                                    .withEmitter(blankEmitter)
                                                    .withSegmentListener(segmentListener)
                                                    .build());
    AWSXRay.clearTraceEntity();
}
 

@Test
public void testAWSXRayServletAsyncListenerEmitsSegmentWhenProcessingEvent() throws IOException, ServletException {
    AWSXRayServletFilter servletFilter = new AWSXRayServletFilter("test");

    AsyncContext asyncContext = Mockito.mock(AsyncContext.class);

    HttpServletRequest request = Mockito.mock(HttpServletRequest.class);
    Mockito.when(request.getRequestURL()).thenReturn(new StringBuffer("test_url"));
    Mockito.when(request.getMethod()).thenReturn("TEST_METHOD");
    Mockito.when(request.isAsyncStarted()).thenReturn(true);
    Mockito.when(request.getAsyncContext()).thenReturn(asyncContext);

    HttpServletResponse response = Mockito.mock(HttpServletResponse.class);

    FilterChain chain = mockChain(request, response);

    AsyncEvent event = Mockito.mock(AsyncEvent.class);
    Mockito.when(event.getSuppliedRequest()).thenReturn(request);
    Mockito.when(event.getSuppliedResponse()).thenReturn(response);

    servletFilter.doFilter(request, response, chain);
    Assert.assertNull(AWSXRay.getTraceEntity());

    AWSXRayServletAsyncListener listener = (AWSXRayServletAsyncListener) Whitebox.getInternalState(servletFilter, "listener");
    listener.onComplete(event);

    Mockito.verify(AWSXRay.getGlobalRecorder().getEmitter(), Mockito.times(1)).sendSegment(Mockito.any());
}
 

@Test
public void testDefaultStreamingStrategyForLambdaTraceContext() {
    DefaultStreamingStrategy defaultStreamingStrategy = new DefaultStreamingStrategy(1);

    //if FacadeSegment size is larger than maxSegmentSize and only the first subsegment is completed, first subsegment will be
    //streamed out
    FacadeSegment facadeSegmentOne = new FacadeSegment(AWSXRay.getGlobalRecorder(), new TraceID(), "",
                                                       TraceHeader.SampleDecision.SAMPLED);
    Subsegment firstSubsegment = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "FirstSubsegment", facadeSegmentOne);
    Subsegment secondSubsegment = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "SecondSubsegment", facadeSegmentOne);
    facadeSegmentOne.addSubsegment(firstSubsegment);
    facadeSegmentOne.addSubsegment(secondSubsegment);

    firstSubsegment.end();

    Assert.assertTrue(facadeSegmentOne.getTotalSize().intValue() == 2);
    defaultStreamingStrategy.streamSome(facadeSegmentOne, AWSXRay.getGlobalRecorder().getEmitter());
    Assert.assertTrue(facadeSegmentOne.getTotalSize().intValue() == 1);

    Subsegment tempOne = facadeSegmentOne.getSubsegments().get(0);
    Assert.assertEquals("SecondSubsegment", tempOne.getName());

    //if FarcadeSegment size is larger than maxSegmentSize and only the second subsegment is completed, second subsegment will
    //be streamed out
    FacadeSegment facadeSegmentTwo = new FacadeSegment(AWSXRay.getGlobalRecorder(), new TraceID(), "",
                                                       TraceHeader.SampleDecision.SAMPLED);
    Subsegment thirdSubsegment = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "ThirdSubsegment", facadeSegmentTwo);
    Subsegment fourthSubsegment = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "FourthSubsegment", facadeSegmentTwo);
    facadeSegmentTwo.addSubsegment(thirdSubsegment);
    facadeSegmentTwo.addSubsegment(fourthSubsegment);

    fourthSubsegment.end();

    Assert.assertTrue(facadeSegmentTwo.getTotalSize().intValue() == 2);
    defaultStreamingStrategy.streamSome(facadeSegmentTwo, AWSXRay.getGlobalRecorder().getEmitter());
    Assert.assertTrue(facadeSegmentTwo.getTotalSize().intValue() == 1);

    Subsegment tempTwo = facadeSegmentTwo.getSubsegments().get(0);
    Assert.assertEquals("ThirdSubsegment", tempTwo.getName());
}
 

@Test
public void testOnBeginSegment() {
    Segment test = AWSXRay.beginSegment("test");
    String beginAnnotation = test.getAnnotations().get("beginTest").toString();

    Assert.assertEquals("isPresent", beginAnnotation);

    AWSXRay.endSegment();
}
 

@Test
public void testOnEndSegment() {
    Segment test = AWSXRay.beginSegment("test");
    AWSXRay.endSegment();
    String endAnnotation = test.getAnnotations().get("endTest").toString();

    Assert.assertEquals("isPresent", endAnnotation);
}
 

@Test
public void testBushyandSpindlySegmentTreeStreaming() {
    TraceID traceId = new TraceID();

    Segment bigSegment = new SegmentImpl(AWSXRay.getGlobalRecorder(), "big", traceId);
    bigSegment.setStartTime(1.0);

    for (int i = 0; i < 5; i++) {
        Subsegment subsegment = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "child" + i, bigSegment);
        subsegment.setStartTime(1.0);
        bigSegment.addSubsegment(subsegment);
        subsegment.end();
    }

    SubsegmentImpl holder = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "big_child0", bigSegment);
    holder.setStartTime(1.0);
    bigSegment.addSubsegment(holder);
    holder.end();

    SubsegmentImpl holder1 = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "big_child1", bigSegment);
    holder1.setStartTime(1.0);
    bigSegment.addSubsegment(holder1);
    holder1.end();

    SubsegmentImpl holder2 = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "big_child2", bigSegment);
    holder2.setStartTime(1.0);
    bigSegment.addSubsegment(holder2);
    holder2.end();

    DefaultStreamingStrategy defaultStreamingStrategy = new DefaultStreamingStrategy(1);
    Assert.assertTrue(defaultStreamingStrategy.requiresStreaming(bigSegment));
    defaultStreamingStrategy.streamSome(bigSegment, AWSXRay.getGlobalRecorder().getEmitter());
    Assert.assertTrue(bigSegment.getReferenceCount() == 0);
}
 

@Test
public void testNestedSubsegmentInjection() {
    SLF4JSegmentListener listener = (SLF4JSegmentListener) AWSXRay.getGlobalRecorder().getSegmentListeners().get(0);
    listener.setPrefix("");
    Segment seg = new SegmentImpl(AWSXRay.getGlobalRecorder(), "test", traceID);
    listener.onSetEntity(null, seg);
    Subsegment sub1 = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "test1", seg);
    listener.onSetEntity(seg, sub1);
    Subsegment sub2 = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "test2", seg);
    listener.onSetEntity(sub1, sub2);

    Assert.assertEquals(traceID.toString() + "@" + sub2.getId(), MDC.get(TRACE_ID_KEY));
}
 

@Test
public void testSubsegmentListeners() {
    AWSXRay.beginSegment("test");
    Subsegment sub = AWSXRay.beginSubsegment("testSub");
    String beginAnnotation = sub.getAnnotations().get("subAnnotation1").toString();
    AWSXRay.endSubsegment();
    String endAnnotation = sub.getAnnotations().get("subAnnotation2").toString();

    Assert.assertEquals("began", beginAnnotation);
    Assert.assertEquals("ended", endAnnotation);
}
 

@Test
public void testDefaultStreamingStrategyRequiresStreaming() {
    DefaultStreamingStrategy defaultStreamingStrategy = new DefaultStreamingStrategy(1);

    Segment smallSegment = new SegmentImpl(AWSXRay.getGlobalRecorder(), "small");
    Assert.assertFalse(defaultStreamingStrategy.requiresStreaming(smallSegment));

    Segment bigSegment = new SegmentImpl(AWSXRay.getGlobalRecorder(), "big");
    bigSegment.addSubsegment(new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "big_child", bigSegment));
    bigSegment.addSubsegment(new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "big_child", bigSegment));
    Assert.assertTrue(defaultStreamingStrategy.requiresStreaming(bigSegment));
}
 

@Test
public void testServletUsesPassedInRecorder() throws IOException, ServletException {
    AWSXRayRecorder customRecorder = Mockito.spy(getMockRecorder());
    AWSXRayServletFilter servletFilter = new AWSXRayServletFilter(new FixedSegmentNamingStrategy("test"), customRecorder);

    AsyncContext asyncContext = Mockito.mock(AsyncContext.class);

    HttpServletRequest request = Mockito.mock(HttpServletRequest.class);
    Mockito.when(request.getRequestURL()).thenReturn(new StringBuffer("test_url"));
    Mockito.when(request.getMethod()).thenReturn("TEST_METHOD");
    Mockito.when(request.isAsyncStarted()).thenReturn(true);
    Mockito.when(request.getAsyncContext()).thenReturn(asyncContext);

    HttpServletResponse response = Mockito.mock(HttpServletResponse.class);

    FilterChain chain = mockChain(request, response);

    AsyncEvent event = Mockito.mock(AsyncEvent.class);
    Mockito.when(event.getSuppliedRequest()).thenReturn(request);
    Mockito.when(event.getSuppliedResponse()).thenReturn(response);

    servletFilter.doFilter(request, response, chain);
    Assert.assertNull(AWSXRay.getTraceEntity());

    AWSXRayServletAsyncListener listener = (AWSXRayServletAsyncListener) Whitebox.getInternalState(servletFilter, "listener");
    listener.onComplete(event);

    Mockito.verify(customRecorder.getEmitter(), Mockito.times(1)).sendSegment(Mockito.any());
}
 

@Test
public void testServletLazilyLoadsRecorder() throws IOException, ServletException {
    AWSXRayServletFilter servletFilter = new AWSXRayServletFilter("test");

    AsyncContext asyncContext = Mockito.mock(AsyncContext.class);
    AWSXRayRecorder customRecorder = Mockito.spy(getMockRecorder());
    AWSXRay.setGlobalRecorder(customRecorder);

    HttpServletRequest request = Mockito.mock(HttpServletRequest.class);
    Mockito.when(request.getRequestURL()).thenReturn(new StringBuffer("test_url"));
    Mockito.when(request.getMethod()).thenReturn("TEST_METHOD");
    Mockito.when(request.isAsyncStarted()).thenReturn(true);
    Mockito.when(request.getAsyncContext()).thenReturn(asyncContext);

    HttpServletResponse response = Mockito.mock(HttpServletResponse.class);

    FilterChain chain = mockChain(request, response);

    AsyncEvent event = Mockito.mock(AsyncEvent.class);
    Mockito.when(event.getSuppliedRequest()).thenReturn(request);
    Mockito.when(event.getSuppliedResponse()).thenReturn(response);

    servletFilter.doFilter(request, response, chain);
    Assert.assertNull(AWSXRay.getTraceEntity());

    AWSXRayServletAsyncListener listener = (AWSXRayServletAsyncListener) Whitebox.getInternalState(servletFilter, "listener");
    listener.onComplete(event);

    Mockito.verify(customRecorder.getEmitter(), Mockito.times(1)).sendSegment(Mockito.any());
}
 

@Test
public void testCaptureSqlException() throws Exception {
    SQLException exception = new SQLException("foo");
    when(delegate.execute(SQL)).thenThrow(exception);
    try {
        statement.execute(SQL);
        fail("Expected exception is not thrown");
    } catch (SQLException th) {
        assertEquals(exception, th);
    } finally {
        assertEquals(exception, AWSXRay.getCurrentSegment().getSubsegments().get(0).getCause().getExceptions().get(0)
                                       .getThrowable());
        assertSubsegment();
    }
}
 

@Test
public void testCaptureRuntimeException() throws Exception {
    RuntimeException exception = new RuntimeException("foo");
    when(delegate.execute(SQL)).thenThrow(exception);
    try {
        statement.execute(SQL);
        fail("Expected exception is not thrown");
    } catch (RuntimeException th) {
        assertEquals(exception, th);
    } finally {
        assertEquals(exception, AWSXRay.getCurrentSegment().getSubsegments().get(0).getCause().getExceptions().get(0)
                                       .getThrowable());
        assertSubsegment();
    }
}
 

private Subsegment createSubsegment() {
    try {
        Connection connection = delegate.getConnection();
        DatabaseMetaData metadata = connection.getMetaData();
        String subsegmentName = DEFAULT_DATABASE_NAME;
        try {
            URI normalizedUri = new URI(new URI(metadata.getURL()).getSchemeSpecificPart());
            subsegmentName = connection.getCatalog() + "@" + normalizedUri.getHost();
        } catch (URISyntaxException e) {
            logger.warn("Unable to parse database URI. Falling back to default '" + DEFAULT_DATABASE_NAME
                        + "' for subsegment name.", e);
        }

        Subsegment subsegment = AWSXRay.beginSubsegment(subsegmentName);
        if (subsegment == null) {
            return null;
        }

        subsegment.setNamespace(Namespace.REMOTE.toString());
        Map<String, Object> sqlParams = new HashMap<>();
        sqlParams.put(URL, metadata.getURL());
        sqlParams.put(USER, metadata.getUserName());
        sqlParams.put(DRIVER_VERSION, metadata.getDriverVersion());
        sqlParams.put(DATABASE_TYPE, metadata.getDatabaseProductName());
        sqlParams.put(DATABASE_VERSION, metadata.getDatabaseProductVersion());
        subsegment.putAllSql(sqlParams);

        return subsegment;
    } catch (SQLException exception) {
        logger.warn("Failed to create X-Ray subsegment for the statement execution.", exception);
        return null;
    }
}
 

@Override
public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
    Subsegment subsegment = null;

    if (isExecution(method)) {
        // only trace on execution methods
        subsegment = createSubsegment();
    }

    logger.debug(
        String.format("Invoking statement execution with X-Ray tracing. Tracing active: %s", subsegment != null));
    try {
        // execute the query "wrapped" in a XRay Subsegment
        return method.invoke(delegate, args);
    } catch (Throwable t) {
        Throwable rootThrowable = t;
        if (t instanceof InvocationTargetException) {
            // the reflection may wrap the actual error with an InvocationTargetException.
            // we want to use the root cause to make the instrumentation seamless
            InvocationTargetException ite = (InvocationTargetException) t;
            if (ite.getTargetException() != null) {
                rootThrowable = ite.getTargetException();
            } else if (ite.getCause() != null) {
                rootThrowable = ite.getCause();
            }
        }

        if (subsegment != null) {
            subsegment.addException(rootThrowable);
        }
        throw rootThrowable;
    } finally {
        if (subsegment != null && isExecution(method)) {
            AWSXRay.endSubsegment();
        }
    }
}
 

@Override
protected void configure() {
  // Disable AWS x-ray in integration tests.
  // See doc: https://docs.aws.amazon.com/xray/latest/devguide/xray-sdk-java-configuration.html#xray-sdk-java-configuration-sysprops
  AWSXRay.setGlobalRecorder(AWSXRayRecorderBuilder.standard().withContextMissingStrategy(new LogErrorContextMissingStrategy()).build());
  bind(ScenarioScope.class).toInstance(new SequentialScenarioScope());
}
 

@Test
public void testSubsegmentInjection() {
    SLF4JSegmentListener listener = (SLF4JSegmentListener) AWSXRay.getGlobalRecorder().getSegmentListeners().get(0);
    listener.setPrefix("");
    Segment seg = new SegmentImpl(AWSXRay.getGlobalRecorder(), "test", traceID);
    listener.onSetEntity(null, seg);
    Subsegment sub = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "test", seg);
    listener.onSetEntity(seg, sub);

    Assert.assertEquals(traceID.toString() + "@" + sub.getId(), MDC.get(TRACE_ID_KEY));
}
 

@Test
public void testUnsampledSegmentInjection() {
    SLF4JSegmentListener listener = (SLF4JSegmentListener) AWSXRay.getGlobalRecorder().getSegmentListeners().get(0);
    listener.setPrefix("");
    Segment seg = new SegmentImpl(AWSXRay.getGlobalRecorder(), "test", traceID);
    seg.setSampled(false);
    listener.onSetEntity(null, seg);

    Assert.assertNull(MDC.get(TRACE_ID_KEY));
}
 

@Test
public void testSegmentInjection() {
    SLF4JSegmentListener listener = (SLF4JSegmentListener) AWSXRay.getGlobalRecorder().getSegmentListeners().get(0);
    listener.setPrefix("");
    Segment seg = new SegmentImpl(AWSXRay.getGlobalRecorder(), "test", traceID);
    listener.onSetEntity(null, seg);

    Assert.assertEquals(traceID.toString() + "@" + seg.getId(), MDC.get(TRACE_ID_KEY));
}
 

@Test
public void testSetPrefix() {
    SLF4JSegmentListener listener = (SLF4JSegmentListener) AWSXRay.getGlobalRecorder().getSegmentListeners().get(0);
    listener.setPrefix("");
    Segment seg = new SegmentImpl(AWSXRay.getGlobalRecorder(), "test", traceID);

    listener.onSetEntity(null, seg);

    Assert.assertEquals(traceID.toString() + "@" + seg.getId(), MDC.get(TRACE_ID_KEY));
}
 

@Test
public void testDefaultPrefix() {
    SLF4JSegmentListener listener = (SLF4JSegmentListener) AWSXRay.getGlobalRecorder().getSegmentListeners().get(0);
    Segment seg = new SegmentImpl(AWSXRay.getGlobalRecorder(), "test", traceID);

    listener.onSetEntity(null, seg);

    Assert.assertEquals(TRACE_ID_KEY + ": " + traceID.toString() + "@" + seg.getId(), MDC.get(TRACE_ID_KEY));
}