33
33
import org .junit .jupiter .api .BeforeAll ;
34
34
import org .junit .jupiter .api .BeforeEach ;
35
35
import org .junit .jupiter .api .Test ;
36
+ import org .junitpioneer .jupiter .SetEnvironmentVariable ;
36
37
import org .mockito .Mock ;
37
38
import org .mockito .MockedStatic ;
38
39
import software .amazon .lambda .powertools .common .internal .LambdaHandlerProcessor ;
@@ -113,6 +114,28 @@ void shouldCaptureNonHandlerMethodWithCustomSegmentName() {
113
114
void shouldCaptureTraces () {
114
115
requestHandler .handleRequest (new Object (), context );
115
116
117
+ assertThat (AWSXRay .getTraceEntity ())
118
+ .isNotNull ();
119
+
120
+ assertThat (AWSXRay .getTraceEntity ().getSubsegmentsCopy ())
121
+ .hasSize (1 )
122
+ .allSatisfy (subsegment ->
123
+ {
124
+ assertThat (subsegment .getAnnotations ())
125
+ .hasSize (2 )
126
+ .containsEntry ("ColdStart" , true )
127
+ .containsEntry ("Service" , "lambdaHandler" );
128
+
129
+ assertThat (subsegment .getMetadata ())
130
+ .hasSize (0 );
131
+ });
132
+ }
133
+
134
+ @ Test
135
+ @ SetEnvironmentVariable (key = "POWERTOOLS_TRACER_CAPTURE_RESPONSE" , value = "true" )
136
+ void shouldCaptureTracesWithResponseMetadata () {
137
+ requestHandler .handleRequest (new Object (), context );
138
+
116
139
assertThat (AWSXRay .getTraceEntity ())
117
140
.isNotNull ();
118
141
@@ -132,6 +155,7 @@ void shouldCaptureTraces() {
132
155
}
133
156
134
157
@ Test
158
+ @ SetEnvironmentVariable (key = "POWERTOOLS_TRACER_CAPTURE_ERROR" , value = "true" )
135
159
void shouldCaptureTracesWithExceptionMetaData () {
136
160
requestHandler = new PowerTracerToolEnabledWithException ();
137
161
@@ -163,6 +187,25 @@ void shouldCaptureTracesWithExceptionMetaData() {
163
187
void shouldCaptureTracesForStream () throws IOException {
164
188
streamHandler .handleRequest (new ByteArrayInputStream ("test" .getBytes ()), new ByteArrayOutputStream (), context );
165
189
190
+ assertThat (AWSXRay .getTraceEntity ())
191
+ .isNotNull ();
192
+
193
+ assertThat (AWSXRay .getTraceEntity ().getSubsegmentsCopy ())
194
+ .hasSize (1 )
195
+ .allSatisfy (subsegment ->
196
+ {
197
+ assertThat (subsegment .getAnnotations ())
198
+ .hasSize (2 )
199
+ .containsEntry ("ColdStart" , true )
200
+ .containsEntry ("Service" , "streamHandler" );
201
+ });
202
+ }
203
+
204
+ @ Test
205
+ @ SetEnvironmentVariable (key = "POWERTOOLS_TRACER_CAPTURE_RESPONSE" , value = "true" )
206
+ void shouldCaptureTracesForStreamWithResponseMetadata () throws IOException {
207
+ streamHandler .handleRequest (new ByteArrayInputStream ("test" .getBytes ()), new ByteArrayOutputStream (), context );
208
+
166
209
assertThat (AWSXRay .getTraceEntity ())
167
210
.isNotNull ();
168
211
0 commit comments