From 10a6e5a727826062dde099187b66c7018f3705f3 Mon Sep 17 00:00:00 2001
From: Jason Del Ponte <delpontej@gmail.com>
Date: Wed, 12 Aug 2020 10:44:36 -0700
Subject: [PATCH] service/ec2: Fix max retries with client customizations
 (#3465)

Fixes #3374 by correcting the EC2 API client's customization for ModifyNetworkInterfaceAttribute and AssignPrivateIpAddresses operations to use the aws.Config.MaxRetries value if set. Previously the API client's customizations would ignore MaxRetries specified in the SDK's aws.Config.MaxRetries field.
---
 CHANGELOG_PENDING.md               |  2 +
 service/ec2/customizations.go      |  6 ++-
 service/ec2/customizations_test.go | 59 ++++++++++++++++++++++++++++++
 3 files changed, 66 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG_PENDING.md b/CHANGELOG_PENDING.md
index 53029d35bf..fdf85cdf10 100644
--- a/CHANGELOG_PENDING.md
+++ b/CHANGELOG_PENDING.md
@@ -12,3 +12,5 @@
   * Fixes [#3410](https://github.com/aws/aws-sdk-go/issues/3410)
 * `codegen`: Export event stream constructor for easier mocking ([#3473](https://github.com/aws/aws-sdk-go/pull/3473))
   * Fixes [#3412](https://github.com/aws/aws-sdk-go/issues/3412) by exporting the operation's EventStream type's constructor function so it can be used to fully initialize fully when mocking out behavior for API operations with event streams.
+* `service/ec2`: Fix max retries with client customizations ([#3465](https://github.com/aws/aws-sdk-go/pull/3465))
+  * Fixes [#3374](https://github.com/aws/aws-sdk-go/issues/3374) by correcting the EC2 API client's customization for ModifyNetworkInterfaceAttribute and AssignPrivateIpAddresses operations to use the aws.Config.MaxRetries value if set. Previously the API client's customizations would ignore MaxRetries specified in the SDK's aws.Config.MaxRetries field.
diff --git a/service/ec2/customizations.go b/service/ec2/customizations.go
index efec8d8a94..3ad3059189 100644
--- a/service/ec2/customizations.go
+++ b/service/ec2/customizations.go
@@ -26,8 +26,12 @@ func init() {
 
 		// only set the retryer on request if config doesn't have a retryer
 		if r.Config.Retryer == nil && (r.Operation.Name == opModifyNetworkInterfaceAttribute || r.Operation.Name == opAssignPrivateIpAddresses) {
+			maxRetries := client.DefaultRetryerMaxNumRetries
+			if m := r.Config.MaxRetries; m != nil && *m != aws.UseServiceDefaultRetries {
+				maxRetries = *m
+			}
 			r.Retryer = client.DefaultRetryer{
-				NumMaxRetries:    client.DefaultRetryerMaxNumRetries,
+				NumMaxRetries:    maxRetries,
 				MinRetryDelay:    customRetryerMinRetryDelay,
 				MinThrottleDelay: customRetryerMinRetryDelay,
 				MaxRetryDelay:    customRetryerMaxRetryDelay,
diff --git a/service/ec2/customizations_test.go b/service/ec2/customizations_test.go
index 08ff2cf0b6..975442227a 100644
--- a/service/ec2/customizations_test.go
+++ b/service/ec2/customizations_test.go
@@ -1,12 +1,19 @@
+// +build go1.7
+
 package ec2_test
 
 import (
+	"bytes"
+	"context"
 	"io/ioutil"
+	"net/http"
 	"net/url"
 	"regexp"
 	"testing"
 
 	"github.com/aws/aws-sdk-go/aws"
+	sdkclient "github.com/aws/aws-sdk-go/aws/client"
+	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/awstesting/unit"
 	"github.com/aws/aws-sdk-go/service/ec2"
 )
@@ -46,3 +53,55 @@ func TestCopySnapshotPresignedURL(t *testing.T) {
 		t.Errorf("expect %v to match, got %v", r.String(), u)
 	}
 }
+
+func TestNoCustomRetryerWithMaxRetries(t *testing.T) {
+	cases := map[string]struct {
+		Config           aws.Config
+		ExpectMaxRetries int
+	}{
+		"With custom retrier": {
+			Config: aws.Config{
+				Retryer: sdkclient.DefaultRetryer{
+					NumMaxRetries: 10,
+				},
+			},
+			ExpectMaxRetries: 10,
+		},
+		"with max retries": {
+			Config: aws.Config{
+				MaxRetries: aws.Int(10),
+			},
+			ExpectMaxRetries: 10,
+		},
+		"no options set": {
+			ExpectMaxRetries: sdkclient.DefaultRetryerMaxNumRetries,
+		},
+	}
+
+	for name, c := range cases {
+		t.Run(name, func(t *testing.T) {
+			client := ec2.New(unit.Session, &aws.Config{
+				DisableParamValidation: aws.Bool(true),
+			}, c.Config.Copy())
+			client.ModifyNetworkInterfaceAttributeWithContext(context.Background(), nil, checkRetryerMaxRetries(t, c.ExpectMaxRetries))
+			client.AssignPrivateIpAddressesWithContext(context.Background(), nil, checkRetryerMaxRetries(t, c.ExpectMaxRetries))
+		})
+	}
+
+}
+
+func checkRetryerMaxRetries(t *testing.T, maxRetries int) func(*request.Request) {
+	return func(r *request.Request) {
+		r.Handlers.Send.Clear()
+		r.Handlers.Send.PushBack(func(rr *request.Request) {
+			if e, a := maxRetries, rr.Retryer.MaxRetries(); e != a {
+				t.Errorf("%s, expect %v max retries, got %v", rr.Operation.Name, e, a)
+			}
+			rr.HTTPResponse = &http.Response{
+				StatusCode: 200,
+				Header:     http.Header{},
+				Body:       ioutil.NopCloser(&bytes.Buffer{}),
+			}
+		})
+	}
+}