Refactor xml response handler to handle SdkPojo

core/protocols/aws-xml-protocol/src/main/java/software/amazon/awssdk/protocols/xml/AwsXmlProtocolFactory.java

@@ -34,6 +34,7 @@
 import software.amazon.awssdk.core.internal.http.CombinedResponseHandler;
 import software.amazon.awssdk.core.metrics.CoreMetric;
 import software.amazon.awssdk.http.SdkHttpFullRequest;
+import software.amazon.awssdk.http.SdkHttpFullResponse;
 import software.amazon.awssdk.protocols.core.ExceptionMetadata;
 import software.amazon.awssdk.protocols.core.OperationInfo;
 import software.amazon.awssdk.protocols.core.OperationMetadataAttribute;
@@ -47,6 +48,7 @@
 import software.amazon.awssdk.protocols.xml.internal.unmarshall.AwsXmlResponseTransformer;
 import software.amazon.awssdk.protocols.xml.internal.unmarshall.AwsXmlUnmarshallingContext;
 import software.amazon.awssdk.protocols.xml.internal.unmarshall.XmlProtocolUnmarshaller;
+import software.amazon.awssdk.protocols.xml.internal.unmarshall.XmlResponseHandler;
 
 /**
  * Factory to generate the various protocol handlers and generators to be used for
@@ -104,10 +106,18 @@ public ProtocolMarshaller<SdkHttpFullRequest> createProtocolMarshaller(Operation
                                     .build();
     }
 
-    public <T extends AwsResponse> HttpResponseHandler<T> createResponseHandler(Supplier<SdkPojo> pojoSupplier,
-                                                                                XmlOperationMetadata staxOperationMetadata) {
-        return timeUnmarshalling(new AwsXmlResponseHandler<>(XML_PROTOCOL_UNMARSHALLER, r -> pojoSupplier.get(),
-                                                             staxOperationMetadata.isHasStreamingSuccessResponse()));
+    public <T extends SdkPojo> HttpResponseHandler<T> createResponseHandler(Supplier<SdkPojo> pojoSupplier,
+                                                                            XmlOperationMetadata staxOperationMetadata) {
+        return createResponseHandler(r -> pojoSupplier.get(), staxOperationMetadata);
+    }
+
+    public <T extends SdkPojo> HttpResponseHandler<T> createResponseHandler(Function<SdkHttpFullResponse, SdkPojo> pojoSupplier,
+                                                                            XmlOperationMetadata staxOperationMetadata) {
+        return timeUnmarshalling(
+            new AwsXmlResponseHandler<>(
+                new XmlResponseHandler<>(
+                    XML_PROTOCOL_UNMARSHALLER, pojoSupplier,
+                    staxOperationMetadata.isHasStreamingSuccessResponse())));
     }
 
     protected <T extends AwsResponse> Function<AwsXmlUnmarshallingContext, T> createResponseTransformer(

core/protocols/aws-xml-protocol/src/main/java/software/amazon/awssdk/protocols/xml/internal/unmarshall/AwsXmlResponseHandler.java

@@ -17,73 +17,42 @@
 
 import static software.amazon.awssdk.awscore.util.AwsHeader.AWS_REQUEST_ID;
 
-import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.function.Function;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.awscore.AwsResponse;
 import software.amazon.awssdk.awscore.AwsResponseMetadata;
 import software.amazon.awssdk.awscore.DefaultAwsResponseMetadata;
-import software.amazon.awssdk.core.SdkPojo;
-import software.amazon.awssdk.core.SdkStandardLogger;
 import software.amazon.awssdk.core.http.HttpResponseHandler;
 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
 import software.amazon.awssdk.http.SdkHttpFullResponse;
 import software.amazon.awssdk.http.SdkHttpResponse;
-import software.amazon.awssdk.utils.Logger;
 import software.amazon.awssdk.utils.http.SdkHttpUtils;
 
 /**
- * Response handler for REST-XML services (Cloudfront, Route53, and S3).
+ * Response handler that adds {@link AwsResponseMetadata} to the response.
  *
  * @param <T> Indicates the type being unmarshalled by this response handler.
  */
 @SdkInternalApi
