Skip to content

Commit bad87a6

Browse files
author
Pankaj Agrawal
committed
test cases for the aspect and utils
1 parent 51d2b07 commit bad87a6

File tree

9 files changed

+370
-13
lines changed

9 files changed

+370
-13
lines changed

src/main/java/software/aws/lambda/tracing/PowerTracer.java

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,11 +29,21 @@ public static void putMetadata(String namespace, String key, Object value) {
2929

3030
public static void withEntitySubsegment(String name, Entity entity, Consumer<Subsegment> subsegment) {
3131
AWSXRay.setTraceEntity(entity);
32-
withSubsegment(name, subsegment);
32+
withEntitySubsegment(SERVICE_NAME, name, entity, subsegment);
33+
}
34+
35+
public static void withEntitySubsegment(String namespace, String name, Entity entity, Consumer<Subsegment> subsegment) {
36+
AWSXRay.setTraceEntity(entity);
37+
withSubsegment(namespace, name, subsegment);
3338
}
3439

3540
public static void withSubsegment(String name, Consumer<Subsegment> subsegment) {
41+
withSubsegment(SERVICE_NAME, name, subsegment);
42+
}
43+
44+
public static void withSubsegment(String namespace, String name, Consumer<Subsegment> subsegment) {
3645
Subsegment segment = AWSXRay.beginSubsegment("## " + name);
46+
segment.setNamespace(namespace);
3747
try {
3848
subsegment.accept(segment);
3949
} finally {

src/main/java/software/aws/lambda/tracing/internal/LambdaTracingAspect.java

Lines changed: 1 addition & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ public Object around(ProceedingJoinPoint pjp,
4242
try {
4343
Object methodReturn = pjp.proceed(proceedArgs);
4444
if (powerToolsTracing.captureResponse()) {
45-
segment.putMetadata(namespace(powerToolsTracing), pjp.getSignature().getName() + " response", response(pjp, methodReturn));
45+
segment.putMetadata(namespace(powerToolsTracing), pjp.getSignature().getName() + " response", methodReturn);
4646
}
4747

4848
IS_COLD_START = false;
@@ -58,15 +58,6 @@ public Object around(ProceedingJoinPoint pjp,
5858
}
5959
}
6060

61-
private Object response(ProceedingJoinPoint pjp, Object methodReturn) {
62-
// TODO should we try to parse output stream? or just not support it?
63-
if (placedOnStreamHandler(pjp)) {
64-
65-
}
66-
67-
return methodReturn;
68-
}
69-
7061
private String namespace(PowerToolTracing powerToolsTracing) {
7162
return powerToolsTracing.namespace().isEmpty() ? SERVICE_NAME : powerToolsTracing.namespace();
7263
}

src/test/java/software/aws/lambda/handlers/PowerTracerToolEnabled.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
public class PowerTracerToolEnabled implements RequestHandler<Object, Object> {
88

99
@Override
10-
@PowerToolTracing
10+
@PowerToolTracing(namespace = "lambdaHandler")
1111
public Object handleRequest(Object input, Context context) {
1212
return null;
1313
}
Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
package software.aws.lambda.handlers;
2+
3+
import java.io.InputStream;
4+
import java.io.OutputStream;
5+
6+
import com.amazonaws.services.lambda.runtime.Context;
7+
import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
8+
import software.aws.lambda.tracing.PowerToolTracing;
9+
10+
public class PowerTracerToolEnabledForStream implements RequestStreamHandler {
11+
12+
@Override
13+
@PowerToolTracing(namespace = "streamHandler")
14+
public void handleRequest(InputStream input, OutputStream output, Context context) {
15+
16+
}
17+
}
Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
package software.aws.lambda.handlers;
2+
3+
import java.io.InputStream;
4+
import java.io.OutputStream;
5+
6+
import com.amazonaws.services.lambda.runtime.Context;
7+
import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
8+
import software.aws.lambda.tracing.PowerToolTracing;
9+
10+
public class PowerTracerToolEnabledForStreamWithNoMetaData implements RequestStreamHandler {
11+
12+
@Override
13+
@PowerToolTracing(captureResponse = false, captureError = false)
14+
public void handleRequest(InputStream input, OutputStream output, Context context) {
15+
16+
}
17+
}
Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
package software.aws.lambda.handlers;
2+
3+
import com.amazonaws.services.lambda.runtime.Context;
4+
import com.amazonaws.services.lambda.runtime.RequestHandler;
5+
import software.aws.lambda.tracing.PowerToolTracing;
6+
7+
public class PowerTracerToolEnabledWithException implements RequestHandler<Object, Object> {
8+
9+
@Override
10+
@PowerToolTracing(namespace = "lambdaHandler")
11+
public Object handleRequest(Object input, Context context) {
12+
throw new RuntimeException("I am failing!");
13+
}
14+
}
Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
package software.aws.lambda.handlers;
2+
3+
import com.amazonaws.services.lambda.runtime.Context;
4+
import com.amazonaws.services.lambda.runtime.RequestHandler;
5+
import software.aws.lambda.tracing.PowerToolTracing;
6+
7+
public class PowerTracerToolEnabledWithNoMetaData implements RequestHandler<Object, Object> {
8+
9+
@Override
10+
@PowerToolTracing(captureResponse = false, captureError = false)
11+
public Object handleRequest(Object input, Context context) {
12+
return null;
13+
}
14+
}
Lines changed: 187 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,192 @@
11
package software.aws.lambda.tracing;
22

3+
import com.amazonaws.services.lambda.runtime.Context;
4+
import com.amazonaws.xray.AWSXRay;
5+
import com.amazonaws.xray.entities.Entity;
6+
import org.junit.jupiter.api.AfterEach;
7+
import org.junit.jupiter.api.BeforeEach;
8+
import org.junit.jupiter.api.Test;
9+
10+
import static org.assertj.core.api.Assertions.assertThat;
11+
import static org.mockito.Mockito.mock;
12+
import static org.mockito.Mockito.verify;
13+
import static software.aws.lambda.tracing.PowerTracer.withEntitySubsegment;
14+
315
class PowerTracerTest {
416

17+
@BeforeEach
18+
void setUp() {
19+
AWSXRay.beginSegment("test");
20+
}
21+
22+
@AfterEach
23+
void tearDown() {
24+
if (AWSXRay.getCurrentSubsegmentOptional().isPresent()) {
25+
AWSXRay.endSubsegment();
26+
}
27+
28+
AWSXRay.endSegment();
29+
}
30+
31+
@Test
32+
void shouldSetAnnotationOnCurrentSubSegment() {
33+
AWSXRay.beginSubsegment("subSegment");
34+
35+
PowerTracer.putAnnotation("key", "val");
36+
37+
assertThat(AWSXRay.getTraceEntity().getAnnotations())
38+
.hasSize(1)
39+
.containsEntry("key", "val");
40+
}
41+
42+
@Test
43+
void shouldNotSetAnnotationIfNoCurrentSubSegment() {
44+
PowerTracer.putAnnotation("key", "val");
45+
46+
assertThat(AWSXRay.getTraceEntity().getAnnotations())
47+
.isEmpty();
48+
}
49+
50+
@Test
51+
void shouldSetMetadataOnCurrentSubSegment() {
52+
AWSXRay.beginSubsegment("subSegment");
53+
54+
PowerTracer.putMetadata("key", "val");
55+
56+
assertThat(AWSXRay.getTraceEntity().getMetadata())
57+
.hasSize(1)
58+
.containsKey("service_undefined")
59+
.satisfies(stringMapMap -> assertThat(stringMapMap.get("service_undefined"))
60+
.containsEntry("key", "val"));
61+
}
62+
63+
@Test
64+
void shouldNotSetMetaDataIfNoCurrentSubSegment() {
65+
PowerTracer.putMetadata("key", "val");
66+
67+
assertThat(AWSXRay.getTraceEntity().getAnnotations())
68+
.isEmpty();
69+
}
70+
71+
@Test
72+
void shouldInvokeCodeBlockWrappedWithinSubsegment() {
73+
Context test = mock(Context.class);
74+
75+
PowerTracer.withSubsegment("testSubSegment", subsegment -> {
76+
subsegment.putAnnotation("key", "val");
77+
subsegment.putMetadata("key", "val");
78+
test.getFunctionName();
79+
});
80+
81+
verify(test).getFunctionName();
82+
83+
assertThat(AWSXRay.getTraceEntity().getSubsegments())
84+
.hasSize(1)
85+
.allSatisfy(subsegment -> {
86+
assertThat(subsegment.getName())
87+
.isEqualTo("## testSubSegment");
88+
89+
assertThat(subsegment.getNamespace())
90+
.isEqualTo("service_undefined");
91+
92+
assertThat(subsegment.getAnnotations())
93+
.hasSize(1)
94+
.containsEntry("key", "val");
95+
96+
assertThat(subsegment.getMetadata())
97+
.hasSize(1);
98+
});
99+
}
100+
101+
@Test
102+
void shouldInvokeCodeBlockWrappedWithinNamespacedSubsegment() {
103+
Context test = mock(Context.class);
104+
105+
PowerTracer.withSubsegment("testNamespace", "testSubSegment", subsegment -> {
106+
subsegment.putAnnotation("key", "val");
107+
subsegment.putMetadata("key", "val");
108+
test.getFunctionName();
109+
});
110+
111+
verify(test).getFunctionName();
112+
113+
assertThat(AWSXRay.getTraceEntity().getSubsegments())
114+
.hasSize(1)
115+
.allSatisfy(subsegment -> {
116+
assertThat(subsegment.getName())
117+
.isEqualTo("## testSubSegment");
118+
119+
assertThat(subsegment.getNamespace())
120+
.isEqualTo("testNamespace");
121+
122+
assertThat(subsegment.getAnnotations())
123+
.hasSize(1)
124+
.containsEntry("key", "val");
125+
126+
assertThat(subsegment.getMetadata())
127+
.hasSize(1);
128+
});
129+
}
130+
131+
@Test
132+
void shouldInvokeCodeBlockWrappedWithinEntitySubsegment() throws InterruptedException {
133+
Context test = mock(Context.class);
134+
135+
Entity traceEntity = AWSXRay.getTraceEntity();
136+
137+
Thread thread = new Thread(() -> withEntitySubsegment("testSubSegment", traceEntity, subsegment -> {
138+
subsegment.putAnnotation("key", "val");
139+
test.getFunctionName();
140+
}));
141+
142+
thread.start();
143+
thread.join();
144+
145+
verify(test).getFunctionName();
146+
147+
assertThat(AWSXRay.getTraceEntity().getSubsegments())
148+
.hasSize(1)
149+
.allSatisfy(subsegment -> {
150+
assertThat(subsegment.getName())
151+
.isEqualTo("## testSubSegment");
152+
153+
assertThat(subsegment.getNamespace())
154+
.isEqualTo("service_undefined");
155+
156+
assertThat(subsegment.getAnnotations())
157+
.hasSize(1)
158+
.containsEntry("key", "val");
159+
});
160+
}
161+
162+
@Test
163+
void shouldInvokeCodeBlockWrappedWithinNamespacedEntitySubsegment() throws InterruptedException {
164+
Context test = mock(Context.class);
165+
166+
Entity traceEntity = AWSXRay.getTraceEntity();
167+
168+
Thread thread = new Thread(() -> withEntitySubsegment("testNamespace", "testSubSegment", traceEntity, subsegment -> {
169+
subsegment.putAnnotation("key", "val");
170+
test.getFunctionName();
171+
}));
172+
173+
thread.start();
174+
thread.join();
175+
176+
verify(test).getFunctionName();
177+
178+
assertThat(AWSXRay.getTraceEntity().getSubsegments())
179+
.hasSize(1)
180+
.allSatisfy(subsegment -> {
181+
assertThat(subsegment.getName())
182+
.isEqualTo("## testSubSegment");
183+
184+
assertThat(subsegment.getNamespace())
185+
.isEqualTo("testNamespace");
186+
187+
assertThat(subsegment.getAnnotations())
188+
.hasSize(1)
189+
.containsEntry("key", "val");
190+
});
191+
}
5192
}

0 commit comments

Comments
 (0)