Repository: aws/amazon-ecs-cli-v2
Pull request #393 (2019-11-16 21:46:25)

Reduce wait time for CF polling from 30s -> 3s
If we pass in our own waiters, we can specify
the wait time that CF waits before checking
the status of creating stacks, updating them, etc.

This helps a lot, especially since we use changesets
and changeset creation is usually much faster than 30s.

a full project init w/ envs went from 8m 40s -> 7m 20s

<!-- Provide summary of changes -->

<!-- Issue number, if available. E.g. "Fixes #31", "Addresses #42, 77" -->

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.



Patch:
diff --git a/internal/pkg/deploy/cloudformation/app.go b/internal/pkg/deploy/cloudformation/app.go
index a638c3ac..18bae105 100644
--- a/internal/pkg/deploy/cloudformation/app.go
+++ b/internal/pkg/deploy/cloudformation/app.go
@@ -4,6 +4,7 @@
 package cloudformation
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/aws/aws-sdk-go/aws"
@@ -22,9 +23,9 @@ func (cf CloudFormation) DeployApp(template, stackName, changeSetName string) er
 
 	// If CreateStack does not return an error we'll wait for StackCreateComplete status.
 	if err == nil {
-		err = cf.client.WaitUntilStackCreateComplete(&cloudformation.DescribeStacksInput{
+		err = cf.client.WaitUntilStackCreateCompleteWithContext(context.Background(), &cloudformation.DescribeStacksInput{
 			StackName: aws.String(stackName),
-		})
+		}, cf.waiters...)
 
 		if err != nil {
 			return fmt.Errorf("wait for stack completion: %w", err)
@@ -62,7 +63,7 @@ func (cf CloudFormation) DeployApp(template, stackName, changeSetName string) er
 		StackName:     aws.String(stackName),
 	}
 
-	err = cf.client.WaitUntilChangeSetCreateComplete(describeChangeSetInput)
+	err = cf.client.WaitUntilChangeSetCreateCompleteWithContext(context.Background(), describeChangeSetInput, cf.waiters...)
 
 	// NOTE: If WaitUntilChangeSetCreateComplete returns an error it's possible that there
 	// are simply no changes between the previous and proposed Stack ChangeSets. We make a call to
@@ -90,9 +91,9 @@ func (cf CloudFormation) DeployApp(template, stackName, changeSetName string) er
 		return fmt.Errorf("execute change set: %w", err)
 	}
 