-public final class AwsXmlResponseHandler<T extends AwsResponse> implements HttpResponseHandler<T> {
+public final class AwsXmlResponseHandler<T> implements HttpResponseHandler<T> {
 
-    private static final Logger log = Logger.loggerFor(AwsXmlResponseHandler.class);
+    private final HttpResponseHandler<T> delegate;
 
-    private final XmlProtocolUnmarshaller unmarshaller;
-    private final Function<SdkHttpFullResponse, SdkPojo> pojoSupplier;
-    private final boolean needsConnectionLeftOpen;
-
-    public AwsXmlResponseHandler(XmlProtocolUnmarshaller unmarshaller,
-                                 Function<SdkHttpFullResponse, SdkPojo> pojoSupplier,
-                                 boolean needsConnectionLeftOpen) {
-        this.unmarshaller = unmarshaller;
-        this.pojoSupplier = pojoSupplier;
-        this.needsConnectionLeftOpen = needsConnectionLeftOpen;
+    public AwsXmlResponseHandler(HttpResponseHandler<T> responseHandler) {
+        this.delegate = responseHandler;
     }
 
     @Override
     public T handle(SdkHttpFullResponse response, ExecutionAttributes executionAttributes) throws Exception {
-        try {
-            return unmarshallResponse(response);
-        } finally {
-            if (!needsConnectionLeftOpen) {
-                closeStream(response);
-            }
-        }
-    }
+        T result = delegate.handle(response, executionAttributes);
 
-    private void closeStream(SdkHttpFullResponse response) {
-        response.content().ifPresent(i -> {
-            try {
-                i.close();
-            } catch (IOException e) {
-                log.warn(() -> "Error closing HTTP content.", e);
-            }
-        });
-    }
+        if (result instanceof AwsResponse) {
+            AwsResponseMetadata responseMetadata = generateResponseMetadata(response);
+            return (T) ((AwsResponse) result).toBuilder().responseMetadata(responseMetadata).build();
+        }
 
-    @SuppressWarnings("unchecked")
-    private T unmarshallResponse(SdkHttpFullResponse response) throws Exception {
-        SdkStandardLogger.REQUEST_LOGGER.trace(() -> "Parsing service response XML.");
-        T result = unmarshaller.unmarshall(pojoSupplier.apply(response), response);
-        SdkStandardLogger.REQUEST_LOGGER.trace(() -> "Done parsing service response.");
-        AwsResponseMetadata responseMetadata = generateResponseMetadata(response);
-        return (T) result.toBuilder().responseMetadata(responseMetadata).build();
+        return result;
     }
 
     /**
@@ -101,6 +70,6 @@ private AwsResponseMetadata generateResponseMetadata(SdkHttpResponse response) {
 
     @Override
     public boolean needsConnectionLeftOpen() {
-        return needsConnectionLeftOpen;
+        return delegate.needsConnectionLeftOpen();
     }
 }

core/protocols/aws-xml-protocol/src/main/java/software/amazon/awssdk/protocols/xml/internal/unmarshall/XmlResponseHandler.java

@@ -0,0 +1,83 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.protocols.xml.internal.unmarshall;
+
+import java.io.IOException;
+import java.util.function.Function;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.core.SdkPojo;
+import software.amazon.awssdk.core.SdkStandardLogger;
+import software.amazon.awssdk.core.http.HttpResponseHandler;
+import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
+import software.amazon.awssdk.http.SdkHttpFullResponse;
+import software.amazon.awssdk.utils.Logger;
+
+/**
+ * Response handler for REST-XML services (Cloudfront, Route53, and S3).
+ *
+ * @param <T> Indicates the type being unmarshalled by this response handler.
+ */
+@SdkInternalApi
+public final class XmlResponseHandler<T extends SdkPojo> implements HttpResponseHandler<T> {
+
+    private static final Logger log = Logger.loggerFor(XmlResponseHandler.class);
+
+    private final XmlProtocolUnmarshaller unmarshaller;
+    private final Function<SdkHttpFullResponse, SdkPojo> pojoSupplier;
+    private final boolean needsConnectionLeftOpen;
+
+    public XmlResponseHandler(XmlProtocolUnmarshaller unmarshaller,
+                              Function<SdkHttpFullResponse, SdkPojo> pojoSupplier,
+                              boolean needsConnectionLeftOpen) {
+        this.unmarshaller = unmarshaller;
+        this.pojoSupplier = pojoSupplier;
+        this.needsConnectionLeftOpen = needsConnectionLeftOpen;
+    }
+
+    @Override
+    public T handle(SdkHttpFullResponse response, ExecutionAttributes executionAttributes) throws Exception {
+        try {
+            return unmarshallResponse(response);
+        } finally {
+            if (!needsConnectionLeftOpen) {
+                closeStream(response);
+            }
+        }
+    }
+
+    private void closeStream(SdkHttpFullResponse response) {
+        response.content().ifPresent(i -> {
+            try {
+                i.close();
+            } catch (IOException e) {
+                log.warn(() -> "Error closing HTTP content.", e);
+            }
+        });
+    }
+
+    @SuppressWarnings("unchecked")
+    private T unmarshallResponse(SdkHttpFullResponse response) throws Exception {
+        SdkStandardLogger.REQUEST_LOGGER.trace(() -> "Parsing service response XML.");
+        T result = unmarshaller.unmarshall(pojoSupplier.apply(response), response);
+        SdkStandardLogger.REQUEST_LOGGER.trace(() -> "Done parsing service response.");
+        return result;
+    }
+
+    @Override
+    public boolean needsConnectionLeftOpen() {
+        return needsConnectionLeftOpen;
+    }
+}

core/protocols/aws-xml-protocol/src/test/java/software/amazon/awssdk/protocols/xml/internal/unmarshall/AwsXmlResponseHandlerTest.java

@@ -0,0 +1,129 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.protocols.xml.internal.unmarshall;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static software.amazon.awssdk.awscore.util.AwsHeader.AWS_REQUEST_ID;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import org.junit.Test;
+import org.mockito.Mockito;
+import software.amazon.awssdk.awscore.AwsResponse;
+import software.amazon.awssdk.core.SdkField;
+import software.amazon.awssdk.core.SdkPojo;
+import software.amazon.awssdk.core.http.HttpResponseHandler;
+import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
+import software.amazon.awssdk.http.AbortableInputStream;
+import software.amazon.awssdk.http.SdkHttpFullResponse;
+import software.amazon.awssdk.utils.ImmutableMap;
+
+public class AwsXmlResponseHandlerTest {
+
+    @Test
+    public void handleResponse_awsResponse_shouldAddResponseMetadata() throws Exception {
+        HttpResponseHandler<FakeResponse> delegate = Mockito.mock(HttpResponseHandler.class);
+        AwsXmlResponseHandler<FakeResponse> responseHandler = new AwsXmlResponseHandler<>(delegate);
+
+        SdkHttpFullResponse response = new TestSdkHttpFullResponse();
+        ExecutionAttributes executionAttributes = new ExecutionAttributes();
+        FakeResponse fakeResponse = FakeResponse.builder().build();
+
+        Mockito.when(delegate.handle(response, executionAttributes)).thenReturn(fakeResponse);
+
+        assertThat(responseHandler.handle(response, executionAttributes)
+            .responseMetadata().requestId()).isEqualTo("1234");
+    }
+
+    @Test
+    public void handleResponse_nonAwsResponse_shouldReturnDirectly() throws Exception {
+        HttpResponseHandler<SdkPojo> delegate = Mockito.mock(HttpResponseHandler.class);
+        AwsXmlResponseHandler<SdkPojo> responseHandler = new AwsXmlResponseHandler<>(delegate);
+
+        SdkHttpFullResponse response = new TestSdkHttpFullResponse();
+        ExecutionAttributes executionAttributes = new ExecutionAttributes();
+        FakeSdkPojo fakeResponse = new FakeSdkPojo();
+
+        Mockito.when(delegate.handle(response, executionAttributes)).thenReturn(fakeResponse);
+
+        assertThat(responseHandler.handle(response, executionAttributes)).isEqualTo(fakeResponse);
+    }
+
+    private static final class FakeSdkPojo implements SdkPojo {
+
+        @Override
+        public List<SdkField<?>> sdkFields() {
+            return Collections.emptyList();
+        }
+    }
+
+    private static final class FakeResponse extends AwsResponse {
+        private FakeResponse(Builder builder) {
+            super(builder);
+        }
+
+        @Override
+        public Builder toBuilder() {
+            return new Builder();
+        }
+
+        @Override
+        public List<SdkField<?>> sdkFields() {
+            return Collections.emptyList();
+        }
+
+        public static Builder builder() {
+            return new Builder();
+        }
+
+        private static final class Builder extends AwsResponse.BuilderImpl {
+
+            @Override
+            public FakeResponse build() {
+                return new FakeResponse(this);
+            }
+        }
+    }
+
+    private static class TestSdkHttpFullResponse implements SdkHttpFullResponse {
+        @Override
+        public Builder toBuilder() {
+            return null;
+        }
+
+        @Override
+        public Optional<AbortableInputStream> content() {
+            return Optional.empty();
+        }
+
+        @Override
+        public Optional<String> statusText() {
+            return Optional.empty();
+        }
+
+        @Override
+        public int statusCode() {
+            return 0;
+        }
+
+        @Override
+        public Map<String, List<String>> headers() {
+            return ImmutableMap.of(AWS_REQUEST_ID, Collections.singletonList("1234"));
+        }
+    }
+}