Java Code Examples for software.amazon.awssdk.core.interceptor.Context#AfterTransmission

@Test
public void sync_streamingOutput_success_allInterceptorMethodsCalled() throws IOException {
    // Given
    ExecutionInterceptor interceptor = mock(NoOpInterceptor.class, CALLS_REAL_METHODS);
    ProtocolRestJsonClient client = client(interceptor);
    StreamingOutputOperationRequest request = StreamingOutputOperationRequest.builder().build();
    stubFor(post(urlPathEqualTo(STREAMING_OUTPUT_PATH)).willReturn(aResponse().withStatus(200).withBody("\0")));

    // When
    client.streamingOutputOperation(request, (r, i) -> {
        assertThat(i.read()).isEqualTo(0);
        // TODO: We have to return "r" here. We should verify other response types are cool once we switch this off of
        // being the unmarshaller
        return r;
    });

    // Expect
    Context.AfterTransmission afterTransmissionArg = captureAfterTransmissionArg(interceptor);
    // TODO: When we don't always close the input stream, make sure we can read the service's '0' response.
    assertThat(afterTransmissionArg.responseBody() != null);
}
 

@Test
public void async_streamingOutput_success_allInterceptorMethodsCalled()
        throws IOException, InterruptedException, ExecutionException, TimeoutException {
    // Given
    ExecutionInterceptor interceptor = mock(NoOpInterceptor.class, CALLS_REAL_METHODS);
    ProtocolRestJsonAsyncClient client = asyncClient(interceptor);
    StreamingOutputOperationRequest request = StreamingOutputOperationRequest.builder().build();
    stubFor(post(urlPathEqualTo(STREAMING_OUTPUT_PATH)).willReturn(aResponse().withStatus(200).withBody("\0")));

    // When
    client.streamingOutputOperation(request, new NoOpAsyncResponseTransformer()).get(10, TimeUnit.SECONDS);

    // Expect
    Context.AfterTransmission afterTransmissionArg = captureAfterTransmissionArgAsync(interceptor);
    assertThat(afterTransmissionArg.responseBody()).isNotNull();
}
 

private Context.AfterTransmission captureAfterTransmissionArg(ExecutionInterceptor interceptor) {
    ArgumentCaptor<Context.AfterTransmission> afterTransmissionArg = ArgumentCaptor.forClass(Context.AfterTransmission.class);

    InOrder inOrder = Mockito.inOrder(interceptor);
    inOrder.verify(interceptor).beforeExecution(any(), any());
    inOrder.verify(interceptor).modifyRequest(any(), any());
    inOrder.verify(interceptor).beforeMarshalling(any(), any());
    inOrder.verify(interceptor).afterMarshalling(any(), any());
    inOrder.verify(interceptor).modifyAsyncHttpContent(any(), any());
    inOrder.verify(interceptor).modifyHttpContent(any(), any());
    inOrder.verify(interceptor).modifyHttpRequest(any(), any());
    inOrder.verify(interceptor).beforeTransmission(any(), any());
    inOrder.verify(interceptor).afterTransmission(afterTransmissionArg.capture(), any());
    inOrder.verify(interceptor).modifyHttpResponse(any(), any());
    inOrder.verify(interceptor).modifyHttpResponseContent(any(), any());
    inOrder.verify(interceptor).beforeUnmarshalling(any(), any());
    inOrder.verify(interceptor).afterUnmarshalling(any(), any());
    inOrder.verify(interceptor).modifyResponse(any(), any());
    inOrder.verify(interceptor).afterExecution(any(), any());
    verifyNoMoreInteractions(interceptor);
    return afterTransmissionArg.getValue();
}
 

private Context.AfterTransmission captureAfterTransmissionArgAsync(ExecutionInterceptor interceptor) {
    ArgumentCaptor<Context.AfterTransmission> afterTransmissionArg = ArgumentCaptor.forClass(Context.AfterTransmission.class);

    InOrder inOrder = Mockito.inOrder(interceptor);
    inOrder.verify(interceptor).beforeExecution(any(), any());
    inOrder.verify(interceptor).modifyRequest(any(), any());
    inOrder.verify(interceptor).beforeMarshalling(any(), any());
    inOrder.verify(interceptor).afterMarshalling(any(), any());
    inOrder.verify(interceptor).modifyAsyncHttpContent(any(), any());
    inOrder.verify(interceptor).modifyHttpContent(any(), any());
    inOrder.verify(interceptor).modifyHttpRequest(any(), any());
    inOrder.verify(interceptor).beforeTransmission(any(), any());
    inOrder.verify(interceptor).afterTransmission(afterTransmissionArg.capture(), any());
    inOrder.verify(interceptor).modifyHttpResponse(any(), any());
    inOrder.verify(interceptor).modifyHttpResponseContent(any(), any());
    inOrder.verify(interceptor).beforeUnmarshalling(any(), any());
    inOrder.verify(interceptor).afterUnmarshalling(any(), any());
    inOrder.verify(interceptor).modifyResponse(any(), any());
    inOrder.verify(interceptor).modifyAsyncHttpResponseContent(any(), any());
    inOrder.verify(interceptor).afterExecution(any(), any());
    verifyNoMoreInteractions(interceptor);
    return afterTransmissionArg.getValue();
}
 

