Java Code Examples for com.amazonaws.xray.AWSXRay#beginSegment()

private Segment segmentInResponseToCode(int code) {
    NoOpResponseHandler responseHandler = new NoOpResponseHandler();
    TracedResponseHandler<String> tracedResponseHandler = new TracedResponseHandler<>(responseHandler);
    HttpResponse httpResponse = new BasicHttpResponse(new BasicStatusLine(HttpVersion.HTTP_1_1, code, ""));

    Segment segment = AWSXRay.beginSegment("test");
    AWSXRay.beginSubsegment("someHttpCall");

    try {
        tracedResponseHandler.handleResponse(httpResponse);
    } catch (IOException e) {
        throw new RuntimeException(e);
    }

    AWSXRay.endSubsegment();
    AWSXRay.endSegment();
    return segment;
}
 

@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 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 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 testLambdaInvokeSubsegmentContainsFunctionName() {
    // Setup test
    AWSLambda lambda = AWSLambdaClientBuilder
        .standard()
        .withRequestHandlers(new TracingHandler())
        .withRegion(Regions.US_EAST_1)
        .withCredentials(new AWSStaticCredentialsProvider(new BasicAWSCredentials("fake", "fake")))
        .build();
    mockHttpClient(lambda, "null"); // Lambda returns "null" on successful fn. with no return value

    // Test logic
    Segment segment = AWSXRay.beginSegment("test");

    InvokeRequest request = new InvokeRequest();
    request.setFunctionName("testFunctionName");
    lambda.invoke(request);

    Assert.assertEquals(1, segment.getSubsegments().size());
    Assert.assertEquals("Invoke", segment.getSubsegments().get(0).getAws().get("operation"));
    Assert.assertEquals("testFunctionName", segment.getSubsegments().get(0).getAws().get("function_name"));
}
 

@Test
public void testS3PutObjectSubsegmentContainsBucketName() {
    // Setup test
    AmazonS3 s3 = AmazonS3ClientBuilder
        .standard()
        .withRequestHandlers(new TracingHandler())
        .withRegion(Regions.US_EAST_1)
        .withCredentials(new AWSStaticCredentialsProvider(new BasicAWSCredentials("fake", "fake")))
        .build();
    mockHttpClient(s3, null);

    String bucket = "test-bucket";
    String key = "test-key";
    // Test logic 
    Segment segment = AWSXRay.beginSegment("test");
    s3.putObject(bucket, key, "This is a test from java");
    Assert.assertEquals(1, segment.getSubsegments().size());
    Assert.assertEquals("PutObject", segment.getSubsegments().get(0).getAws().get("operation"));
    Assert.assertEquals(bucket, segment.getSubsegments().get(0).getAws().get("bucket_name"));
    Assert.assertEquals(key, segment.getSubsegments().get(0).getAws().get("key"));
}
 

@Test
public void testSNSPublish() {
    // Setup test
    // reference : https://docs.aws.amazon.com/sns/latest/api/API_Publish.html
    final String publishResponse = "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n" +
                                   "<PublishResponse xmlns=\"http://sns.amazonaws.com/doc/2010-03-31/\">" +
                                   "<PublishResult><MessageId>94f20ce6-13c5-43a0-9a9e-ca52d816e90b</MessageId>" +
                                   "</PublishResult>" +
                                   "</PublishResponse>";
    final String topicArn = "testTopicArn";
    AmazonSNS sns = AmazonSNSClientBuilder
        .standard()
        .withRequestHandlers(new TracingHandler())
        .withRegion(Regions.US_EAST_1)
        .withCredentials(new AWSStaticCredentialsProvider(new BasicAWSCredentials("fake", "fake")))
        .build();
    mockHttpClient(sns, publishResponse);
    // Test logic 
    Segment segment = AWSXRay.beginSegment("test");
    sns.publish(new PublishRequest(topicArn, "testMessage"));
    Assert.assertEquals(1, segment.getSubsegments().size());
    Assert.assertEquals("Publish", segment.getSubsegments().get(0).getAws().get("operation"));
    Assert.assertEquals(topicArn, segment.getSubsegments().get(0).getAws().get("topic_arn"));
}
 

