1+ package software .amazon .lambda .powertools .testsuite ;
2+
3+
4+ import java .io .ByteArrayInputStream ;
5+ import java .io .ByteArrayOutputStream ;
6+ import java .io .IOException ;
7+ import java .lang .reflect .InvocationTargetException ;
8+ import java .lang .reflect .Method ;
9+ import java .nio .channels .FileChannel ;
10+ import java .nio .charset .StandardCharsets ;
11+ import java .nio .file .Files ;
12+ import java .nio .file .Paths ;
13+ import java .nio .file .StandardOpenOption ;
14+ import java .util .Map ;
15+
16+ import com .amazonaws .services .lambda .runtime .Context ;
17+ import com .amazonaws .services .lambda .runtime .events .SQSEvent ;
18+ import com .amazonaws .services .lambda .runtime .events .models .s3 .S3EventNotification ;
19+ import com .amazonaws .services .s3 .AmazonS3 ;
20+ import com .amazonaws .services .s3 .model .S3Object ;
21+ import com .amazonaws .xray .AWSXRay ;
22+ import com .fasterxml .jackson .core .JsonProcessingException ;
23+ import com .fasterxml .jackson .databind .ObjectMapper ;
24+ import org .apache .logging .log4j .Level ;
25+ import org .apache .logging .log4j .ThreadContext ;
26+ import org .junit .jupiter .api .AfterEach ;
27+ import org .junit .jupiter .api .BeforeEach ;
28+ import org .junit .jupiter .api .Test ;
29+ import org .mockito .Mock ;
30+ import software .amazon .lambda .powertools .core .internal .LambdaHandlerProcessor ;
31+ import software .amazon .lambda .powertools .logging .internal .LambdaLoggingAspect ;
32+ import software .amazon .lambda .powertools .sqs .internal .SqsLargeMessageAspect ;
33+ import software .amazon .lambda .powertools .testsuite .handler .LoggingOrderMessageHandler ;
34+ import software .amazon .lambda .powertools .testsuite .handler .TracingLoggingStreamMessageHandler ;
35+
36+ import static java .util .Collections .emptyMap ;
37+ import static java .util .Collections .singletonList ;
38+ import static org .apache .commons .lang3 .reflect .FieldUtils .writeStaticField ;
39+ import static org .assertj .core .api .Assertions .assertThat ;
40+ import static org .assertj .core .api .Assertions .fail ;
41+ import static org .mockito .Mockito .when ;
42+ import static org .mockito .MockitoAnnotations .openMocks ;
43+
44+ public class LoggingOrderTest {
45+
46+ private static final String BUCKET_NAME = "ms-extended-sqs-client" ;
47+ private static final String BUCKET_KEY = "c71eb2ae-37e0-4265-8909-32f4153faddf" ;
48+
49+ @ Mock
50+ private Context context ;
51+
52+ @ Mock
53+ private AmazonS3 amazonS3 ;
54+
55+ @ BeforeEach
56+ void setUp () throws IllegalAccessException , IOException , NoSuchMethodException , InvocationTargetException {
57+ openMocks (this );
58+ writeStaticField (SqsLargeMessageAspect .class , "amazonS3" , amazonS3 , true );
59+ ThreadContext .clearAll ();
60+ writeStaticField (LambdaHandlerProcessor .class , "IS_COLD_START" , null , true );
61+ setupContext ();
62+ //Make sure file is cleaned up before running full stack logging regression
63+ FileChannel .open (Paths .get ("target/logfile.json" ), StandardOpenOption .WRITE ).truncate (0 ).close ();
64+ resetLogLevel (Level .INFO );
65+ AWSXRay .beginSegment (LoggingOrderTest .class .getName ());
66+ }
67+
68+ @ AfterEach
69+ void tearDown () {
70+ AWSXRay .endSegment ();
71+ }
72+
73+ /**
74+ * The SQSEvent payload will be altered by the @SqsLargeMessage annotation. Logging of the event should happen
75+ * after the event has been altered
76+ */
77+ @ Test
78+ public void testThatLoggingAnnotationActsLast () throws IOException {
79+ S3Object s3Response = new S3Object ();
80+ s3Response .setObjectContent (new ByteArrayInputStream ("A big message" .getBytes ()));
81+
82+ when (amazonS3 .getObject (BUCKET_NAME , BUCKET_KEY )).thenReturn (s3Response );
83+ SQSEvent sqsEvent = messageWithBody ("[\" software.amazon.payloadoffloading.PayloadS3Pointer\" ,{\" s3BucketName\" :\" " + BUCKET_NAME + "\" ,\" s3Key\" :\" " + BUCKET_KEY + "\" }]" );
84+
85+ LoggingOrderMessageHandler requestHandler = new LoggingOrderMessageHandler ();
86+ requestHandler .handleRequest (sqsEvent , context );
87+
88+ assertThat (Files .lines (Paths .get ("target/logfile.json" )))
89+ .hasSize (2 )
90+ .satisfies (line -> {
91+ Map <String , Object > actual = parseToMap (line .get (0 ));
92+
93+ String message = actual .get ("message" ).toString ();
94+
95+ assertThat (message )
96+ .contains ("A big message" );
97+ });
98+ }
99+
100+ @ Test
101+ public void testLoggingAnnotationActsAfterTracingForStreamingHandler () throws IOException {
102+
103+ ByteArrayOutputStream output = new ByteArrayOutputStream ();
104+ S3EventNotification s3EventNotification = s3EventNotification ();
105+
106+ TracingLoggingStreamMessageHandler handler = new TracingLoggingStreamMessageHandler ();
107+ handler .handleRequest (new ByteArrayInputStream (new ObjectMapper ().writeValueAsBytes (s3EventNotification )), output , context );
108+
109+ assertThat (new String (output .toByteArray (), StandardCharsets .UTF_8 ))
110+ .isNotEmpty ();
111+ }
112+
113+ private void setupContext () {
114+ when (context .getFunctionName ()).thenReturn ("testFunction" );
115+ when (context .getInvokedFunctionArn ()).thenReturn ("testArn" );
116+ when (context .getFunctionVersion ()).thenReturn ("1" );
117+ when (context .getMemoryLimitInMB ()).thenReturn (10 );
118+ when (context .getAwsRequestId ()).thenReturn ("RequestId" );
119+ }
120+
121+ private void resetLogLevel (Level level ) throws NoSuchMethodException , IllegalAccessException , InvocationTargetException {
122+ Method resetLogLevels = LambdaLoggingAspect .class .getDeclaredMethod ("resetLogLevels" , Level .class );
123+ resetLogLevels .setAccessible (true );
124+ resetLogLevels .invoke (null , level );
125+ writeStaticField (LambdaLoggingAspect .class , "LEVEL_AT_INITIALISATION" , level , true );
126+ }
127+
128+ private Map <String , Object > parseToMap (String stringAsJson ) {
129+ try {
130+ return new ObjectMapper ().readValue (stringAsJson , Map .class );
131+ } catch (JsonProcessingException e ) {
132+ fail ("Failed parsing logger line " + stringAsJson );
133+ return emptyMap ();
134+ }
135+ }
136+
137+ private S3EventNotification s3EventNotification () {
138+ S3EventNotification .S3EventNotificationRecord record = new S3EventNotification .S3EventNotificationRecord ("us-west-2" ,
139+ "ObjectCreated:Put" ,
140+ "aws:s3" ,
141+ null ,
142+ "2.1" ,
143+ new S3EventNotification .RequestParametersEntity ("127.0.0.1" ),
144+ new S3EventNotification .ResponseElementsEntity ("C3D13FE58DE4C810" , "FMyUVURIY8/IgAtTv8xRjskZQpcIZ9KG4V5Wp6S7S/JRWeUWerMUE5JgHvANOjpD" ),
145+ new S3EventNotification .S3Entity ("testConfigRule" ,
146+ new S3EventNotification .S3BucketEntity ("mybucket" ,
147+ new S3EventNotification .UserIdentityEntity ("A3NL1KOZZKExample" ),
148+ "arn:aws:s3:::mybucket" ),
149+ new S3EventNotification .S3ObjectEntity ("HappyFace.jpg" ,
150+ 1024L ,
151+ "d41d8cd98f00b204e9800998ecf8427e" ,
152+ "096fKKXTRTtl3on89fVO.nfljtsv6qko" ,
153+ "0055AED6DCD90281E5" ),
154+ "1.0" ),
155+ new S3EventNotification .UserIdentityEntity ("AIDAJDPLRKLG7UEXAMPLE" )
156+ );
157+
158+ return new S3EventNotification (singletonList (record ));
159+ }
160+
161+ private SQSEvent messageWithBody (String messageBody ) {
162+ SQSEvent .SQSMessage sqsMessage = new SQSEvent .SQSMessage ();
163+ sqsMessage .setBody (messageBody );
164+ SQSEvent sqsEvent = new SQSEvent ();
165+ sqsEvent .setRecords (singletonList (sqsMessage ));
166+ return sqsEvent ;
167+ }
168+ }
0 commit comments