private void validateArgs(Context.AfterTransmission context,
                          String expectedStringMemberValue, String expectedIntegerHeaderValue,
                          String expectedResponseIntegerHeaderValue) {
    validateArgs(context, expectedStringMemberValue, expectedIntegerHeaderValue);
    assertThat(context.httpResponse().firstMatchingHeader("x-amz-integer"))
            .isEqualTo(Optional.ofNullable(expectedResponseIntegerHeaderValue));
}
 

/**
 * After sending an http request. Will be called multiple times in the case of retries.
 */
@Override public void afterTransmission(Context.AfterTransmission context,
    ExecutionAttributes executionAttributes) {
  Span span = executionAttributes.getAttribute(SPAN);
  if (span == null) {
    // An evil interceptor deleted our attribute.
    return;
  }
  span.annotate("wr");
}
 

@Override
public void afterTransmission(Context.AfterTransmission context, ExecutionAttributes attrs) {
  SdkHttpResponse response = context.httpResponse();
  IpcLogEntry logEntry = attrs.getAttribute(LOG_ENTRY)
      .markEnd()
      .withHttpStatus(response.statusCode());
  attrs.putAttribute(STATUS_IS_SET, true);

  response.headers().forEach((k, vs) -> vs.forEach(v -> logEntry.addResponseHeader(k, v)));
}
 

@Override
public void afterTransmission(Context.AfterTransmission context, ExecutionAttributes executionAttributes) {
    SdkHttpFullRequest request = (SdkHttpFullRequest) context.httpRequest();
    assertThat(context.httpRequest().headers().get("x-amz-content-sha256").get(0)).doesNotContain("UNSIGNED-PAYLOAD");
}
 