@Test
public void testShouldNotTraceXRaySamplingOperations() {
    com.amazonaws.services.xray.AWSXRay xray = AWSXRayClientBuilder
        .standard()
        .withRequestHandlers(new TracingHandler()).withRegion(Regions.US_EAST_1)
        .withCredentials(new AWSStaticCredentialsProvider(new BasicAWSCredentials("fake", "fake")))
        .build();
    mockHttpClient(xray, null);

    Segment segment = AWSXRay.beginSegment("test");
    xray.getSamplingRules(new GetSamplingRulesRequest());
    Assert.assertEquals(0, segment.getSubsegments().size());

    xray.getSamplingTargets(new GetSamplingTargetsRequest());
    Assert.assertEquals(0, segment.getSubsegments().size());
}
 

@Before
public void setup() {
    Emitter blankEmitter = Mockito.mock(Emitter.class);
    Mockito.doReturn(true).when(blankEmitter).sendSegment(Mockito.anyObject());

    AWSXRay.setGlobalRecorder(
            AWSXRayRecorderBuilder.standard()
                    .withEmitter(blankEmitter)
                    .build()
    );
    AWSXRay.clearTraceEntity();
    AWSXRay.beginSegment("test");
}
 

@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 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 testS3CopyObjectSubsegmentContainsBucketName() {
    // Setup test
    final String copyResponse = "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n" +
                                "<CopyObjectResult xmlns=\"http://s3.amazonaws.com/doc/2006-03-01/\">" +
                                "<LastModified>2018-01-21T10:09:54.000Z</LastModified><ETag>" +
                                "&quot;31748afd7b576119d3c2471f39fc7a55&quot;</ETag>" +
                                "</CopyObjectResult>";
    AmazonS3 s3 = AmazonS3ClientBuilder
        .standard()
        .withRequestHandlers(new TracingHandler())
        .withRegion(Regions.US_EAST_1)
        .withCredentials(new AWSStaticCredentialsProvider(new BasicAWSCredentials("fake", "fake")))
        .build();
    mockHttpClient(s3, copyResponse);

    String bucket = "source-bucket";
    String key = "source-key";
    String dstBucket = "dest-bucket";
    String dstKey = "dest-key";
    // Test logic 
    Segment segment = AWSXRay.beginSegment("test");
    s3.copyObject(bucket, key, dstBucket, dstKey);
    Assert.assertEquals(1, segment.getSubsegments().size());
    Assert.assertEquals("CopyObject", segment.getSubsegments().get(0).getAws().get("operation"));
    Assert.assertEquals(bucket, segment.getSubsegments().get(0).getAws().get("source_bucket_name"));
    Assert.assertEquals(key, segment.getSubsegments().get(0).getAws().get("source_key"));
    Assert.assertEquals(dstBucket, segment.getSubsegments().get(0).getAws().get("destination_bucket_name"));
    Assert.assertEquals(dstKey, segment.getSubsegments().get(0).getAws().get("destination_key"));
}
 

@Test
void invokeTest() {
  AWSXRay.beginSegment("blank-java-test");
  String path = "src/test/resources/event.json";
  String eventString = loadJsonFile(path);
  SQSEvent event = gson.fromJson(eventString, SQSEvent.class);
  Context context = new TestContext();
  String requestId = context.getAwsRequestId();
  Handler handler = new Handler();
  String result = handler.handleRequest(event, context);
  assertTrue(result.contains("totalCodeSize"));
  AWSXRay.endSegment();
}
 

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

    Segment segment = AWSXRay.beginSegment("big");


    Subsegment subsegment = new SubsegmentImpl(AWSXRay.getGlobalRecorder(), "subsegment1", segment);
    subsegment.setStartTime(1.0);
    segment.addSubsegment(subsegment);
    subsegment.end();

    Thread thread1 = new Thread(() -> {
        AWSXRay.setTraceEntity(segment);
        AWSXRay.beginSubsegment("thread1");
        AWSXRay.endSubsegment();
    });
    Thread thread2 = new Thread(() -> {
        AWSXRay.setTraceEntity(segment);
        AWSXRay.beginSubsegment("thread2");
        AWSXRay.endSubsegment();
    });

    thread1.start();
    thread2.start();
    for (Thread thread : new Thread[]{thread1, thread2}) {
        try {
            thread.join();
        } catch (InterruptedException e) {
            return;
        }
    }

    Assert.assertTrue(AWSXRay.getTraceEntity().getName().equals("big"));
    //asserts that all subsegments are added correctly.
    Assert.assertTrue(AWSXRay.getTraceEntity().getTotalSize().intValue() == 3);

    defaultStreamingStrategy.streamSome(segment, AWSXRay.getGlobalRecorder().getEmitter());

    Assert.assertTrue(segment.getTotalSize().intValue() == 0);
}