Append User-Agent for Waiters operations

codegen/src/main/java/software/amazon/awssdk/codegen/poet/waiters/BaseWaiterClassSpec.java

@@ -29,6 +29,7 @@
 import com.squareup.javapoet.ParameterizedTypeName;
 import com.squareup.javapoet.TypeName;
 import com.squareup.javapoet.TypeSpec;
+import com.squareup.javapoet.TypeVariableName;
 import com.squareup.javapoet.WildcardTypeName;
 import java.time.Duration;
 import java.util.ArrayList;
@@ -37,19 +38,23 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.function.Consumer;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import javax.lang.model.element.Modifier;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.annotations.ThreadSafe;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.awscore.exception.AwsServiceException;
 import software.amazon.awssdk.codegen.emitters.tasks.WaitersRuntimeGeneratorTask;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
 import software.amazon.awssdk.codegen.model.service.Acceptor;
 import software.amazon.awssdk.codegen.model.service.WaiterDefinition;
 import software.amazon.awssdk.codegen.poet.ClassSpec;
+import software.amazon.awssdk.codegen.poet.PoetExtensions;
 import software.amazon.awssdk.codegen.poet.PoetUtils;
+import software.amazon.awssdk.core.ApiName;
 import software.amazon.awssdk.core.internal.waiters.WaiterAttribute;
 import software.amazon.awssdk.core.retry.backoff.FixedDelayBackoffStrategy;
 import software.amazon.awssdk.core.waiters.PollingStrategy;
@@ -62,18 +67,22 @@
  * Base class containing common logic shared between the sync waiter class and the async waiter class
  */
 public abstract class BaseWaiterClassSpec implements ClassSpec {
+
+    private static final String WAITERS_USER_AGENT = "waiter";
     private final IntermediateModel model;
     private final String modelPackage;
     private final Map<String, WaiterDefinition> waiters;
     private final ClassName waiterClassName;
     private final JmesPathAcceptorGenerator jmesPathAcceptorGenerator;
+    private final PoetExtensions poetExtensions;
 
     public BaseWaiterClassSpec(IntermediateModel model, ClassName waiterClassName) {
         this.model = model;
         this.modelPackage = model.getMetadata().getFullModelPackageName();
         this.waiters = model.getWaiters();
         this.waiterClassName = waiterClassName;
         this.jmesPathAcceptorGenerator = new JmesPathAcceptorGenerator(waitersRuntimeClass());
+        this.poetExtensions = new PoetExtensions(model);
     }
 
     @Override
@@ -105,6 +114,7 @@ public TypeSpec poetSpec() {
                                             .build());
 
         typeSpecBuilder.addType(builder());
+        typeSpecBuilder.addMethod(applyWaitersUserAgentMethod(poetExtensions, model));
         return typeSpecBuilder.build();
     }
 