private void expectAllMethodsCalled(ExecutionInterceptor interceptor,
                                    SdkRequest inputRequest,
                                    Exception expectedException,
                                    boolean isAsync) {
    ArgumentCaptor<ExecutionAttributes> attributes = ArgumentCaptor.forClass(ExecutionAttributes.class);

    ArgumentCaptor<Context.BeforeExecution> beforeExecutionArg = ArgumentCaptor.forClass(Context.BeforeExecution.class);
    ArgumentCaptor<Context.BeforeMarshalling> modifyRequestArg = ArgumentCaptor.forClass(Context.BeforeMarshalling.class);
    ArgumentCaptor<Context.BeforeMarshalling> beforeMarshallingArg = ArgumentCaptor.forClass(Context.BeforeMarshalling.class);
    ArgumentCaptor<Context.AfterMarshalling> afterMarshallingArg = ArgumentCaptor.forClass(Context.AfterMarshalling.class);
    ArgumentCaptor<Context.BeforeTransmission> modifyHttpRequestArg = ArgumentCaptor.forClass(Context.BeforeTransmission.class);
    ArgumentCaptor<Context.BeforeTransmission> modifyHttpContentArg = ArgumentCaptor.forClass(Context.BeforeTransmission.class);
    ArgumentCaptor<Context.BeforeTransmission> modifyHttpContentAsyncArg = ArgumentCaptor.forClass(Context.BeforeTransmission.class);
    ArgumentCaptor<Context.BeforeTransmission> beforeTransmissionArg = ArgumentCaptor.forClass(Context.BeforeTransmission.class);
    ArgumentCaptor<Context.AfterTransmission> afterTransmissionArg = ArgumentCaptor.forClass(Context.AfterTransmission.class);
    ArgumentCaptor<Context.BeforeUnmarshalling> modifyHttpResponseArg = ArgumentCaptor.forClass(Context.BeforeUnmarshalling.class);
    ArgumentCaptor<Context.BeforeUnmarshalling> modifyHttpResponseContentArg = ArgumentCaptor.forClass(Context.BeforeUnmarshalling.class);
    ArgumentCaptor<Context.BeforeUnmarshalling> beforeUnmarshallingArg = ArgumentCaptor.forClass(Context.BeforeUnmarshalling.class);
    ArgumentCaptor<Context.BeforeUnmarshalling> modifyAsyncHttpResponseContent = ArgumentCaptor.forClass(Context.BeforeUnmarshalling.class);
    ArgumentCaptor<Context.AfterUnmarshalling> afterUnmarshallingArg = ArgumentCaptor.forClass(Context.AfterUnmarshalling.class);
    ArgumentCaptor<Context.AfterExecution> modifyResponseArg = ArgumentCaptor.forClass(Context.AfterExecution.class);
    ArgumentCaptor<Context.AfterExecution> afterExecutionArg = ArgumentCaptor.forClass(Context.AfterExecution.class);

    // Verify methods are called in the right order
    InOrder inOrder = Mockito.inOrder(interceptor);
    inOrder.verify(interceptor).beforeExecution(beforeExecutionArg.capture(), attributes.capture());
    inOrder.verify(interceptor).modifyRequest(modifyRequestArg.capture(), attributes.capture());
    inOrder.verify(interceptor).beforeMarshalling(beforeMarshallingArg.capture(), attributes.capture());
    inOrder.verify(interceptor).afterMarshalling(afterMarshallingArg.capture(), attributes.capture());
    inOrder.verify(interceptor).modifyAsyncHttpContent(modifyHttpContentAsyncArg.capture(), attributes.capture());
    inOrder.verify(interceptor).modifyHttpContent(modifyHttpContentArg.capture(), attributes.capture());
    inOrder.verify(interceptor).modifyHttpRequest(modifyHttpRequestArg.capture(), attributes.capture());
    inOrder.verify(interceptor).beforeTransmission(beforeTransmissionArg.capture(), attributes.capture());
    inOrder.verify(interceptor).afterTransmission(afterTransmissionArg.capture(), attributes.capture());
    inOrder.verify(interceptor).modifyHttpResponse(modifyHttpResponseArg.capture(), attributes.capture());
    inOrder.verify(interceptor).modifyHttpResponseContent(modifyHttpResponseContentArg.capture(), attributes.capture());
    inOrder.verify(interceptor).beforeUnmarshalling(beforeUnmarshallingArg.capture(), attributes.capture());
    if (isAsync) {
        inOrder.verify(interceptor).modifyAsyncHttpResponseContent(modifyAsyncHttpResponseContent.capture(), attributes.capture());
    }
    inOrder.verify(interceptor).afterUnmarshalling(afterUnmarshallingArg.capture(), attributes.capture());
    inOrder.verify(interceptor).modifyResponse(modifyResponseArg.capture(), attributes.capture());
    inOrder.verify(interceptor).afterExecution(afterExecutionArg.capture(), attributes.capture());
    if (expectedException != null) {
        ArgumentCaptor<Context.FailedExecution> failedExecutionArg = ArgumentCaptor.forClass(Context.FailedExecution.class);
        inOrder.verify(interceptor).modifyException(failedExecutionArg.capture(), attributes.capture());
        inOrder.verify(interceptor).onExecutionFailure(failedExecutionArg.capture(), attributes.capture());
        verifyFailedExecutionMethodCalled(failedExecutionArg, true);
        assertThat(failedExecutionArg.getValue().exception()).isEqualTo(expectedException);
    }
    verifyNoMoreInteractions(interceptor);

    // Verify beforeExecution gets untouched request
    assertThat(beforeExecutionArg.getValue().request()).isSameAs(inputRequest);

    // Verify methods were given correct parameters
    validateArgs(beforeExecutionArg.getValue(), null);
    validateArgs(modifyRequestArg.getValue(), null);
    validateArgs(beforeMarshallingArg.getValue(), "1");
    validateArgs(afterMarshallingArg.getValue(), "1", null);
    validateArgs(modifyHttpRequestArg.getValue(), "1", null);
    validateArgs(beforeTransmissionArg.getValue(), "1", "2");
    validateArgs(afterTransmissionArg.getValue(), "1", "2", null);
    validateArgs(modifyHttpResponseArg.getValue(), "1", "2", null);
    validateArgs(beforeUnmarshallingArg.getValue(), "1", "2", "3");
    validateArgs(afterUnmarshallingArg.getValue(), "1", "2", "3", null);
    validateArgs(modifyResponseArg.getValue(), "1", "2", "3", null);
    validateArgs(afterExecutionArg.getValue(), "1", "2", "3", "4");

    // Verify same execution attributes were used for all method calls
    assertThat(attributes.getAllValues()).containsOnly(attributes.getAllValues().get(0));
}
 

@Override
public void afterTransmission(Context.AfterTransmission context, ExecutionAttributes executionAttributes) {
    wait(afterTransmissionWait);
}