-	if err := cf.client.WaitUntilStackUpdateComplete(&cloudformation.DescribeStacksInput{
+	if err := cf.client.WaitUntilStackUpdateCompleteWithContext(context.Background(), &cloudformation.DescribeStacksInput{
 		StackName: aws.String(stackName),
-	}); err != nil {
+	}, cf.waiters...); err != nil {
 		return fmt.Errorf("wait for stack update: %w", err)
 	}
 
diff --git a/internal/pkg/deploy/cloudformation/changeset.go b/internal/pkg/deploy/cloudformation/changeset.go
index 2a488811..5bc96126 100644
--- a/internal/pkg/deploy/cloudformation/changeset.go
+++ b/internal/pkg/deploy/cloudformation/changeset.go
@@ -4,9 +4,11 @@
 package cloudformation
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/service/cloudformation"
 	"github.com/aws/aws-sdk-go/service/cloudformation/cloudformationiface"
 	"github.com/google/uuid"
@@ -27,7 +29,8 @@ type changeSet struct {
 	statusReason    string
 	changes         []*cloudformation.Change
 
-	c cloudformationiface.CloudFormationAPI
+	c       cloudformationiface.CloudFormationAPI
+	waiters []request.WaiterOption
 }
 
 func (set *changeSet) String() string {
@@ -35,10 +38,12 @@ func (set *changeSet) String() string {
 }
 
 func (set *changeSet) waitForCreation() error {
-	if err := set.c.WaitUntilChangeSetCreateComplete(&cloudformation.DescribeChangeSetInput{
+	describeChangeSetInput := &cloudformation.DescribeChangeSetInput{
 		ChangeSetName: aws.String(set.name),
 		StackName:     aws.String(set.stackID),
-	}); err != nil {
+	}
+
+	if err := set.c.WaitUntilChangeSetCreateCompleteWithContext(context.Background(), describeChangeSetInput, set.waiters...); err != nil {
 		return fmt.Errorf("failed to wait for changeSet creation %s: %w", set, err)
 	}
 	return nil
diff --git a/internal/pkg/deploy/cloudformation/cloudformation.go b/internal/pkg/deploy/cloudformation/cloudformation.go
index 66aa48ac..3fb0afd7 100644
--- a/internal/pkg/deploy/cloudformation/cloudformation.go
+++ b/internal/pkg/deploy/cloudformation/cloudformation.go
@@ -5,6 +5,7 @@
 package cloudformation
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"strings"
@@ -14,6 +15,7 @@ import (
 	"github.com/aws/amazon-ecs-cli-v2/templates"
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
+	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/cloudformation"
 	"github.com/aws/aws-sdk-go/service/cloudformation/cloudformationiface"
@@ -51,6 +53,7 @@ func (cf cfClientBuilder) Client(region string) cloudformationiface.CloudFormati
 // CloudFormation wraps the CloudFormationAPI interface
 type CloudFormation struct {
 	regionalClientProvider regionalClientProvider
+	waiters                []request.WaiterOption
 	client                 cloudformationiface.CloudFormationAPI
 	box                    packd.Box
 }
@@ -60,10 +63,19 @@ func New(sess *session.Session) CloudFormation {
 	cb := cfClientBuilder{
 		session: sess,
 	}
+
+	waiterOptions := []request.WaiterOption{
+		// Poll for CloudFormation updates every 3 seconds.
+		request.WithWaiterDelay(request.ConstantWaiterDelay(3 * time.Second)),
+		// Wait for at most 90 mins for any CloudFormation action.
+		request.WithWaiterMaxAttempts(1800),
+	}
+
 	return CloudFormation{
 		regionalClientProvider: cb,
 		client:                 cb.Client(*sess.Config.Region),
 		box:                    templates.Box(),
+		waiters:                waiterOptions,
 	}
 }
 
@@ -139,7 +151,7 @@ func (cf CloudFormation) waitForStackCreation(stackConfig stackConfiguration) (*
 		StackName: aws.String(stackConfig.StackName()),
 	}
 
-	if err := cf.client.WaitUntilStackCreateComplete(describeStackInput); err != nil {
+	if err := cf.client.WaitUntilStackCreateCompleteWithContext(context.Background(), describeStackInput, cf.waiters...); err != nil {
 		return nil, fmt.Errorf("failed to create stack %s: %w", stackConfig.StackName(), err)
 	}
 
@@ -221,6 +233,7 @@ func (cf CloudFormation) createChangeSet(in *cloudformation.CreateChangeSetInput
 		name:    aws.StringValue(out.Id),
 		stackID: aws.StringValue(out.StackId),
 		c:       cf.client,
+		waiters: cf.waiters,
 	}, nil
 }
 
diff --git a/internal/pkg/deploy/cloudformation/project.go b/internal/pkg/deploy/cloudformation/project.go
index 57a64dc5..49be5401 100644
--- a/internal/pkg/deploy/cloudformation/project.go
+++ b/internal/pkg/deploy/cloudformation/project.go
@@ -4,6 +4,7 @@
 package cloudformation
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"time"
@@ -93,7 +94,7 @@ func (cf CloudFormation) DelegateDNSPermissions(project *archer.Project, account
 		return fmt.Errorf("updating project to allow DNS delegation: %w", err)
 	}
 
-	return cf.client.WaitUntilStackUpdateComplete(&describeStack)
+	return cf.client.WaitUntilStackUpdateCompleteWithContext(context.Background(), &describeStack, cf.waiters...)
 }
 
 // GetProjectResourcesByRegion fetches all the regional resources for a particular region.


Test patch:
diff --git a/internal/pkg/deploy/cloudformation/app_test.go b/internal/pkg/deploy/cloudformation/app_test.go
index d4e741fe..f6d8af71 100644
--- a/internal/pkg/deploy/cloudformation/app_test.go
+++ b/internal/pkg/deploy/cloudformation/app_test.go
@@ -20,13 +20,13 @@ func TestDeployApp(t *testing.T) {
 	mockError := errors.New("mockError")
 
 	testCases := map[string]struct {
-		mockCreateStack                      func(t *testing.T, in *cloudformation.CreateStackInput) (*cloudformation.CreateStackOutput, error)
-		mockWaitUntilStackCreateComplete     func(t *testing.T, in *cloudformation.DescribeStacksInput) error
-		mockCreateChangeSet                  func(t *testing.T, in *cloudformation.CreateChangeSetInput) (*cloudformation.CreateChangeSetOutput, error)
-		mockWaitUntilChangeSetCreateComplete func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error
-		mockExecuteChangeSet                 func(t *testing.T, in *cloudformation.ExecuteChangeSetInput) (*cloudformation.ExecuteChangeSetOutput, error)
-		mockWaitUntilStackUpdateComplete     func(t *testing.T, in *cloudformation.DescribeStacksInput) error
-		mockDescribeChangeSet                func(t *testing.T, in *cloudformation.DescribeChangeSetInput) (*cloudformation.DescribeChangeSetOutput, error)
+		mockCreateStack                                 func(t *testing.T, in *cloudformation.CreateStackInput) (*cloudformation.CreateStackOutput, error)
+		mockWaitUntilStackCreateCompleteWithContext     func(t *testing.T, in *cloudformation.DescribeStacksInput) error
+		mockCreateChangeSet                             func(t *testing.T, in *cloudformation.CreateChangeSetInput) (*cloudformation.CreateChangeSetOutput, error)
+		mockWaitUntilChangeSetCreateCompleteWithContext func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error
+		mockExecuteChangeSet                            func(t *testing.T, in *cloudformation.ExecuteChangeSetInput) (*cloudformation.ExecuteChangeSetOutput, error)
+		mockWaitUntilStackUpdateCompleteWithContext     func(t *testing.T, in *cloudformation.DescribeStacksInput) error
+		mockDescribeChangeSet                           func(t *testing.T, in *cloudformation.DescribeChangeSetInput) (*cloudformation.DescribeChangeSetOutput, error)
 
 		wantErr error
 	}{
@@ -40,7 +40,7 @@ func TestDeployApp(t *testing.T) {
 
 				return &cloudformation.CreateStackOutput{}, nil
 			},
-			mockWaitUntilStackCreateComplete: func(t *testing.T, in *cloudformation.DescribeStacksInput) error {
+			mockWaitUntilStackCreateCompleteWithContext: func(t *testing.T, in *cloudformation.DescribeStacksInput) error {
 				t.Helper()
 
 				require.Equal(t, mockStackName, *in.StackName)
@@ -69,7 +69,7 @@ func TestDeployApp(t *testing.T) {
 
 				return &cloudformation.CreateChangeSetOutput{}, nil
 			},
-			mockWaitUntilChangeSetCreateComplete: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
+			mockWaitUntilChangeSetCreateCompleteWithContext: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
 				t.Helper()
 
 				require.Equal(t, mockChangeSetName, *in.ChangeSetName)
@@ -85,7 +85,7 @@ func TestDeployApp(t *testing.T) {
 
 				return &cloudformation.ExecuteChangeSetOutput{}, nil
 			},
-			mockWaitUntilStackUpdateComplete: func(t *testing.T, in *cloudformation.DescribeStacksInput) error {
+			mockWaitUntilStackUpdateCompleteWithContext: func(t *testing.T, in *cloudformation.DescribeStacksInput) error {
 				t.Helper()
 
 				require.Equal(t, mockStackName, *in.StackName)
@@ -114,7 +114,7 @@ func TestDeployApp(t *testing.T) {
 
 				return &cloudformation.CreateChangeSetOutput{}, nil
 			},
-			mockWaitUntilChangeSetCreateComplete: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
+			mockWaitUntilChangeSetCreateCompleteWithContext: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
 				t.Helper()
 
 				require.Equal(t, mockChangeSetName, *in.ChangeSetName)
@@ -154,7 +154,7 @@ func TestDeployApp(t *testing.T) {
 
 				return &cloudformation.CreateChangeSetOutput{}, nil
 			},
-			mockWaitUntilChangeSetCreateComplete: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
+			mockWaitUntilChangeSetCreateCompleteWithContext: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
 				t.Helper()
 
 				require.Equal(t, mockChangeSetName, *in.ChangeSetName)
@@ -180,13 +180,13 @@ func TestDeployApp(t *testing.T) {
 				client: mockCloudFormation{
 					t: t,
 
-					mockCreateStack:                      tc.mockCreateStack,
-					mockWaitUntilStackCreateComplete:     tc.mockWaitUntilStackCreateComplete,
-					mockCreateChangeSet:                  tc.mockCreateChangeSet,
-					mockWaitUntilChangeSetCreateComplete: tc.mockWaitUntilChangeSetCreateComplete,
-					mockExecuteChangeSet:                 tc.mockExecuteChangeSet,
-					mockWaitUntilStackUpdateComplete:     tc.mockWaitUntilStackUpdateComplete,
-					mockDescribeChangeSet:                tc.mockDescribeChangeSet,
+					mockCreateStack: tc.mockCreateStack,
+					mockWaitUntilStackCreateCompleteWithContext:     tc.mockWaitUntilStackCreateCompleteWithContext,
+					mockCreateChangeSet:                             tc.mockCreateChangeSet,
+					mockWaitUntilChangeSetCreateCompleteWithContext: tc.mockWaitUntilChangeSetCreateCompleteWithContext,
+					mockExecuteChangeSet:                            tc.mockExecuteChangeSet,
+					mockWaitUntilStackUpdateCompleteWithContext:     tc.mockWaitUntilStackUpdateCompleteWithContext,
+					mockDescribeChangeSet:                           tc.mockDescribeChangeSet,
 				},
 			}
 
diff --git a/internal/pkg/deploy/cloudformation/cloudformation_integration_test.go b/internal/pkg/deploy/cloudformation/cloudformation_integration_test.go
index 045a7143..b3e44469 100644
--- a/internal/pkg/deploy/cloudformation/cloudformation_integration_test.go
+++ b/internal/pkg/deploy/cloudformation/cloudformation_integration_test.go
@@ -1,5 +1,4 @@
 // +build integration
-
 // Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
@@ -408,7 +407,7 @@ func Test_Environment_Deployment_Integration(t *testing.T) {
 			},
 			"CFNExecutionRoleARN": func(output *awsCF.Output) {
 				require.Equal(t,
-					fmt.Sprintf("%s-CRNExecutionRoleARN", envStackName),
+					fmt.Sprintf("%s-CFNExecutionRoleARN", envStackName),
 					*output.ExportName,
 					"Should export CRNExecutionRole ARN")
 
diff --git a/internal/pkg/deploy/cloudformation/cloudformation_test.go b/internal/pkg/deploy/cloudformation/cloudformation_test.go
index 03daa561..5577b65c 100644
--- a/internal/pkg/deploy/cloudformation/cloudformation_test.go
+++ b/internal/pkg/deploy/cloudformation/cloudformation_test.go
@@ -4,6 +4,7 @@
 package cloudformation
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"regexp"
@@ -12,6 +13,7 @@ import (
 	"github.com/aws/amazon-ecs-cli-v2/internal/pkg/deploy/cloudformation/stack"
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
+	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/service/cloudformation"
 	"github.com/aws/aws-sdk-go/service/cloudformation/cloudformationiface"
 	"github.com/gobuffalo/packd"
@@ -29,32 +31,28 @@ const (
 type mockCloudFormation struct {
 	cloudformationiface.CloudFormationAPI
 
-	t                                    *testing.T
-	mockCreateChangeSet                  func(t *testing.T, in *cloudformation.CreateChangeSetInput) (*cloudformation.CreateChangeSetOutput, error)
-	mockWaitUntilChangeSetCreateComplete func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error
-	mockExecuteChangeSet                 func(t *testing.T, in *cloudformation.ExecuteChangeSetInput) (*cloudformation.ExecuteChangeSetOutput, error)
-	mockDescribeChangeSet                func(t *testing.T, in *cloudformation.DescribeChangeSetInput) (*cloudformation.DescribeChangeSetOutput, error)
-	mockWaitUntilStackCreateComplete     func(t *testing.T, in *cloudformation.DescribeStacksInput) error
-	mockDescribeStacks                   func(t *testing.T, in *cloudformation.DescribeStacksInput) (*cloudformation.DescribeStacksOutput, error)
-	mockCreateStackSet                   func(t *testing.T, in *cloudformation.CreateStackSetInput) (*cloudformation.CreateStackSetOutput, error)
-	mockDescribeStackSet                 func(t *testing.T, in *cloudformation.DescribeStackSetInput) (*cloudformation.DescribeStackSetOutput, error)
-	mockUpdateStackSet                   func(t *testing.T, in *cloudformation.UpdateStackSetInput) (*cloudformation.UpdateStackSetOutput, error)
-	mockListStackInstances               func(t *testing.T, in *cloudformation.ListStackInstancesInput) (*cloudformation.ListStackInstancesOutput, error)
-	mockCreateStackInstances             func(t *testing.T, in *cloudformation.CreateStackInstancesInput) (*cloudformation.CreateStackInstancesOutput, error)
-	mockDescribeStackSetOperation        func(t *testing.T, in *cloudformation.DescribeStackSetOperationInput) (*cloudformation.DescribeStackSetOperationOutput, error)
-	mockDescribeStackEvents              func(t *testing.T, in *cloudformation.DescribeStackEventsInput) (*cloudformation.DescribeStackEventsOutput, error)
-	mockCreateStack                      func(t *testing.T, in *cloudformation.CreateStackInput) (*cloudformation.CreateStackOutput, error)
-	mockWaitUntilStackUpdateComplete     func(t *testing.T, in *cloudformation.DescribeStacksInput) error
+	t                                               *testing.T
+	mockCreateChangeSet                             func(t *testing.T, in *cloudformation.CreateChangeSetInput) (*cloudformation.CreateChangeSetOutput, error)
+	mockWaitUntilChangeSetCreateCompleteWithContext func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error
+	mockExecuteChangeSet                            func(t *testing.T, in *cloudformation.ExecuteChangeSetInput) (*cloudformation.ExecuteChangeSetOutput, error)
+	mockDescribeChangeSet                           func(t *testing.T, in *cloudformation.DescribeChangeSetInput) (*cloudformation.DescribeChangeSetOutput, error)
+	mockWaitUntilStackCreateCompleteWithContext     func(t *testing.T, in *cloudformation.DescribeStacksInput) error
+	mockDescribeStacks                              func(t *testing.T, in *cloudformation.DescribeStacksInput) (*cloudformation.DescribeStacksOutput, error)
+	mockCreateStackSet                              func(t *testing.T, in *cloudformation.CreateStackSetInput) (*cloudformation.CreateStackSetOutput, error)
+	mockDescribeStackSet                            func(t *testing.T, in *cloudformation.DescribeStackSetInput) (*cloudformation.DescribeStackSetOutput, error)
+	mockUpdateStackSet                              func(t *testing.T, in *cloudformation.UpdateStackSetInput) (*cloudformation.UpdateStackSetOutput, error)
+	mockListStackInstances                          func(t *testing.T, in *cloudformation.ListStackInstancesInput) (*cloudformation.ListStackInstancesOutput, error)
+	mockCreateStackInstances                        func(t *testing.T, in *cloudformation.CreateStackInstancesInput) (*cloudformation.CreateStackInstancesOutput, error)
+	mockDescribeStackSetOperation                   func(t *testing.T, in *cloudformation.DescribeStackSetOperationInput) (*cloudformation.DescribeStackSetOperationOutput, error)
+	mockDescribeStackEvents                         func(t *testing.T, in *cloudformation.DescribeStackEventsInput) (*cloudformation.DescribeStackEventsOutput, error)
+	mockCreateStack                                 func(t *testing.T, in *cloudformation.CreateStackInput) (*cloudformation.CreateStackOutput, error)
+	mockWaitUntilStackUpdateCompleteWithContext     func(t *testing.T, in *cloudformation.DescribeStacksInput) error
 }
 
 func (cf mockCloudFormation) CreateChangeSet(in *cloudformation.CreateChangeSetInput) (*cloudformation.CreateChangeSetOutput, error) {
 	return cf.mockCreateChangeSet(cf.t, in)
 }
 
-func (cf mockCloudFormation) WaitUntilChangeSetCreateComplete(in *cloudformation.DescribeChangeSetInput) error {
-	return cf.mockWaitUntilChangeSetCreateComplete(cf.t, in)
-}
-
 func (cf mockCloudFormation) ExecuteChangeSet(in *cloudformation.ExecuteChangeSetInput) (*cloudformation.ExecuteChangeSetOutput, error) {
 	return cf.mockExecuteChangeSet(cf.t, in)
 }
@@ -63,10 +61,6 @@ func (cf mockCloudFormation) DescribeChangeSet(in *cloudformation.DescribeChange
 	return cf.mockDescribeChangeSet(cf.t, in)
 }
 
-func (cf mockCloudFormation) WaitUntilStackCreateComplete(in *cloudformation.DescribeStacksInput) error {
-	return cf.mockWaitUntilStackCreateComplete(cf.t, in)
-}
-
 func (cf mockCloudFormation) DescribeStacks(in *cloudformation.DescribeStacksInput) (*cloudformation.DescribeStacksOutput, error) {
 	return cf.mockDescribeStacks(cf.t, in)
 }
@@ -103,8 +97,16 @@ func (cf mockCloudFormation) CreateStack(in *cloudformation.CreateStackInput) (*
 	return cf.mockCreateStack(cf.t, in)
 }
 
-func (cf mockCloudFormation) WaitUntilStackUpdateComplete(in *cloudformation.DescribeStacksInput) error {
-	return cf.mockWaitUntilStackUpdateComplete(cf.t, in)
+func (cf mockCloudFormation) WaitUntilStackUpdateCompleteWithContext(context context.Context, in *cloudformation.DescribeStacksInput, opts ...request.WaiterOption) error {
+	return cf.mockWaitUntilStackUpdateCompleteWithContext(cf.t, in)
+}
+
+func (cf mockCloudFormation) WaitUntilChangeSetCreateCompleteWithContext(context context.Context, in *cloudformation.DescribeChangeSetInput, opts ...request.WaiterOption) error {
+	return cf.mockWaitUntilChangeSetCreateCompleteWithContext(cf.t, in)
+}
+
+func (cf mockCloudFormation) WaitUntilStackCreateCompleteWithContext(context context.Context, in *cloudformation.DescribeStacksInput, opts ...request.WaiterOption) error {
+	return cf.mockWaitUntilStackCreateCompleteWithContext(cf.t, in)
 }
 
 type mockStackConfiguration struct {
@@ -177,7 +179,7 @@ func TestDeploy(t *testing.T) {
 							StackId: aws.String(mockStackID),
 						}, nil
 					},
-					mockWaitUntilChangeSetCreateComplete: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
+					mockWaitUntilChangeSetCreateCompleteWithContext: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
 						return errors.New("some AWS error")
 					},
 				},
@@ -196,7 +198,7 @@ func TestDeploy(t *testing.T) {
 							StackId: aws.String(mockStackID),
 						}, nil
 					},
-					mockWaitUntilChangeSetCreateComplete: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
+					mockWaitUntilChangeSetCreateCompleteWithContext: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
 						return nil
 					},
 					mockDescribeChangeSet: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) (*cloudformation.DescribeChangeSetOutput, error) {
@@ -218,7 +220,7 @@ func TestDeploy(t *testing.T) {
 							StackId: aws.String(mockStackID),
 						}, nil
 					},
-					mockWaitUntilChangeSetCreateComplete: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
+					mockWaitUntilChangeSetCreateCompleteWithContext: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
 						return nil
 					},
 					mockDescribeChangeSet: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) (*cloudformation.DescribeChangeSetOutput, error) {
@@ -243,7 +245,7 @@ func TestDeploy(t *testing.T) {
 							StackId: aws.String(mockStackID),
 						}, nil
 					},
-					mockWaitUntilChangeSetCreateComplete: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
+					mockWaitUntilChangeSetCreateCompleteWithContext: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
 						return nil
 					},
 					mockDescribeChangeSet: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) (*cloudformation.DescribeChangeSetOutput, error) {
@@ -268,7 +270,7 @@ func TestDeploy(t *testing.T) {
 							StackId: aws.String(mockStackID),
 						}, nil
 					},
-					mockWaitUntilChangeSetCreateComplete: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
+					mockWaitUntilChangeSetCreateCompleteWithContext: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
 						return nil
 					},
 					mockDescribeChangeSet: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) (*cloudformation.DescribeChangeSetOutput, error) {
@@ -300,7 +302,7 @@ func TestDeploy(t *testing.T) {
 							StackId: aws.String(mockStackID),
 						}, nil
 					},
-					mockWaitUntilChangeSetCreateComplete: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
+					mockWaitUntilChangeSetCreateCompleteWithContext: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
 						return nil
 					},
 					mockDescribeChangeSet: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) (*cloudformation.DescribeChangeSetOutput, error) {
@@ -330,7 +332,7 @@ func TestDeploy(t *testing.T) {
 							StackId: aws.String(mockStackID),
 						}, nil
 					},
-					mockWaitUntilChangeSetCreateComplete: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
+					mockWaitUntilChangeSetCreateCompleteWithContext: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
 						require.Equal(t, mockStackID, *in.StackName)
 						require.Equal(t, mockChangeSetID, *in.ChangeSetName)
 						return nil
@@ -407,7 +409,7 @@ func getMockWaitStackCreateCFClient(t *testing.T, stackName string, shouldThrowE
 	return CloudFormation{
 		client: &mockCloudFormation{
 			t: t,
-			mockWaitUntilStackCreateComplete: func(t *testing.T, input *cloudformation.DescribeStacksInput) error {
+			mockWaitUntilStackCreateCompleteWithContext: func(t *testing.T, input *cloudformation.DescribeStacksInput) error {
 				require.Equal(t, stackName, *input.StackName)
 				if shouldThrowError {
 					return fmt.Errorf("some AWS error")
diff --git a/internal/pkg/deploy/cloudformation/env_test.go b/internal/pkg/deploy/cloudformation/env_test.go
index 49e65193..0e107cd2 100644
--- a/internal/pkg/deploy/cloudformation/env_test.go
+++ b/internal/pkg/deploy/cloudformation/env_test.go
@@ -23,9 +23,9 @@ func TestStreamEnvironmentCreation(t *testing.T) {
 	testCases := map[string]struct {
 		in *deploy.CreateEnvironmentInput
 
-		mockWaitUntilStackCreateComplete func(t *testing.T, in *cloudformation.DescribeStacksInput) error
-		mockDescribeStacks               func(t *testing.T, in *cloudformation.DescribeStacksInput) (*cloudformation.DescribeStacksOutput, error)
-		mockDescribeStackEvents          func(t *testing.T, in *cloudformation.DescribeStackEventsInput) (*cloudformation.DescribeStackEventsOutput, error)
+		mockWaitUntilStackCreateCompleteWithContext func(t *testing.T, in *cloudformation.DescribeStacksInput) error
+		mockDescribeStacks                          func(t *testing.T, in *cloudformation.DescribeStacksInput) (*cloudformation.DescribeStacksOutput, error)
+		mockDescribeStackEvents                     func(t *testing.T, in *cloudformation.DescribeStackEventsInput) (*cloudformation.DescribeStackEventsOutput, error)
 
 		wantedEvents []deploy.ResourceEvent
 		wantedResult deploy.CreateEnvironmentResponse
@@ -36,7 +36,7 @@ func TestStreamEnvironmentCreation(t *testing.T) {
 				Name:    "test",
 			},
 
-			mockWaitUntilStackCreateComplete: func(t *testing.T, in *cloudformation.DescribeStacksInput) error {
+			mockWaitUntilStackCreateCompleteWithContext: func(t *testing.T, in *cloudformation.DescribeStacksInput) error {
 				require.Equal(t, "phonetool-test", *in.StackName, "stack names should be equal to each other")
 				return errors.New("wait until error")
 			},
@@ -75,7 +75,7 @@ func TestStreamEnvironmentCreation(t *testing.T) {
 				Name:    "test",
 			},
 
-			mockWaitUntilStackCreateComplete: func(t *testing.T, in *cloudformation.DescribeStacksInput) error {
+			mockWaitUntilStackCreateCompleteWithContext: func(t *testing.T, in *cloudformation.DescribeStacksInput) error {
 				require.Equal(t, "phonetool-test", *in.StackName, "stack names should be equal to each other")
 				return errors.New("wait until error")
 			},
@@ -96,7 +96,7 @@ func TestStreamEnvironmentCreation(t *testing.T) {
 				Name:    "test",
 			},
 
-			mockWaitUntilStackCreateComplete: func(t *testing.T, in *cloudformation.DescribeStacksInput) error {
+			mockWaitUntilStackCreateCompleteWithContext: func(t *testing.T, in *cloudformation.DescribeStacksInput) error {
 				require.Equal(t, "phonetool-test", *in.StackName, "stack names should be equal to each other")
 				return nil
 			},
@@ -151,10 +151,10 @@ func TestStreamEnvironmentCreation(t *testing.T) {
 			// GIVEN
 			cf := CloudFormation{
 				client: &mockCloudFormation{
-					t:                                t,
-					mockDescribeStackEvents:          tc.mockDescribeStackEvents,
-					mockDescribeStacks:               tc.mockDescribeStacks,
-					mockWaitUntilStackCreateComplete: tc.mockWaitUntilStackCreateComplete,
+					t:                       t,
+					mockDescribeStackEvents: tc.mockDescribeStackEvents,
+					mockDescribeStacks:      tc.mockDescribeStacks,
+					mockWaitUntilStackCreateCompleteWithContext: tc.mockWaitUntilStackCreateCompleteWithContext,
 				},
 				box: emptyEnvBox(),
 			}
diff --git a/internal/pkg/deploy/cloudformation/project_test.go b/internal/pkg/deploy/cloudformation/project_test.go
index ef918c28..3f6f9fec 100644
--- a/internal/pkg/deploy/cloudformation/project_test.go
+++ b/internal/pkg/deploy/cloudformation/project_test.go
@@ -1041,10 +1041,10 @@ func TestDelegateDNSPermissions(t *testing.T) {
 							ExecutionStatus: aws.String(cloudformation.ExecutionStatusAvailable),
 						}, nil
 					},
-					mockWaitUntilChangeSetCreateComplete: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
+					mockWaitUntilChangeSetCreateCompleteWithContext: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
 						return nil
 					},
-					mockWaitUntilStackUpdateComplete: func(t *testing.T, in *cloudformation.DescribeStacksInput) error {
+					mockWaitUntilStackUpdateCompleteWithContext: func(t *testing.T, in *cloudformation.DescribeStacksInput) error {
 						return nil
 					},
 				}
@@ -1106,7 +1106,7 @@ func TestDelegateDNSPermissions(t *testing.T) {
 							ExecutionStatus: aws.String(cloudformation.ExecutionStatusAvailable),
 						}, nil
 					},
-					mockWaitUntilChangeSetCreateComplete: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
+					mockWaitUntilChangeSetCreateCompleteWithContext: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
 						return nil
 					},
 				}
@@ -1141,11 +1141,11 @@ func getMockSuccessfulDeployCFClient(t *testing.T, stackName string) *mockCloudF
 				StackId: aws.String(stackName),
 			}, nil
 		},
-		mockWaitUntilChangeSetCreateComplete: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
+		mockWaitUntilChangeSetCreateCompleteWithContext: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) error {
 
 			return nil
 		},
-		mockWaitUntilStackCreateComplete: func(t *testing.T, input *cloudformation.DescribeStacksInput) error {
+		mockWaitUntilStackCreateCompleteWithContext: func(t *testing.T, input *cloudformation.DescribeStacksInput) error {
 			return nil
 		},
 		mockDescribeChangeSet: func(t *testing.T, in *cloudformation.DescribeChangeSetInput) (*cloudformation.DescribeChangeSetOutput, error) {