Skip to content

Commit 733edbf

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

File tree

9 files changed

+372
-13
lines changed

9 files changed

+372
-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: 189 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,194 @@
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(map ->
60+
assertThat(map.get("service_undefined"))
61+
.containsEntry("key", "val")
62+
);
63+
}
64+
65+
@Test
66+
void shouldNotSetMetaDataIfNoCurrentSubSegment() {
67+
PowerTracer.putMetadata("key", "val");
68+
69+
assertThat(AWSXRay.getTraceEntity().getAnnotations())
70+
.isEmpty();
71+
}
72+
73+
@Test
74+
void shouldInvokeCodeBlockWrappedWithinSubsegment() {
75+
Context test = mock(Context.class);
76+
77+
PowerTracer.withSubsegment("testSubSegment", subsegment -> {
78+
subsegment.putAnnotation("key", "val");
79+
subsegment.putMetadata("key", "val");
80+
test.getFunctionName();
81+
});
82+
83+
verify(test).getFunctionName();
84+
85+
assertThat(AWSXRay.getTraceEntity().getSubsegments())
86+
.hasSize(1)
87+
.allSatisfy(subsegment -> {
88+
assertThat(subsegment.getName())
89+
.isEqualTo("## testSubSegment");
90+
91+
assertThat(subsegment.getNamespace())
92+
.isEqualTo("service_undefined");
93+
94+
assertThat(subsegment.getAnnotations())
95+
.hasSize(1)
96+
.containsEntry("key", "val");
97+
98+
assertThat(subsegment.getMetadata())
99+
.hasSize(1);
100+
});
101+
}
102+
103+
@Test
104+
void shouldInvokeCodeBlockWrappedWithinNamespacedSubsegment() {
105+
Context test = mock(Context.class);
106+
107+
PowerTracer.withSubsegment("testNamespace", "testSubSegment", subsegment -> {
108+
subsegment.putAnnotation("key", "val");
109+
subsegment.putMetadata("key", "val");
110+
test.getFunctionName();
111+
});
112+
113+
verify(test).getFunctionName();
114+
115+
assertThat(AWSXRay.getTraceEntity().getSubsegments())
116+
.hasSize(1)
117+
.allSatisfy(subsegment -> {
118+
assertThat(subsegment.getName())
119+
.isEqualTo("## testSubSegment");
120+
121+
assertThat(subsegment.getNamespace())
122+
.isEqualTo("testNamespace");
123+
124+
assertThat(subsegment.getAnnotations())
125+
.hasSize(1)
126+
.containsEntry("key", "val");
127+
128+
assertThat(subsegment.getMetadata())
129+
.hasSize(1);
130+
});
131+
}
132+
133+
@Test
134+
void shouldInvokeCodeBlockWrappedWithinEntitySubsegment() throws InterruptedException {
135+
Context test = mock(Context.class);
136+
137+
Entity traceEntity = AWSXRay.getTraceEntity();
138+
139+
Thread thread = new Thread(() -> withEntitySubsegment("testSubSegment", traceEntity, subsegment -> {
140+
subsegment.putAnnotation("key", "val");
141+
test.getFunctionName();
142+
}));
143+
144+
thread.start();
145+
thread.join();
146+
147+
verify(test).getFunctionName();
148+
149+
assertThat(AWSXRay.getTraceEntity().getSubsegments())
150+
.hasSize(1)
151+
.allSatisfy(subsegment -> {
152+
assertThat(subsegment.getName())
153+
.isEqualTo("## testSubSegment");
154+
155+
assertThat(subsegment.getNamespace())
156+
.isEqualTo("service_undefined");
157+
158+
assertThat(subsegment.getAnnotations())
159+
.hasSize(1)
160+
.containsEntry("key", "val");
161+
});
162+
}
163+
164+
@Test
165+
void shouldInvokeCodeBlockWrappedWithinNamespacedEntitySubsegment() throws InterruptedException {
166+
Context test = mock(Context.class);
167+
168+
Entity traceEntity = AWSXRay.getTraceEntity();
169+
170+
Thread thread = new Thread(() -> withEntitySubsegment("testNamespace", "testSubSegment", traceEntity, subsegment -> {
171+
subsegment.putAnnotation("key", "val");
172+
test.getFunctionName();
173+
}));
174+
175+
thread.start();
176+
thread.join();
177+
178+
verify(test).getFunctionName();
179+
180+
assertThat(AWSXRay.getTraceEntity().getSubsegments())
181+
.hasSize(1)
182+
.allSatisfy(subsegment -> {
183+
assertThat(subsegment.getName())
184+
.isEqualTo("## testSubSegment");
185+
186+
assertThat(subsegment.getNamespace())
187+
.isEqualTo("testNamespace");
188+
189+
assertThat(subsegment.getAnnotations())
190+
.hasSize(1)
191+
.containsEntry("key", "val");
192+
});
193+
}
5194
}

0 commit comments

Comments
 (0)