fix tests

aws-powertools · Jun 25, 2024 · 74fa661 · 74fa661
1 parent 46e76f9
commit 74fa661
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 1 deletion.
diff --git a/powertools-metrics/src/main/java/software/amazon/lambda/powertools/metrics/MetricsUtils.java b/powertools-metrics/src/main/java/software/amazon/lambda/powertools/metrics/MetricsUtils.java
@@ -14,7 +14,6 @@
 
 package software.amazon.lambda.powertools.metrics;
 
-import static java.util.Objects.requireNonNull;
 import static java.util.Optional.ofNullable;
 import static software.amazon.lambda.powertools.common.internal.LambdaHandlerProcessor.getXrayTraceId;
 import static software.amazon.lambda.powertools.metrics.internal.LambdaMetricsAspect.REQUEST_ID_PROPERTY;

diff --git a/powertools-tracing/pom.xml b/powertools-tracing/pom.xml
@@ -78,6 +78,11 @@
             <artifactId>junit-jupiter-engine</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.junit-pioneer</groupId>
+            <artifactId>junit-pioneer</artifactId>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>org.apache.commons</groupId>
             <artifactId>commons-lang3</artifactId>

diff --git a/...test/java/software/amazon/lambda/powertools/tracing/internal/LambdaTracingAspectTest.java b/...test/java/software/amazon/lambda/powertools/tracing/internal/LambdaTracingAspectTest.java
@@ -33,6 +33,7 @@
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junitpioneer.jupiter.SetEnvironmentVariable;
 import org.mockito.Mock;
 import org.mockito.MockedStatic;
 import software.amazon.lambda.powertools.common.internal.LambdaHandlerProcessor;
@@ -113,6 +114,28 @@ void shouldCaptureNonHandlerMethodWithCustomSegmentName() {
     void shouldCaptureTraces() {
         requestHandler.handleRequest(new Object(), context);
 
+        assertThat(AWSXRay.getTraceEntity())
+                .isNotNull();
+
+        assertThat(AWSXRay.getTraceEntity().getSubsegmentsCopy())
+                .hasSize(1)
+                .allSatisfy(subsegment ->
+                {
+                    assertThat(subsegment.getAnnotations())
+                            .hasSize(2)
+                            .containsEntry("ColdStart", true)
+                            .containsEntry("Service", "lambdaHandler");
+
+                    assertThat(subsegment.getMetadata())
+                            .hasSize(0);
+                });
+    }
+
+    @Test
+    @SetEnvironmentVariable(key = "POWERTOOLS_TRACER_CAPTURE_RESPONSE", value = "true")
+    void shouldCaptureTracesWithResponseMetadata() {
+        requestHandler.handleRequest(new Object(), context);
+
         assertThat(AWSXRay.getTraceEntity())
                 .isNotNull();
 
@@ -132,6 +155,7 @@ void shouldCaptureTraces() {
     }
 
     @Test
+    @SetEnvironmentVariable(key = "POWERTOOLS_TRACER_CAPTURE_ERROR", value = "true")
     void shouldCaptureTracesWithExceptionMetaData() {
         requestHandler = new PowerTracerToolEnabledWithException();
 
@@ -163,6 +187,25 @@ void shouldCaptureTracesWithExceptionMetaData() {
     void shouldCaptureTracesForStream() throws IOException {
         streamHandler.handleRequest(new ByteArrayInputStream("test".getBytes()), new ByteArrayOutputStream(), context);
 
+        assertThat(AWSXRay.getTraceEntity())
+                .isNotNull();
+
+        assertThat(AWSXRay.getTraceEntity().getSubsegmentsCopy())
+                .hasSize(1)
+                .allSatisfy(subsegment ->
+                {
+                    assertThat(subsegment.getAnnotations())
+                            .hasSize(2)
+                            .containsEntry("ColdStart", true)
+                            .containsEntry("Service", "streamHandler");
+                });
+    }
+
+    @Test
+    @SetEnvironmentVariable(key = "POWERTOOLS_TRACER_CAPTURE_RESPONSE", value = "true")
+    void shouldCaptureTracesForStreamWithResponseMetadata() throws IOException {
+        streamHandler.handleRequest(new ByteArrayInputStream("test".getBytes()), new ByteArrayOutputStream(), context);
+
         assertThat(AWSXRay.getTraceEntity())
                 .isNotNull();