|
17 | 17 |
|
18 | 18 | import static org.assertj.core.api.Assertions.assertThat; |
19 | 19 |
|
20 | | -import java.util.List; |
21 | | -import java.util.concurrent.atomic.AtomicReference; |
22 | 20 | import org.junit.jupiter.api.Test; |
23 | | -import org.slf4j.MDC; |
24 | 21 | import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider; |
25 | 22 | import software.amazon.awssdk.awscore.interceptor.TraceIdExecutionInterceptor; |
26 | | -import software.amazon.awssdk.core.interceptor.Context; |
27 | | -import software.amazon.awssdk.core.interceptor.ExecutionAttributes; |
28 | | -import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; |
29 | 23 | import software.amazon.awssdk.http.AbortableInputStream; |
30 | 24 | import software.amazon.awssdk.http.HttpExecuteResponse; |
31 | | -import software.amazon.awssdk.http.SdkHttpRequest; |
32 | 25 | import software.amazon.awssdk.http.SdkHttpResponse; |
33 | 26 | import software.amazon.awssdk.regions.Region; |
34 | | -import software.amazon.awssdk.services.protocolrestjson.ProtocolRestJsonAsyncClient; |
35 | 27 | import software.amazon.awssdk.services.protocolrestjson.ProtocolRestJsonClient; |
36 | 28 | import software.amazon.awssdk.testutils.EnvironmentVariableHelper; |
37 | | -import software.amazon.awssdk.testutils.service.http.MockAsyncHttpClient; |
38 | 29 | import software.amazon.awssdk.testutils.service.http.MockSyncHttpClient; |
39 | 30 | import software.amazon.awssdk.utils.StringInputStream; |
40 | 31 |
|
@@ -65,182 +56,4 @@ public void traceIdInterceptorIsEnabled() { |
65 | 56 | } |
66 | 57 | }); |
67 | 58 | } |
68 | | - |
69 | | - @Test |
70 | | - public void traceIdInterceptorPreservesTraceIdAcrossRetries() { |
71 | | - EnvironmentVariableHelper.run(env -> { |
72 | | - env.set("AWS_LAMBDA_FUNCTION_NAME", "foo"); |
73 | | - MDC.put("AWS_LAMBDA_X_TRACE_ID", "mdc-trace-123"); |
74 | | - |
75 | | - try (MockAsyncHttpClient mockHttpClient = new MockAsyncHttpClient(); |
76 | | - ProtocolRestJsonAsyncClient client = ProtocolRestJsonAsyncClient.builder() |
77 | | - .region(Region.US_WEST_2) |
78 | | - .credentialsProvider(AnonymousCredentialsProvider.create()) |
79 | | - .httpClient(mockHttpClient) |
80 | | - .build()) { |
81 | | - |
82 | | - mockHttpClient.stubResponses( |
83 | | - HttpExecuteResponse.builder() |
84 | | - .response(SdkHttpResponse.builder().statusCode(500).build()) |
85 | | - .responseBody(AbortableInputStream.create(new StringInputStream("{}"))) |
86 | | - .build(), |
87 | | - HttpExecuteResponse.builder() |
88 | | - .response(SdkHttpResponse.builder().statusCode(500).build()) |
89 | | - .responseBody(AbortableInputStream.create(new StringInputStream("{}"))) |
90 | | - .build(), |
91 | | - HttpExecuteResponse.builder().response(SdkHttpResponse.builder().statusCode(200).build()) |
92 | | - .responseBody(AbortableInputStream.create(new StringInputStream("{}"))) |
93 | | - .build()); |
94 | | - |
95 | | - client.allTypes().join(); |
96 | | - |
97 | | - List<SdkHttpRequest> requests = mockHttpClient.getRequests(); |
98 | | - assertThat(requests).hasSize(3); |
99 | | - |
100 | | - assertThat(requests.get(0).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123"); |
101 | | - assertThat(requests.get(1).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123"); |
102 | | - assertThat(requests.get(2).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123"); |
103 | | - |
104 | | - } finally { |
105 | | - MDC.clear(); |
106 | | - } |
107 | | - }); |
108 | | - } |
109 | | - |
110 | | - @Test |
111 | | - public void traceIdInterceptorPreservesTraceIdAcrossChainedFutures() { |
112 | | - EnvironmentVariableHelper.run(env -> { |
113 | | - env.set("AWS_LAMBDA_FUNCTION_NAME", "foo"); |
114 | | - MDC.put("AWS_LAMBDA_X_TRACE_ID", "mdc-trace-123"); |
115 | | - |
116 | | - try (MockAsyncHttpClient mockHttpClient = new MockAsyncHttpClient(); |
117 | | - ProtocolRestJsonAsyncClient client = ProtocolRestJsonAsyncClient.builder() |
118 | | - .region(Region.US_WEST_2) |
119 | | - .credentialsProvider(AnonymousCredentialsProvider.create()) |
120 | | - .httpClient(mockHttpClient) |
121 | | - .build()) { |
122 | | - |
123 | | - mockHttpClient.stubResponses( |
124 | | - HttpExecuteResponse.builder() |
125 | | - .response(SdkHttpResponse.builder().statusCode(200).build()) |
126 | | - .responseBody(AbortableInputStream.create(new StringInputStream("{}"))) |
127 | | - .build(), |
128 | | - HttpExecuteResponse.builder() |
129 | | - .response(SdkHttpResponse.builder().statusCode(200).build()) |
130 | | - .responseBody(AbortableInputStream.create(new StringInputStream("{}"))) |
131 | | - .build() |
132 | | - ); |
133 | | - |
134 | | - client.allTypes() |
135 | | - .thenRun(() -> { |
136 | | - client.allTypes().join(); |
137 | | - }) |
138 | | - .join(); |
139 | | - |
140 | | - List<SdkHttpRequest> requests = mockHttpClient.getRequests(); |
141 | | - |
142 | | - assertThat(requests).hasSize(2); |
143 | | - |
144 | | - assertThat(requests.get(0).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123"); |
145 | | - assertThat(requests.get(1).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123"); |
146 | | - |
147 | | - } finally { |
148 | | - MDC.clear(); |
149 | | - } |
150 | | - }); |
151 | | - } |
152 | | - |
153 | | - @Test |
154 | | - public void traceIdInterceptorPreservesTraceIdAcrossExceptionallyCompletedFutures() { |
155 | | - EnvironmentVariableHelper.run(env -> { |
156 | | - env.set("AWS_LAMBDA_FUNCTION_NAME", "foo"); |
157 | | - MDC.put("AWS_LAMBDA_X_TRACE_ID", "mdc-trace-123"); |
158 | | - |
159 | | - try (MockAsyncHttpClient mockHttpClient = new MockAsyncHttpClient(); |
160 | | - ProtocolRestJsonAsyncClient client = ProtocolRestJsonAsyncClient.builder() |
161 | | - .region(Region.US_WEST_2) |
162 | | - .credentialsProvider(AnonymousCredentialsProvider.create()) |
163 | | - .httpClient(mockHttpClient) |
164 | | - .build()) { |
165 | | - |
166 | | - mockHttpClient.stubResponses( |
167 | | - HttpExecuteResponse.builder() |
168 | | - .response(SdkHttpResponse.builder().statusCode(400).build()) |
169 | | - .responseBody(AbortableInputStream.create(new StringInputStream("{}"))) |
170 | | - .build(), |
171 | | - HttpExecuteResponse.builder() |
172 | | - .response(SdkHttpResponse.builder().statusCode(200).build()) |
173 | | - .responseBody(AbortableInputStream.create(new StringInputStream("{}"))) |
174 | | - .build() |
175 | | - ); |
176 | | - |
177 | | - client.allTypes() |
178 | | - .exceptionally(throwable -> { |
179 | | - client.allTypes().join(); |
180 | | - return null; |
181 | | - }).join(); |
182 | | - |
183 | | - List<SdkHttpRequest> requests = mockHttpClient.getRequests(); |
184 | | - |
185 | | - assertThat(requests).hasSize(2); |
186 | | - |
187 | | - assertThat(requests.get(0).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123"); |
188 | | - assertThat(requests.get(1).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123"); |
189 | | - |
190 | | - } finally { |
191 | | - MDC.clear(); |
192 | | - } |
193 | | - }); |
194 | | - } |
195 | | - |
196 | | - @Test |
197 | | - public void traceIdInterceptorPreservesTraceIdAcrossExceptionallyCompletedFuturesThrownInPreExecution() { |
198 | | - EnvironmentVariableHelper.run(env -> { |
199 | | - env.set("AWS_LAMBDA_FUNCTION_NAME", "foo"); |
200 | | - MDC.put("AWS_LAMBDA_X_TRACE_ID", "mdc-trace-123"); |
201 | | - |
202 | | - ExecutionInterceptor throwingInterceptor = new ExecutionInterceptor() { |
203 | | - private boolean hasThrown = false; |
204 | | - |
205 | | - @Override |
206 | | - public void beforeMarshalling(Context.BeforeMarshalling context, ExecutionAttributes executionAttributes) { |
207 | | - if (!hasThrown) { |
208 | | - hasThrown = true; |
209 | | - throw new RuntimeException("failing in pre execution"); |
210 | | - } |
211 | | - } |
212 | | - }; |
213 | | - |
214 | | - try (MockAsyncHttpClient mockHttpClient = new MockAsyncHttpClient(); |
215 | | - ProtocolRestJsonAsyncClient client = ProtocolRestJsonAsyncClient.builder() |
216 | | - .region(Region.US_WEST_2) |
217 | | - .credentialsProvider(AnonymousCredentialsProvider.create()) |
218 | | - .overrideConfiguration(o -> o.addExecutionInterceptor(throwingInterceptor)) |
219 | | - .httpClient(mockHttpClient) |
220 | | - .build()) { |
221 | | - |
222 | | - mockHttpClient.stubResponses( |
223 | | - HttpExecuteResponse.builder() |
224 | | - .response(SdkHttpResponse.builder().statusCode(200).build()) |
225 | | - .responseBody(AbortableInputStream.create(new StringInputStream("{}"))) |
226 | | - .build() |
227 | | - ); |
228 | | - |
229 | | - client.allTypes() |
230 | | - .exceptionally(throwable -> { |
231 | | - client.allTypes().join(); |
232 | | - return null; |
233 | | - }).join(); |
234 | | - |
235 | | - List<SdkHttpRequest> requests = mockHttpClient.getRequests(); |
236 | | - |
237 | | - assertThat(requests).hasSize(1); |
238 | | - assertThat(requests.get(0).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123"); |
239 | | - |
240 | | - } finally { |
241 | | - MDC.clear(); |
242 | | - } |
243 | | - }); |
244 | | - } |
245 | 59 | } |
246 | | - |
0 commit comments