@@ -277,7 +287,7 @@ private MethodSpec waiterOperation(Map.Entry<String, WaiterDefinition> waiterDef
             .addParameter(requestType, opModel.getInput().getVariableName())
             .addModifiers(PUBLIC)
             .addAnnotation(Override.class)
-            .addStatement("return $L.$L(() -> client.$N($N))",
+            .addStatement("return $L.$L(() -> client.$N(applyWaitersUserAgent($N)))",
                           waiterFieldName(waiterMethodName),
                           waiterClassName.simpleName().equals("Waiter") ? "run" : "runAsync",
                           lowercaseFirstChar(waiterDefinition.getValue().getOperation()),
@@ -326,6 +336,40 @@ private MethodSpec.Builder methodSignatureWithReturnType(String waiterMethodName
                          .returns(getWaiterResponseType(opModel));
     }
 
+    static MethodSpec applyWaitersUserAgentMethod(PoetExtensions poetExtensions, IntermediateModel model) {
+
+        TypeVariableName typeVariableName =
+            TypeVariableName.get("T", poetExtensions.getModelClass(model.getSdkRequestBaseClassName()));
+
+        ParameterizedTypeName parameterizedTypeName = ParameterizedTypeName
+            .get(ClassName.get(Consumer.class), ClassName.get(AwsRequestOverrideConfiguration.Builder.class));
+
+        CodeBlock codeBlock = CodeBlock.builder()
+                                       .addStatement("$T userAgentApplier = b -> b.addApiName($T.builder().version"
+                                                     + "($S).name($S).build())",
+                                                     parameterizedTypeName, ApiName.class,
+                                                     WAITERS_USER_AGENT,
+                                                     "hll")
+                                       .addStatement("$T overrideConfiguration =\n"
+                                                     + "            request.overrideConfiguration().map(c -> c.toBuilder()"
+                                                     + ".applyMutation"
+                                                     + "(userAgentApplier).build())\n"
+                                                     + "            .orElse((AwsRequestOverrideConfiguration.builder()"
+                                                     + ".applyMutation"
+                                                     + "(userAgentApplier).build()))", AwsRequestOverrideConfiguration.class)
+                                       .addStatement("return (T) request.toBuilder().overrideConfiguration"
+                                                     + "(overrideConfiguration).build()")
+                                       .build();
+
+        return MethodSpec.methodBuilder("applyWaitersUserAgent")
+                         .addModifiers(Modifier.PRIVATE)
+                         .addParameter(typeVariableName, "request")
+                         .addTypeVariable(typeVariableName)
+                         .addCode(codeBlock)
+                         .returns(typeVariableName)
+                         .build();
+    }
+
     private String getWaiterMethodName(String waiterMethodName) {
         return "waitUntil" + waiterMethodName;
     }

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/waiters/query-async-waiter-class.java

@@ -7,10 +7,13 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.function.Consumer;
 import software.amazon.awssdk.annotations.Generated;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.annotations.ThreadSafe;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.awscore.exception.AwsServiceException;
+import software.amazon.awssdk.core.ApiName;
 import software.amazon.awssdk.core.internal.waiters.WaiterAttribute;
 import software.amazon.awssdk.core.retry.backoff.FixedDelayBackoffStrategy;
 import software.amazon.awssdk.core.waiters.AsyncWaiter;
@@ -21,6 +24,7 @@
 import software.amazon.awssdk.services.query.QueryAsyncClient;
 import software.amazon.awssdk.services.query.model.APostOperationRequest;
 import software.amazon.awssdk.services.query.model.APostOperationResponse;
+import software.amazon.awssdk.services.query.model.QueryRequest;
 import software.amazon.awssdk.services.query.waiters.internal.WaitersRuntime;
 import software.amazon.awssdk.utils.AttributeMap;
 import software.amazon.awssdk.utils.SdkAutoCloseable;
@@ -77,7 +81,7 @@ private static String errorCode(Throwable error) {
     @Override
     public CompletableFuture<WaiterResponse<APostOperationResponse>> waitUntilPostOperationSuccess(
         APostOperationRequest aPostOperationRequest) {
-        return postOperationSuccessWaiter.runAsync(() -> client.aPostOperation(aPostOperationRequest));
+        return postOperationSuccessWaiter.runAsync(() -> client.aPostOperation(applyWaitersUserAgent(aPostOperationRequest)));
     }
 
     private static List<WaiterAcceptor<? super APostOperationResponse>> postOperationSuccessWaiterAcceptors() {
@@ -102,6 +106,15 @@ public static QueryAsyncWaiter.Builder builder() {
         return new DefaultBuilder();
     }
 
+    private <T extends QueryRequest> T applyWaitersUserAgent(T request) {
+        Consumer<AwsRequestOverrideConfiguration.Builder> userAgentApplier = b -> b.addApiName(ApiName.builder()
+                                                                                                      .version("waiter").name("hll").build());
+        AwsRequestOverrideConfiguration overrideConfiguration = request.overrideConfiguration()
+                                                                       .map(c -> c.toBuilder().applyMutation(userAgentApplier).build())
+                                                                       .orElse((AwsRequestOverrideConfiguration.builder().applyMutation(userAgentApplier).build()));
+        return (T) request.toBuilder().overrideConfiguration(overrideConfiguration).build();
+    }
+
     public static final class DefaultBuilder implements QueryAsyncWaiter.Builder {
         private QueryAsyncClient client;
 

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/waiters/query-sync-waiter-class.java

@@ -4,10 +4,13 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import java.util.function.Consumer;
 import software.amazon.awssdk.annotations.Generated;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.annotations.ThreadSafe;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.awscore.exception.AwsServiceException;
+import software.amazon.awssdk.core.ApiName;
 import software.amazon.awssdk.core.internal.waiters.WaiterAttribute;
 import software.amazon.awssdk.core.retry.backoff.FixedDelayBackoffStrategy;
 import software.amazon.awssdk.core.waiters.PollingStrategy;
@@ -18,6 +21,7 @@
 import software.amazon.awssdk.services.query.QueryClient;
 import software.amazon.awssdk.services.query.model.APostOperationRequest;
 import software.amazon.awssdk.services.query.model.APostOperationResponse;
+import software.amazon.awssdk.services.query.model.QueryRequest;
 import software.amazon.awssdk.services.query.waiters.internal.WaitersRuntime;
 import software.amazon.awssdk.utils.AttributeMap;
 import software.amazon.awssdk.utils.SdkAutoCloseable;
@@ -59,7 +63,7 @@ private static String errorCode(Throwable error) {
 
     @Override
     public WaiterResponse<APostOperationResponse> waitUntilPostOperationSuccess(APostOperationRequest aPostOperationRequest) {
-        return postOperationSuccessWaiter.run(() -> client.aPostOperation(aPostOperationRequest));
+        return postOperationSuccessWaiter.run(() -> client.aPostOperation(applyWaitersUserAgent(aPostOperationRequest)));
     }
 
     private static List<WaiterAcceptor<? super APostOperationResponse>> postOperationSuccessWaiterAcceptors() {
@@ -84,6 +88,15 @@ public static QueryWaiter.Builder builder() {
         return new DefaultBuilder();
     }
 
+    private <T extends QueryRequest> T applyWaitersUserAgent(T request) {
+        Consumer<AwsRequestOverrideConfiguration.Builder> userAgentApplier = b -> b.addApiName(ApiName.builder()
+                                                                                                      .version("waiter").name("hll").build());
+        AwsRequestOverrideConfiguration overrideConfiguration = request.overrideConfiguration()
+                                                                       .map(c -> c.toBuilder().applyMutation(userAgentApplier).build())
+                                                                       .orElse((AwsRequestOverrideConfiguration.builder().applyMutation(userAgentApplier).build()));
+        return (T) request.toBuilder().overrideConfiguration(overrideConfiguration).build();
+    }
+
     public static final class DefaultBuilder implements QueryWaiter.Builder {
         private QueryClient client;
 

services/autoscaling/src/it/java/software/amazon/awssdk/services/autoscaling/waiters/AutoScalingWaiterTest.java

@@ -16,6 +16,7 @@
 package software.amazon.awssdk.services.autoscaling.waiters;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 import static software.amazon.awssdk.services.autoscaling.model.LifecycleState.IN_SERVICE;
@@ -68,7 +69,7 @@ public void waitUntilGroupInServiceWorks() {
                                                                                      i -> i.lifecycleState(IN_SERVICE)))
                                              .build();
 
-        when(client.describeAutoScalingGroups(request)).thenReturn(response1, response2);
+        when(client.describeAutoScalingGroups(any(DescribeAutoScalingGroupsRequest.class))).thenReturn(response1, response2);
 
         AutoScalingWaiter waiter = AutoScalingWaiter.builder()
                                                     .pollingStrategy(PollingStrategy.builder()

services/dynamodb/src/test/java/software/amazon/awssdk/services/dynamodb/WaitersUserAgentTest.java

@@ -0,0 +1,112 @@
+/*
+ * 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.services.dynamodb;
+
+import static com.github.tomakehurst.wiremock.client.WireMock.aResponse;
+import static com.github.tomakehurst.wiremock.client.WireMock.any;
+import static com.github.tomakehurst.wiremock.client.WireMock.stubFor;
+import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo;
+import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.options;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.times;
+
+import com.github.tomakehurst.wiremock.common.ConsoleNotifier;
+import com.github.tomakehurst.wiremock.junit.WireMockRule;
+import java.net.URI;
+import java.util.concurrent.CompletableFuture;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Matchers;
+import org.mockito.Mockito;
+import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
+import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
+import software.amazon.awssdk.core.interceptor.Context;
+import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
+import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
+import software.amazon.awssdk.core.waiters.WaiterResponse;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.dynamodb.model.DescribeTableRequest;
+import software.amazon.awssdk.services.dynamodb.model.DescribeTableResponse;
+import software.amazon.awssdk.services.dynamodb.waiters.DynamoDbAsyncWaiter;
+import software.amazon.awssdk.services.dynamodb.waiters.DynamoDbWaiter;
+
+public class WaitersUserAgentTest {
+
+    @Rule
+    public WireMockRule mockServer = new WireMockRule(options().notifier(new ConsoleNotifier(true)).dynamicHttpsPort());
+
+    private DynamoDbClient dynamoDbClient;
+    private DynamoDbAsyncClient dynamoDbAsyncClient;
+    private ExecutionInterceptor interceptor;
+
+    @Before
+    public void setup() {
+        interceptor = Mockito.spy(AbstractExecutionInterceptor.class);
+
+        dynamoDbClient = DynamoDbClient.builder()
+                                       .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("test",
+                                                                                                                        "test")))
+                                       .region(Region.US_WEST_2)
+                                       .endpointOverride(URI.create("http://localhost:" + mockServer.port()))
+                                       .overrideConfiguration(c -> c.addExecutionInterceptor(interceptor))
+                                       .build();
+
+        dynamoDbAsyncClient = DynamoDbAsyncClient.builder()
+                                                 .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials
+                                                                                                           .create("test",
+                                                                                                                   "test")))
+                                                 .region(Region.US_WEST_2)
+                                                 .endpointOverride(URI.create("http://localhost:" + mockServer.port()))
+                                                 .overrideConfiguration(c -> c.addExecutionInterceptor(interceptor))
+                                                 .build();
+    }
+
+    @Test
+    public void syncWaiters_shouldHaveWaitersUserAgent() {
+        stubFor(any(urlEqualTo("/")).willReturn(aResponse().withStatus(500)));
+
+        DynamoDbWaiter waiter = dynamoDbClient.waiter();
+        assertThatThrownBy(() -> waiter.waitUntilTableExists(DescribeTableRequest.builder().tableName("table").build())).isNotNull();
+
+        ArgumentCaptor<Context.BeforeTransmission> context = ArgumentCaptor.forClass(Context.BeforeTransmission.class);
+        Mockito.verify(interceptor).beforeTransmission(context.capture(), Matchers.any());
+
+        assertTrue(context.getValue().httpRequest().headers().get("User-Agent").toString().contains("waiter"));
+    }
+
+    @Test
+    public void asyncWaiters_shouldHaveWaitersUserAgent() {
+        DynamoDbAsyncWaiter waiter = dynamoDbAsyncClient.waiter();
+        CompletableFuture<WaiterResponse<DescribeTableResponse>> responseFuture = waiter.waitUntilTableExists(DescribeTableRequest.builder().tableName("table").build());
+
+        ArgumentCaptor<Context.BeforeTransmission> context = ArgumentCaptor.forClass(Context.BeforeTransmission.class);
+        Mockito.verify(interceptor).beforeTransmission(context.capture(), Matchers.any());
+
+        assertTrue(context.getValue().httpRequest().headers().get("User-Agent").toString().contains("waiter"));
+
+        responseFuture.cancel(true);
+    }
+
+    public static abstract class AbstractExecutionInterceptor implements ExecutionInterceptor {
+        @Override
+        public void beforeTransmission(Context.BeforeTransmission context, ExecutionAttributes executionAttributes) {
+            throw new RuntimeException("Interrupting the request.");
+        }
+    }
+}

services/ecs/src/it/java/software/amazon/awssdk/services/ecs/waiters/EcsWaiterTest.java

@@ -16,6 +16,7 @@
 package software.amazon.awssdk.services.ecs.waiters;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -57,7 +58,7 @@ public void waitUntilServicesStableWorks() {
                                                     .runningCount(2))
                                     .build();
 
-        when(client.describeServices(request)).thenReturn(response1, response2);
+        when(client.describeServices(any(DescribeServicesRequest.class))).thenReturn(response1, response2);
 
         EcsWaiter waiter = EcsWaiter.builder()
                                     .pollingStrategy(PollingStrategy.builder()