Repository: keikoproj/upgrade-manager (go)

Issue:
RollingUpgrade may rotate multiple nodes if it is restarted
**Is this a BUG REPORT or FEATURE REQUEST?**:
Bug

**What happened**:
It seems if the rolling upgrade controller is restarted/killed during postDrainWait, the new pod might pick different instances to start rotating (randomUpdate), and we can end up with multiple nodes in the same scaling group which are getting drained.

**What you expected to happen**:
RollingUpgrade should annotate the node it's working on and resume same nodes if it gets killed

**How to reproduce it (as minimally and precisely as possible)**:
Start a rolling upgrade, and delete the controller pod once you see nodes in SchedulingDisabled.

With 3 rollingupgrade CR submitted - should have 3 draining nodes only:
```
NAME                                           STATUS                     ROLES    AGE   VERSION
ip-222-111-232-192.us-west-2.compute.internal   Ready                      master   10d   v1.14.8
ip-222-111-232-228.us-west-2.compute.internal   Ready                      node     10d   v1.14.8
ip-222-111-232-236.us-west-2.compute.internal   Ready                      node     10d   v1.14.8
ip-222-111-232-243.us-west-2.compute.internal   Ready,SchedulingDisabled   node     10d   v1.14.8
ip-222-111-232-49.us-west-2.compute.internal    Ready,SchedulingDisabled   node     25m   v1.14.8
ip-222-111-233-105.us-west-2.compute.internal   Ready                      node     10d   v1.14.8
ip-222-111-233-182.us-west-2.compute.internal   Ready                      node     10d   v1.14.8
ip-222-111-233-246.us-west-2.compute.internal   Ready                      master   10d   v1.14.8
ip-222-111-233-59.us-west-2.compute.internal    Ready                      node     26m   v1.14.8
ip-222-111-233-99.us-west-2.compute.internal    Ready,SchedulingDisabled   node     25m   v1.14.8
ip-222-111-234-114.us-west-2.compute.internal   NotReady                   node     4s    v1.14.8
ip-222-111-234-126.us-west-2.compute.internal   Ready                      node     26m   v1.14.8
ip-222-111-234-145.us-west-2.compute.internal   Ready,SchedulingDisabled   node     10d   v1.14.8
ip-222-111-234-158.us-west-2.compute.internal   Ready                      master   10d   v1.14.8
ip-222-111-234-250.us-west-2.compute.internal   Ready                      node     25m   v1.14.8
ip-222-111-234-57.us-west-2.compute.internal    Ready,SchedulingDisabled   node     10d   v1.14.8
```

Pull request description:
RollingUpgrade idempotency fix
Fixes #37 

This PR makes the following changes:
- Before selecting nodes for restack, we make one API call (per reconcile) to find instances with tag `upgrademgr.keikoproj.io/state = in-progress` and then compare against CR's ASG seen instances.
- If instances as such are found, they override the node selector, otherwise we proceed with the node selector normally.
- Every instance worked on gets tagged with `upgrademgr.keikoproj.io/state = in-progress`, and set to `completed` after termination is done


Patch:
diff --git a/README.md b/README.md
index da7b41b..8f8cca4 100644
--- a/README.md
+++ b/README.md
@@ -63,7 +63,8 @@ If you already have an existing cluster created using kops, follow the instructi
 {
     "Effect": "Allow",
     "Action": [
-        "ec2:TerminateInstances",
+        "ec2:CreateTags",
+        "ec2:DescribeInstances",
         "autoscaling:DescribeAutoScalingGroups",
         "autoscaling:TerminateInstanceInAutoScalingGroup"
     ],
diff --git a/controllers/helpers.go b/controllers/helpers.go
index 2e45dae..2f49fde 100644
--- a/controllers/helpers.go
+++ b/controllers/helpers.go
@@ -1,7 +1,11 @@
 package controllers
 
 import (
+	"fmt"
+	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/autoscaling"
+	"github.com/aws/aws-sdk-go/service/ec2"
+	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
 	upgrademgrv1alpha1 "github.com/keikoproj/upgrade-manager/api/v1alpha1"
 	"log"
 	"strconv"
@@ -32,6 +36,58 @@ func getMaxUnavailable(strategy upgrademgrv1alpha1.UpdateStrategy, totalNodes in
 	return maxUnavailable
 }
 
+func tagEC2instance(instanceID, tagKey, tagValue string, client ec2iface.EC2API) error {
+	input := &ec2.CreateTagsInput{
+		Resources: aws.StringSlice([]string{instanceID}),
+		Tags: []*ec2.Tag{
+			{
+				Key:   aws.String(tagKey),
+				Value: aws.String(tagValue),
+			},
+		},
+	}
+	_, err := client.CreateTags(input)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func getTaggedInstances(tagKey, tagValue string, client ec2iface.EC2API) ([]string, error) {
+	instances := []string{}
+	key := fmt.Sprintf("tag:%v", tagKey)
+	input := &ec2.DescribeInstancesInput{
+		Filters: []*ec2.Filter{
+			{
+				Name:   aws.String(key),
+				Values: aws.StringSlice([]string{tagValue}),
+			},
+		},
+	}
+
+	err := client.DescribeInstancesPages(input, func(page *ec2.DescribeInstancesOutput, lastPage bool) bool {
+		for _, res := range page.Reservations {
+			for _, instance := range res.Instances {
+				instances = append(instances, aws.StringValue(instance.InstanceId))
+			}
+		}
+		return page.NextToken != nil
+	})
+	if err != nil {
+		return instances, err
+	}
+	return instances, nil
+}
+
+func contains(s []string, e string) bool {
+	for _, a := range s {
+		if a == e {
+			return true
+		}
+	}
+	return false
+}
+
 // getNextAvailableInstances checks the cluster state store for the instance state
 // and returns the next set of instances available for update
 func getNextAvailableInstances(
diff --git a/controllers/rollingupgrade_controller.go b/controllers/rollingupgrade_controller.go
index 46fae51..9f79ef8 100644
--- a/controllers/rollingupgrade_controller.go
+++ b/controllers/rollingupgrade_controller.go
@@ -37,7 +37,8 @@ import (
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/autoscaling"
 	"github.com/aws/aws-sdk-go/service/autoscaling/autoscalingiface"
-
+	"github.com/aws/aws-sdk-go/service/ec2"
+	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
 	"github.com/go-logr/logr"
 	upgrademgrv1alpha1 "github.com/keikoproj/upgrade-manager/api/v1alpha1"
 	log "github.com/keikoproj/upgrade-manager/pkg/log"
@@ -65,6 +66,8 @@ const (
 	ClearCompletedFrequency = "1d"
 	// ClearErrorFrequency is the time after which an errored rollingUpgrade object is deleted.
 	ClearErrorFrequency = "7d"
+	// EC2StateTagKey is the EC2 tag key for indicating the state
+	EC2StateTagKey = "upgrademgr.keikoproj.io/state"
 
 	// Environment variable keys
 	asgNameKey      = "ASG_NAME"
@@ -96,6 +99,8 @@ var DefaultRetryer = awsclient.DefaultRetryer{
 type RollingUpgradeReconciler struct {
 	client.Client
 	Log             logr.Logger
+	EC2Client       ec2iface.EC2API
+	ASGClient       autoscalingiface.AutoScalingAPI
 	generatedClient *kubernetes.Clientset
 	Asg             *autoscaling.Group
 	NodeList        *corev1.NodeList
@@ -156,7 +161,7 @@ func (r *RollingUpgradeReconciler) postDrainHelper(ruObj *upgrademgrv1alpha1.Rol
 			return errors.New(msg)
 		}
 	}
-	log.Printf("%s: Waiting for postDrainDelay", ruObj.Name)
+	log.Printf("%s: Waiting for postDrainDelay of %v seconds", ruObj.Name, ruObj.Spec.PostDrainDelaySeconds)
 	time.Sleep(time.Duration(ruObj.Spec.PostDrainDelaySeconds) * time.Second)
 
 	if ruObj.Spec.PostDrain.PostWaitScript != "" {
@@ -261,14 +266,14 @@ func (r *RollingUpgradeReconciler) WaitForTermination(nodeName string, nodeInter
 }
 
 // TerminateNode actually terminates the given node.
-func (r *RollingUpgradeReconciler) TerminateNode(ruObj *upgrademgrv1alpha1.RollingUpgrade, instanceID string, svc autoscalingiface.AutoScalingAPI) error {
+func (r *RollingUpgradeReconciler) TerminateNode(ruObj *upgrademgrv1alpha1.RollingUpgrade, instanceID string) error {
 
 	input := &autoscaling.TerminateInstanceInAutoScalingGroupInput{
 		InstanceId:                     aws.String(instanceID),
 		ShouldDecrementDesiredCapacity: aws.Bool(false),
 	}
 
-	result, err := svc.TerminateInstanceInAutoScalingGroup(input)
+	result, err := r.ASGClient.TerminateInstanceInAutoScalingGroup(input)
 	if err != nil {
 		if aerr, ok := err.(awserr.Error); ok {
 			if strings.Contains(aerr.Message(), "not found") {
@@ -290,6 +295,7 @@ func (r *RollingUpgradeReconciler) TerminateNode(ruObj *upgrademgrv1alpha1.Rolli
 
 	log.Printf("%s: Termination output: %v\n", ruObj.Name, result)
 
+	log.Printf("starting post termination sleep for %v seconds", ruObj.Spec.NodeIntervalSeconds)
 	time.Sleep(time.Duration(ruObj.Spec.NodeIntervalSeconds) * time.Second)
 	if ruObj.Spec.PostTerminate.Script != "" {
 		out, err := runScript(ruObj.Spec.PostTerminate.Script, false, ruObj.Name)
@@ -335,7 +341,7 @@ func (r *RollingUpgradeReconciler) getNodeFromAsg(i *autoscaling.Instance, nodeL
 	return nil
 }
 
-func (r *RollingUpgradeReconciler) populateAsg(ruObj *upgrademgrv1alpha1.RollingUpgrade, svc autoscalingiface.AutoScalingAPI) error {
+func (r *RollingUpgradeReconciler) populateAsg(ruObj *upgrademgrv1alpha1.RollingUpgrade) error {
 	// Initialize a session in the given region that the SDK will use to load
 	// credentials.
 	input := &autoscaling.DescribeAutoScalingGroupsInput{
@@ -344,7 +350,7 @@ func (r *RollingUpgradeReconciler) populateAsg(ruObj *upgrademgrv1alpha1.Rolling
 		},
 	}
 
-	result, err := svc.DescribeAutoScalingGroups(input)
+	result, err := r.ASGClient.DescribeAutoScalingGroups(input)
 	if err != nil {
 		log.Println(err.Error())
 		return errors.New(ruObj.Name + ": Failed to describe autoscaling group: " + err.Error())
@@ -399,7 +405,21 @@ func loadEnvironmentVariables(ruObj *upgrademgrv1alpha1.RollingUpgrade, nodeInst
 	return nil
 }
 
-func (r *RollingUpgradeReconciler) runRestack(ctx *context.Context, ruObj *upgrademgrv1alpha1.RollingUpgrade, svc autoscalingiface.AutoScalingAPI, KubeCtlCall string) (int, error) {
+func (r *RollingUpgradeReconciler) getInProgressInstances(instances []*autoscaling.Instance) ([]*autoscaling.Instance, error) {
+	var inProgressInstances []*autoscaling.Instance
+	taggedInstances, err := getTaggedInstances(EC2StateTagKey, "in-progress", r.EC2Client)
+	if err != nil {
+		return inProgressInstances, err
+	}
+	for _, instance := range instances {
+		if contains(taggedInstances, aws.StringValue(instance.InstanceId)) {
+			inProgressInstances = append(inProgressInstances, instance)
+		}
+	}
+	return inProgressInstances, nil
+}
+
+func (r *RollingUpgradeReconciler) runRestack(ctx *context.Context, ruObj *upgrademgrv1alpha1.RollingUpgrade, KubeCtlCall string) (int, error) {
 	// Setting default values for the Strategy in rollup object
 	r.setDefaultsForRollingUpdateStrategy(ruObj)
 
@@ -427,11 +447,25 @@ func (r *RollingUpgradeReconciler) runRestack(ctx *context.Context, ruObj *upgra
 	r.ClusterState.initializeAsg(*asg.AutoScalingGroupName, asg.Instances)
 
 	currentLaunchConfigName := aws.StringValue(asg.LaunchConfigurationName)
-
 	processedInstances := 0
+
+	inProgress, err := r.getInProgressInstances(asg.Instances)
+	if err != nil {
+		log.Errorf("Failed to acquire in-progress instances, %v", err)
+	}
+
 	for processedInstances < totalNodes {
-		// Fetch instances to update from node selector
-		instances := nodeSelector.SelectNodesForRestack(r.ClusterState)
+		instances := []*autoscaling.Instance{}
+		if len(inProgress) == 0 {
+			// Fetch instances to update from node selector
+			instances = nodeSelector.SelectNodesForRestack(r.ClusterState)
+			log.Printf("selected instances for rotation: %+v", instances)
+		} else {
+			// Prefer in progress instances over new ones
+			instances = inProgress
+			inProgress = []*autoscaling.Instance{}
+			log.Printf("found in progress instances: %+v", instances)
+		}
 
 		if instances == nil {
 			errorMessage := fmt.Sprintf(
@@ -445,7 +479,7 @@ func (r *RollingUpgradeReconciler) runRestack(ctx *context.Context, ruObj *upgra
 		}
 
 		// update the instances
-		err := r.UpdateInstances(ctx, ruObj, instances, currentLaunchConfigName, KubeCtlCall, svc)
+		err := r.UpdateInstances(ctx, ruObj, instances, currentLaunchConfigName, KubeCtlCall)
 		processedInstances += len(instances)
 		if err != nil {
 			return processedInstances, err
@@ -517,8 +551,10 @@ func (r *RollingUpgradeReconciler) Process(ctx *context.Context,
 	if err != nil {
 		log.Fatalf("failed to create asg client, %v", err)
 	}
-	svc := autoscaling.New(sess)
-	err = r.populateAsg(ruObj, svc)
+	r.ASGClient = autoscaling.New(sess)
+	r.EC2Client = ec2.New(sess)
+
+	err = r.populateAsg(ruObj)
 	if err != nil {
 		return r.finishExecution(StatusError, 0, ctx, ruObj)
 	}
@@ -545,7 +581,7 @@ func (r *RollingUpgradeReconciler) Process(ctx *context.Context,
 	r.Update(*ctx, ruObj)
 
 	// Run the restack that acutally performs the rolling update.
-	nodesProcessed, err := r.runRestack(ctx, ruObj, svc, KubeCtlBinary)
+	nodesProcessed, err := r.runRestack(ctx, ruObj, KubeCtlBinary)
 	if err != nil {
 		return r.finishExecution(StatusError, nodesProcessed, ctx, ruObj)
 	}
@@ -627,6 +663,15 @@ func (r *RollingUpgradeReconciler) SetupWithManager(mgr ctrl.Manager) error {
 		Complete(r)
 }
 
+func (r *RollingUpgradeReconciler) setStateTag(instanceID string, state string) error {
+	log.Printf("setting instance %v state to %v", instanceID, state)
+	err := tagEC2instance(instanceID, EC2StateTagKey, state, r.EC2Client)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
 // validateRollingUpgradeObj validates rollup object for the type, maxUnavailable and drainTimeout
 func (r *RollingUpgradeReconciler) validateRollingUpgradeObj(ruObj *upgrademgrv1alpha1.RollingUpgrade) error {
 	strategy := ruObj.Spec.Strategy
@@ -710,8 +755,7 @@ func (r *RollingUpgradeReconciler) UpdateInstances(ctx *context.Context,
 	ruObj *upgrademgrv1alpha1.RollingUpgrade,
 	instances []*autoscaling.Instance,
 	currentLaunchConfigName string,
-	KubeCtlCall string,
-	svc autoscalingiface.AutoScalingAPI) error {
+	KubeCtlCall string) error {
 
 	totalNodes := len(instances)
 	if totalNodes == 0 {
@@ -721,7 +765,7 @@ func (r *RollingUpgradeReconciler) UpdateInstances(ctx *context.Context,
 	ch := make(chan error)
 
 	for _, instance := range instances {
-		go r.UpdateInstance(ctx, ruObj, instance, currentLaunchConfigName, KubeCtlCall, svc, ch)
+		go r.UpdateInstance(ctx, ruObj, instance, currentLaunchConfigName, KubeCtlCall, ch)
 	}
 
 	// wait for upgrades to complete
@@ -754,7 +798,6 @@ func (r *RollingUpgradeReconciler) UpdateInstance(ctx *context.Context,
 	i *autoscaling.Instance,
 	currentLaunchConfigName string,
 	KubeCtlCall string,
-	svc autoscalingiface.AutoScalingAPI,
 	ch chan error) {
 
 	targetLaunchConfigName := aws.StringValue(i.LaunchConfigurationName)
@@ -786,6 +829,13 @@ func (r *RollingUpgradeReconciler) UpdateInstance(ctx *context.Context,
 		return
 	}
 
+	// set the EC2 tag indicating the state to in-progress
+	err = r.setStateTag(targetInstanceID, "in-progress")
+	if err != nil {
+		ch <- err
+		return
+	}
+
 	// Drain and wait for draining node.
 	err = r.DrainNode(ruObj, nodeName, KubeCtlCall, ruObj.Spec.Strategy.DrainTimeout)
 	if err != nil {
@@ -794,7 +844,7 @@ func (r *RollingUpgradeReconciler) UpdateInstance(ctx *context.Context,
 	}
 
 	// Terminate instance.
-	err = r.TerminateNode(ruObj, targetInstanceID, svc)
+	err = r.TerminateNode(ruObj, targetInstanceID)
 	if err != nil {
 		ch <- err
 		return
@@ -810,6 +860,13 @@ func (r *RollingUpgradeReconciler) UpdateInstance(ctx *context.Context,
 		log.Warnf("termination waiter completed but %s is still joined, will proceed with upgrade", nodeName)
 	}
 
+	// set the EC2 tag indicating the state to completed
+	err = r.setStateTag(targetInstanceID, "completed")
+	if err != nil {
+		ch <- err
+		return
+	}
+
 	ruObj.Status.NodesProcessed = ruObj.Status.NodesProcessed + 1
 	r.Update(*ctx, ruObj)
 


Test patch:
diff --git a/controllers/rollingupgrade_controller_test.go b/controllers/rollingupgrade_controller_test.go
index bcb84f6..cb26e38 100644
--- a/controllers/rollingupgrade_controller_test.go
+++ b/controllers/rollingupgrade_controller_test.go
@@ -18,9 +18,11 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/service/autoscaling/autoscalingiface"
+	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
 	"github.com/pkg/errors"
 
 	"github.com/aws/aws-sdk-go/service/autoscaling"
+	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/onsi/gomega"
 	"golang.org/x/net/context"
 	corev1 "k8s.io/api/core/v1"
@@ -292,6 +294,11 @@ func TestDrainNodePostDrainFailureToDrain(t *testing.T) {
 	g.Expect(err).To(gomega.Not(gomega.BeNil()))
 }
 
+type MockEC2 struct {
+	ec2iface.EC2API
+	reservations []*ec2.Reservation
+}
+
 type MockAutoscalingGroup struct {
 	autoscalingiface.AutoScalingAPI
 	errorFlag       bool
@@ -299,6 +306,23 @@ type MockAutoscalingGroup struct {
 	errorInstanceId string
 }
 
+func (m MockEC2) CreateTags(input *ec2.CreateTagsInput) (*ec2.CreateTagsOutput, error) {
+	return &ec2.CreateTagsOutput{}, nil
+}
+
+func (m MockEC2) DescribeInstances(input *ec2.DescribeInstancesInput) (*ec2.DescribeInstancesOutput, error) {
+	return &ec2.DescribeInstancesOutput{Reservations: m.reservations}, nil
+}
+
+func (m MockEC2) DescribeInstancesPages(input *ec2.DescribeInstancesInput, callback func(*ec2.DescribeInstancesOutput, bool) bool) error {
+	page, err := m.DescribeInstances(input)
+	if err != nil {
+		return err
+	}
+	callback(page, false)
+	return nil
+}
+
 func (mockAutoscalingGroup MockAutoscalingGroup) TerminateInstanceInAutoScalingGroup(input *autoscaling.TerminateInstanceInAutoScalingGroupInput) (*autoscaling.TerminateInstanceInAutoScalingGroupOutput, error) {
 	output := &autoscaling.TerminateInstanceInAutoScalingGroupOutput{}
 	if mockAutoscalingGroup.errorFlag {
@@ -314,14 +338,57 @@ func (mockAutoscalingGroup MockAutoscalingGroup) TerminateInstanceInAutoScalingG
 	return output, nil
 }
 
+func TestGetInProgressInstances(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+	mockInstances := []*autoscaling.Instance{
+		{
+			InstanceId: aws.String("i-0123foo"),
+		},
+		{
+			InstanceId: aws.String("i-0123bar"),
+		},
+	}
+	expectedInstance := &autoscaling.Instance{
+		InstanceId: aws.String("i-0123foo"),
+	}
+	mockReservations := []*ec2.Reservation{
+		{
+			Instances: []*ec2.Instance{
+				{
+					InstanceId: aws.String("i-0123foo"),
+				},
+			},
+		},
+	}
+	reconciler := &RollingUpgradeReconciler{
+		ClusterState: NewClusterState(),
+		EC2Client:    MockEC2{reservations: mockReservations},
+		ASGClient: MockAutoscalingGroup{
+			errorFlag: false,
+			awsErr:    nil,
+		},
+	}
+	inProgressInstances, err := reconciler.getInProgressInstances(mockInstances)
+	g.Expect(err).To(gomega.BeNil())
+	g.Expect(inProgressInstances).To(gomega.ContainElement(expectedInstance))
+	g.Expect(inProgressInstances).To(gomega.HaveLen(1))
+}
+
 func TestTerminateNodeSuccess(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 	mockNode := "some-node-id"
 
 	ruObj := &upgrademgrv1alpha1.RollingUpgrade{ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"}}
-	rcRollingUpgrade := &RollingUpgradeReconciler{ClusterState: NewClusterState()}
-	mockAutoscalingGroup := MockAutoscalingGroup{errorFlag: false, awsErr: nil}
-	err := rcRollingUpgrade.TerminateNode(ruObj, mockNode, mockAutoscalingGroup)
+	rcRollingUpgrade := &RollingUpgradeReconciler{
+		ClusterState: NewClusterState(),
+		EC2Client:    MockEC2{},
+		ASGClient: MockAutoscalingGroup{
+			errorFlag: false,
+			awsErr:    nil,
+		},
+	}
+
+	err := rcRollingUpgrade.TerminateNode(ruObj, mockNode)
 	g.Expect(err).To(gomega.BeNil())
 }
 
@@ -330,12 +397,16 @@ func TestTerminateNodeErrorNotFound(t *testing.T) {
 	mockNode := "some-node-id"
 
 	ruObj := &upgrademgrv1alpha1.RollingUpgrade{ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"}}
-	rcRollingUpgrade := &RollingUpgradeReconciler{ClusterState: NewClusterState()}
 	mockAutoscalingGroup := MockAutoscalingGroup{errorFlag: true, awsErr: awserr.New("InvalidInstanceID.NotFound",
 		"ValidationError: Instance Id not found - No managed instance found for instance ID i-0bba",
 		nil)}
+	rcRollingUpgrade := &RollingUpgradeReconciler{
+		ClusterState: NewClusterState(),
+		ASGClient:    mockAutoscalingGroup,
+		EC2Client:    MockEC2{},
+	}
 
-	err := rcRollingUpgrade.TerminateNode(ruObj, mockNode, mockAutoscalingGroup)
+	err := rcRollingUpgrade.TerminateNode(ruObj, mockNode)
 	g.Expect(err).To(gomega.BeNil())
 }
 
@@ -344,12 +415,16 @@ func TestTerminateNodeErrorOtherError(t *testing.T) {
 	mockNode := "some-node-id"
 
 	ruObj := &upgrademgrv1alpha1.RollingUpgrade{ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"}}
-	rcRollingUpgrade := &RollingUpgradeReconciler{ClusterState: NewClusterState()}
 	mockAutoscalingGroup := MockAutoscalingGroup{errorFlag: true, awsErr: awserr.New("some-other-aws-error",
 		"some message",
 		errors.New("some error"))}
 
-	err := rcRollingUpgrade.TerminateNode(ruObj, mockNode, mockAutoscalingGroup)
+	rcRollingUpgrade := &RollingUpgradeReconciler{
+		ClusterState: NewClusterState(),
+		ASGClient:    mockAutoscalingGroup,
+		EC2Client:    MockEC2{},
+	}
+	err := rcRollingUpgrade.TerminateNode(ruObj, mockNode)
 	g.Expect(err.Error()).To(gomega.ContainSubstring("some error"))
 }
 
@@ -359,10 +434,13 @@ func TestTerminateNodePostTerminateScriptSuccess(t *testing.T) {
 
 	ruObj := &upgrademgrv1alpha1.RollingUpgrade{ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"}}
 	ruObj.Spec.PostTerminate.Script = "echo hello!"
-	rcRollingUpgrade := &RollingUpgradeReconciler{ClusterState: NewClusterState()}
 	mockAutoscalingGroup := MockAutoscalingGroup{errorFlag: false, awsErr: nil}
-
-	err := rcRollingUpgrade.TerminateNode(ruObj, mockNode, mockAutoscalingGroup)
+	rcRollingUpgrade := &RollingUpgradeReconciler{
+		ClusterState: NewClusterState(),
+		ASGClient:    mockAutoscalingGroup,
+		EC2Client:    MockEC2{},
+	}
+	err := rcRollingUpgrade.TerminateNode(ruObj, mockNode)
 	g.Expect(err).To(gomega.BeNil())
 }
 
@@ -372,10 +450,13 @@ func TestTerminateNodePostTerminateScriptErrorNotFoundFromServer(t *testing.T) {
 
 	ruObj := &upgrademgrv1alpha1.RollingUpgrade{ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"}}
 	ruObj.Spec.PostTerminate.Script = "echo 'Error from server (NotFound)'; exit 1"
-	rcRollingUpgrade := &RollingUpgradeReconciler{ClusterState: NewClusterState()}
 	mockAutoscalingGroup := MockAutoscalingGroup{errorFlag: false, awsErr: nil}
-
-	err := rcRollingUpgrade.TerminateNode(ruObj, mockNode, mockAutoscalingGroup)
+	rcRollingUpgrade := &RollingUpgradeReconciler{
+		ClusterState: NewClusterState(),
+		ASGClient:    mockAutoscalingGroup,
+		EC2Client:    MockEC2{},
+	}
+	err := rcRollingUpgrade.TerminateNode(ruObj, mockNode)
 	g.Expect(err).To(gomega.BeNil())
 }
 
@@ -385,10 +466,13 @@ func TestTerminateNodePostTerminateScriptErrorOtherError(t *testing.T) {
 
 	ruObj := &upgrademgrv1alpha1.RollingUpgrade{ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"}}
 	ruObj.Spec.PostTerminate.Script = "exit 1"
-	rcRollingUpgrade := &RollingUpgradeReconciler{ClusterState: NewClusterState()}
 	mockAutoscalingGroup := MockAutoscalingGroup{errorFlag: false, awsErr: nil}
-
-	err := rcRollingUpgrade.TerminateNode(ruObj, mockNode, mockAutoscalingGroup)
+	rcRollingUpgrade := &RollingUpgradeReconciler{
+		ClusterState: NewClusterState(),
+		ASGClient:    mockAutoscalingGroup,
+		EC2Client:    MockEC2{},
+	}
+	err := rcRollingUpgrade.TerminateNode(ruObj, mockNode)
 	g.Expect(err).To(gomega.Not(gomega.BeNil()))
 	g.Expect(err.Error()).To(gomega.HavePrefix("Failed to run postTerminate script: "))
 }
@@ -537,9 +621,12 @@ func TestPopulateAsgSuccess(t *testing.T) {
 		TypeMeta: metav1.TypeMeta{Kind: "RollingUpgrade", APIVersion: "v1alpha1"},
 		Spec:     upgrademgrv1alpha1.RollingUpgradeSpec{AsgName: "correct-asg"}}
 
-	mockAsgAPI := &MockAutoScalingAPI{}
-	rcRollingUpgrade := &RollingUpgradeReconciler{ClusterState: NewClusterState()}
-	err := rcRollingUpgrade.populateAsg(ruObj, mockAsgAPI)
+	rcRollingUpgrade := &RollingUpgradeReconciler{
+		ClusterState: NewClusterState(),
+		ASGClient:    &MockAutoScalingAPI{},
+		EC2Client:    MockEC2{},
+	}
+	err := rcRollingUpgrade.populateAsg(ruObj)
 
 	g.Expect(err).To(gomega.BeNil())
 
@@ -558,9 +645,12 @@ func TestPopulateAsgTooMany(t *testing.T) {
 		TypeMeta: metav1.TypeMeta{Kind: "RollingUpgrade", APIVersion: "v1alpha1"},
 		Spec:     upgrademgrv1alpha1.RollingUpgradeSpec{AsgName: "too-many"}}
 
-	mockAsgAPI := &MockAutoScalingAPI{}
-	rcRollingUpgrade := &RollingUpgradeReconciler{ClusterState: NewClusterState()}
-	err := rcRollingUpgrade.populateAsg(ruObj, mockAsgAPI)
+	rcRollingUpgrade := &RollingUpgradeReconciler{
+		ClusterState: NewClusterState(),
+		ASGClient:    &MockAutoScalingAPI{},
+		EC2Client:    MockEC2{},
+	}
+	err := rcRollingUpgrade.populateAsg(ruObj)
 
 	g.Expect(err).To(gomega.Not(gomega.BeNil()))
 	g.Expect(err.Error()).To(gomega.Equal("Too many asgs"))
@@ -573,9 +663,12 @@ func TestPopulateAsgNone(t *testing.T) {
 		TypeMeta: metav1.TypeMeta{Kind: "RollingUpgrade", APIVersion: "v1alpha1"},
 		Spec:     upgrademgrv1alpha1.RollingUpgradeSpec{AsgName: "no-asg-at-all"}}
 
-	mockAsgAPI := &MockAutoScalingAPI{}
-	rcRollingUpgrade := &RollingUpgradeReconciler{ClusterState: NewClusterState()}
-	err := rcRollingUpgrade.populateAsg(ruObj, mockAsgAPI)
+	rcRollingUpgrade := &RollingUpgradeReconciler{
+		ClusterState: NewClusterState(),
+		ASGClient:    &MockAutoScalingAPI{},
+		EC2Client:    MockEC2{},
+	}
+	err := rcRollingUpgrade.populateAsg(ruObj)
 
 	g.Expect(err).To(gomega.Not(gomega.BeNil()))
 	g.Expect(err.Error()).To(gomega.Equal("No ASG found"))
@@ -715,7 +808,6 @@ func TestRunRestackSuccessOneNode(t *testing.T) {
 		ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"},
 		Spec:       upgrademgrv1alpha1.RollingUpgradeSpec{AsgName: someAsg},
 	}
-	mockAutoscalingGroup := MockAutoscalingGroup{}
 
 	mgr, err := manager.New(cfg, manager.Options{})
 	g.Expect(err).NotTo(gomega.HaveOccurred())
@@ -730,6 +822,8 @@ func TestRunRestackSuccessOneNode(t *testing.T) {
 	nodeList := corev1.NodeList{Items: []corev1.Node{fooNode1, fooNode2, correctNode}}
 	rcRollingUpgrade := &RollingUpgradeReconciler{
 		Client:          mgr.GetClient(),
+		ASGClient:       MockAutoscalingGroup{},
+		EC2Client:       MockEC2{},
 		generatedClient: kubernetes.NewForConfigOrDie(mgr.GetConfig()),
 		admissionMap:    sync.Map{},
 		ruObjNameToASG:  sync.Map{},
@@ -740,7 +834,7 @@ func TestRunRestackSuccessOneNode(t *testing.T) {
 
 	ctx := context.TODO()
 
-	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, mockAutoscalingGroup, "exit 0;")
+	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, "exit 0;")
 	g.Expect(nodesProcessed).To(gomega.Equal(1))
 	g.Expect(err).To(gomega.BeNil())
 }
@@ -762,7 +856,6 @@ func TestRunRestackSuccessMultipleNodes(t *testing.T) {
 
 	ruObj := &upgrademgrv1alpha1.RollingUpgrade{ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"},
 		Spec: upgrademgrv1alpha1.RollingUpgradeSpec{AsgName: someAsg}}
-	mockAutoscalingGroup := MockAutoscalingGroup{}
 
 	mgr, err := manager.New(cfg, manager.Options{})
 	g.Expect(err).NotTo(gomega.HaveOccurred())
@@ -778,6 +871,8 @@ func TestRunRestackSuccessMultipleNodes(t *testing.T) {
 	nodeList := corev1.NodeList{Items: []corev1.Node{fooNode1, fooNode2, correctNode, correctNode2}}
 	rcRollingUpgrade := &RollingUpgradeReconciler{
 		Client:          mgr.GetClient(),
+		ASGClient:       MockAutoscalingGroup{},
+		EC2Client:       MockEC2{},
 		generatedClient: kubernetes.NewForConfigOrDie(mgr.GetConfig()),
 		admissionMap:    sync.Map{},
 		ruObjNameToASG:  sync.Map{},
@@ -788,7 +883,7 @@ func TestRunRestackSuccessMultipleNodes(t *testing.T) {
 
 	ctx := context.TODO()
 
-	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, mockAutoscalingGroup, "exit 0;")
+	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, "exit 0;")
 	g.Expect(nodesProcessed).To(gomega.Equal(2))
 	g.Expect(err).To(gomega.BeNil())
 }
@@ -807,7 +902,6 @@ func TestRunRestackSameLaunchConfig(t *testing.T) {
 
 	ruObj := &upgrademgrv1alpha1.RollingUpgrade{ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"},
 		Spec: upgrademgrv1alpha1.RollingUpgradeSpec{AsgName: someAsg}}
-	mockAutoscalingGroup := MockAutoscalingGroup{}
 
 	mgr, err := manager.New(cfg, manager.Options{})
 	g.Expect(err).NotTo(gomega.HaveOccurred())
@@ -815,6 +909,8 @@ func TestRunRestackSameLaunchConfig(t *testing.T) {
 
 	rcRollingUpgrade := &RollingUpgradeReconciler{
 		Client:          mgr.GetClient(),
+		ASGClient:       MockAutoscalingGroup{},
+		EC2Client:       MockEC2{},
 		generatedClient: kubernetes.NewForConfigOrDie(mgr.GetConfig()),
 		admissionMap:    sync.Map{},
 		ruObjNameToASG:  sync.Map{},
@@ -825,7 +921,7 @@ func TestRunRestackSameLaunchConfig(t *testing.T) {
 	ctx := context.TODO()
 
 	// This execution should not perform drain or termination, but should pass
-	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, mockAutoscalingGroup, KubeCtlBinary)
+	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, KubeCtlBinary)
 	g.Expect(nodesProcessed).To(gomega.Equal(1))
 	g.Expect(err).To(gomega.BeNil())
 }
@@ -834,12 +930,15 @@ func TestRunRestackRollingUpgradeNotInMap(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 
 	ruObj := &upgrademgrv1alpha1.RollingUpgrade{ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"}}
-	mockAutoscalingGroup := MockAutoscalingGroup{}
-	rcRollingUpgrade := &RollingUpgradeReconciler{ClusterState: NewClusterState()}
+	rcRollingUpgrade := &RollingUpgradeReconciler{
+		ClusterState: NewClusterState(),
+		ASGClient:    MockAutoscalingGroup{},
+		EC2Client:    MockEC2{},
+	}
 	ctx := context.TODO()
 
 	g.Expect(rcRollingUpgrade.ruObjNameToASG.Load(ruObj.Name)).To(gomega.BeNil())
-	int, err := rcRollingUpgrade.runRestack(&ctx, ruObj, mockAutoscalingGroup, KubeCtlBinary)
+	int, err := rcRollingUpgrade.runRestack(&ctx, ruObj, KubeCtlBinary)
 	g.Expect(int).To(gomega.Equal(0))
 	g.Expect(err).To(gomega.Not(gomega.BeNil()))
 	g.Expect(err.Error()).To(gomega.HavePrefix("Failed to find rollingUpgrade name in map."))
@@ -860,7 +959,6 @@ func TestRunRestackRollingUpgradeNodeNameNotFound(t *testing.T) {
 
 	ruObj := &upgrademgrv1alpha1.RollingUpgrade{ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"},
 		Spec: upgrademgrv1alpha1.RollingUpgradeSpec{AsgName: someAsg}}
-	mockAutoscalingGroup := MockAutoscalingGroup{}
 
 	mgr, err := manager.New(cfg, manager.Options{})
 	g.Expect(err).NotTo(gomega.HaveOccurred())
@@ -869,6 +967,8 @@ func TestRunRestackRollingUpgradeNodeNameNotFound(t *testing.T) {
 	emptyNodeList := corev1.NodeList{}
 	rcRollingUpgrade := &RollingUpgradeReconciler{
 		Client:          mgr.GetClient(),
+		ASGClient:       MockAutoscalingGroup{},
+		EC2Client:       MockEC2{},
 		generatedClient: kubernetes.NewForConfigOrDie(mgr.GetConfig()),
 		admissionMap:    sync.Map{},
 		ruObjNameToASG:  sync.Map{},
@@ -880,7 +980,7 @@ func TestRunRestackRollingUpgradeNodeNameNotFound(t *testing.T) {
 	ctx := context.TODO()
 
 	// This execution gets past the different launch config check, but fails to be found at the node level
-	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, mockAutoscalingGroup, KubeCtlBinary)
+	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, KubeCtlBinary)
 	g.Expect(nodesProcessed).To(gomega.Equal(1))
 	g.Expect(err).To(gomega.BeNil())
 }
@@ -900,7 +1000,6 @@ func TestRunRestackNoNodeName(t *testing.T) {
 
 	ruObj := &upgrademgrv1alpha1.RollingUpgrade{ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"},
 		Spec: upgrademgrv1alpha1.RollingUpgradeSpec{AsgName: someAsg}}
-	mockAutoscalingGroup := MockAutoscalingGroup{}
 
 	mgr, err := manager.New(cfg, manager.Options{})
 	g.Expect(err).NotTo(gomega.HaveOccurred())
@@ -914,6 +1013,8 @@ func TestRunRestackNoNodeName(t *testing.T) {
 	nodeList := corev1.NodeList{Items: []corev1.Node{fooNode1, fooNode2, correctNode}}
 	rcRollingUpgrade := &RollingUpgradeReconciler{
 		Client:          mgr.GetClient(),
+		ASGClient:       MockAutoscalingGroup{},
+		EC2Client:       MockEC2{},
 		generatedClient: kubernetes.NewForConfigOrDie(mgr.GetConfig()),
 		admissionMap:    sync.Map{},
 		ruObjNameToASG:  sync.Map{},
@@ -925,7 +1026,7 @@ func TestRunRestackNoNodeName(t *testing.T) {
 	ctx := context.TODO()
 
 	// This execution gets past the different launch config check, but since there is no node name, it is skipped
-	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, mockAutoscalingGroup, KubeCtlBinary)
+	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, KubeCtlBinary)
 	g.Expect(nodesProcessed).To(gomega.Equal(1))
 	g.Expect(err).To(gomega.BeNil())
 }
@@ -958,7 +1059,6 @@ func TestRunRestackDrainNodeFail(t *testing.T) {
 			PreDrain: somePreDrain,
 		},
 	}
-	mockAutoscalingGroup := MockAutoscalingGroup{}
 
 	mgr, err := manager.New(cfg, manager.Options{})
 	g.Expect(err).NotTo(gomega.HaveOccurred())
@@ -973,6 +1073,8 @@ func TestRunRestackDrainNodeFail(t *testing.T) {
 	nodeList := corev1.NodeList{Items: []corev1.Node{fooNode1, fooNode2, correctNode}}
 	rcRollingUpgrade := &RollingUpgradeReconciler{
 		Client:          mgr.GetClient(),
+		ASGClient:       MockAutoscalingGroup{},
+		EC2Client:       MockEC2{},
 		generatedClient: kubernetes.NewForConfigOrDie(mgr.GetConfig()),
 		admissionMap:    sync.Map{},
 		ruObjNameToASG:  sync.Map{},
@@ -984,7 +1086,7 @@ func TestRunRestackDrainNodeFail(t *testing.T) {
 	ctx := context.TODO()
 
 	// This execution gets past the different launch config check, but fails to drain the node because of a predrain failing script
-	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, mockAutoscalingGroup, KubeCtlBinary)
+	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, KubeCtlBinary)
 	g.Expect(nodesProcessed).To(gomega.Equal(1))
 	g.Expect(err.Error()).To(gomega.HavePrefix("Error updating instances, ErrorCount: 1, Errors: ["))
 }
@@ -1022,6 +1124,8 @@ func TestRunRestackTerminateNodeFail(t *testing.T) {
 	nodeList := corev1.NodeList{Items: []corev1.Node{fooNode1, fooNode2, correctNode}}
 	rcRollingUpgrade := &RollingUpgradeReconciler{
 		Client:          mgr.GetClient(),
+		ASGClient:       mockAutoscalingGroup,
+		EC2Client:       MockEC2{},
 		generatedClient: kubernetes.NewForConfigOrDie(mgr.GetConfig()),
 		admissionMap:    sync.Map{},
 		ruObjNameToASG:  sync.Map{},
@@ -1033,7 +1137,7 @@ func TestRunRestackTerminateNodeFail(t *testing.T) {
 	ctx := context.TODO()
 
 	// This execution gets past the different launch config check, but fails to terminate node
-	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, mockAutoscalingGroup, "exit 0;")
+	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, "exit 0;")
 	g.Expect(nodesProcessed).To(gomega.Equal(1))
 	g.Expect(err.Error()).To(gomega.HavePrefix("Error updating instances, ErrorCount: 1, Errors: ["))
 	g.Expect(err.Error()).To(gomega.ContainSubstring("some error"))
@@ -1077,7 +1181,6 @@ func TestUniformAcrossAzUpdateSuccessMultipleNodes(t *testing.T) {
 			},
 		},
 	}
-	mockAutoscalingGroup := MockAutoscalingGroup{}
 
 	mgr, err := manager.New(cfg, manager.Options{})
 	g.Expect(err).NotTo(gomega.HaveOccurred())
@@ -1112,6 +1215,8 @@ func TestUniformAcrossAzUpdateSuccessMultipleNodes(t *testing.T) {
 	}}
 	rcRollingUpgrade := &RollingUpgradeReconciler{
 		Client:          mgr.GetClient(),
+		ASGClient:       MockAutoscalingGroup{},
+		EC2Client:       MockEC2{},
 		generatedClient: kubernetes.NewForConfigOrDie(mgr.GetConfig()),
 		admissionMap:    sync.Map{},
 		ruObjNameToASG:  sync.Map{},
@@ -1122,7 +1227,7 @@ func TestUniformAcrossAzUpdateSuccessMultipleNodes(t *testing.T) {
 
 	ctx := context.TODO()
 
-	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, mockAutoscalingGroup, "exit 0;")
+	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, "exit 0;")
 	g.Expect(nodesProcessed).To(gomega.Equal(9))
 	g.Expect(err).To(gomega.BeNil())
 }
@@ -1144,7 +1249,6 @@ func TestUpdateInstances(t *testing.T) {
 
 	ruObj := &upgrademgrv1alpha1.RollingUpgrade{ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"},
 		Spec: upgrademgrv1alpha1.RollingUpgradeSpec{AsgName: someAsg}}
-	mockAutoscalingGroup := MockAutoscalingGroup{}
 
 	mgr, err := manager.New(cfg, manager.Options{})
 	g.Expect(err).NotTo(gomega.HaveOccurred())
@@ -1160,6 +1264,8 @@ func TestUpdateInstances(t *testing.T) {
 	nodeList := corev1.NodeList{Items: []corev1.Node{fooNode1, fooNode2, correctNode, correctNode2}}
 	rcRollingUpgrade := &RollingUpgradeReconciler{
 		Client:          mgr.GetClient(),
+		ASGClient:       MockAutoscalingGroup{},
+		EC2Client:       MockEC2{},
 		generatedClient: kubernetes.NewForConfigOrDie(mgr.GetConfig()),
 		admissionMap:    sync.Map{},
 		ruObjNameToASG:  sync.Map{},
@@ -1171,7 +1277,7 @@ func TestUpdateInstances(t *testing.T) {
 	ctx := context.TODO()
 
 	err = rcRollingUpgrade.UpdateInstances(&ctx,
-		ruObj, mockAsg.Instances, "A", "exit 0;", mockAutoscalingGroup)
+		ruObj, mockAsg.Instances, "A", "exit 0;")
 	g.Expect(err).ShouldNot(gomega.HaveOccurred())
 }
 
@@ -1212,6 +1318,8 @@ func TestUpdateInstancesError(t *testing.T) {
 	nodeList := corev1.NodeList{Items: []corev1.Node{fooNode1, fooNode2, correctNode, correctNode2}}
 	rcRollingUpgrade := &RollingUpgradeReconciler{
 		Client:          mgr.GetClient(),
+		ASGClient:       mockAutoScalingGroup,
+		EC2Client:       MockEC2{},
 		generatedClient: kubernetes.NewForConfigOrDie(mgr.GetConfig()),
 		admissionMap:    sync.Map{},
 		ruObjNameToASG:  sync.Map{},
@@ -1223,7 +1331,7 @@ func TestUpdateInstancesError(t *testing.T) {
 	ctx := context.TODO()
 
 	err = rcRollingUpgrade.UpdateInstances(&ctx,
-		ruObj, mockAsg.Instances, "A", "exit 0;", mockAutoScalingGroup)
+		ruObj, mockAsg.Instances, "A", "exit 0;")
 	g.Expect(err).Should(gomega.HaveOccurred())
 	g.Expect(err).Should(gomega.BeAssignableToTypeOf(&UpdateInstancesError{}))
 	if updateInstancesError, ok := err.(*UpdateInstancesError); ok {
@@ -1271,6 +1379,8 @@ func TestUpdateInstancesPartialError(t *testing.T) {
 	nodeList := corev1.NodeList{Items: []corev1.Node{fooNode1, fooNode2, correctNode, correctNode2}}
 	rcRollingUpgrade := &RollingUpgradeReconciler{
 		Client:          mgr.GetClient(),
+		ASGClient:       mockAutoScalingGroup,
+		EC2Client:       MockEC2{},
 		generatedClient: kubernetes.NewForConfigOrDie(mgr.GetConfig()),
 		admissionMap:    sync.Map{},
 		ruObjNameToASG:  sync.Map{},
@@ -1282,7 +1392,7 @@ func TestUpdateInstancesPartialError(t *testing.T) {
 	ctx := context.TODO()
 
 	err = rcRollingUpgrade.UpdateInstances(&ctx,
-		ruObj, mockAsg.Instances, "A", "exit 0;", mockAutoScalingGroup)
+		ruObj, mockAsg.Instances, "A", "exit 0;")
 	g.Expect(err).Should(gomega.HaveOccurred())
 	g.Expect(err).Should(gomega.BeAssignableToTypeOf(&UpdateInstancesError{}))
 	if updateInstancesError, ok := err.(*UpdateInstancesError); ok {
@@ -1309,7 +1419,7 @@ func TestUpdateInstancesWithZeroInstances(t *testing.T) {
 	ctx := context.TODO()
 
 	err = rcRollingUpgrade.UpdateInstances(&ctx,
-		nil, nil, "A", "exit 0;", nil)
+		nil, nil, "A", "exit 0;")
 	g.Expect(err).ShouldNot(gomega.HaveOccurred())
 }
 
@@ -1905,7 +2015,6 @@ func TestRunRestackNoNodeInAsg(t *testing.T) {
 			Strategy: upgrademgrv1alpha1.UpdateStrategy{Type: upgrademgrv1alpha1.RandomUpdateStrategy},
 		},
 	}
-	mockAutoscalingGroup := MockAutoscalingGroup{}
 
 	mgr, err := manager.New(cfg, manager.Options{})
 	g.Expect(err).NotTo(gomega.HaveOccurred())
@@ -1914,6 +2023,8 @@ func TestRunRestackNoNodeInAsg(t *testing.T) {
 	nodeList := corev1.NodeList{Items: []corev1.Node{}}
 	rcRollingUpgrade := &RollingUpgradeReconciler{
 		Client:          mgr.GetClient(),
+		ASGClient:       MockAutoscalingGroup{},
+		EC2Client:       MockEC2{},
 		generatedClient: kubernetes.NewForConfigOrDie(mgr.GetConfig()),
 		admissionMap:    sync.Map{},
 		NodeList:        &nodeList,
@@ -1924,7 +2035,7 @@ func TestRunRestackNoNodeInAsg(t *testing.T) {
 	ctx := context.TODO()
 
 	// This execution gets past the different launch config check, but since there is no node name, it is skipped
-	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, mockAutoscalingGroup, KubeCtlBinary)
+	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, KubeCtlBinary)
 	g.Expect(nodesProcessed).To(gomega.Equal(0))
 	g.Expect(err).To(gomega.BeNil())
 }
@@ -1988,7 +2099,6 @@ func TestRunRestackWithNodesLessThanMaxUnavailable(t *testing.T) {
 			},
 		},
 	}
-	mockAutoscalingGroup := MockAutoscalingGroup{}
 
 	mgr, err := manager.New(cfg, manager.Options{})
 	g.Expect(err).NotTo(gomega.HaveOccurred())
@@ -1996,6 +2106,8 @@ func TestRunRestackWithNodesLessThanMaxUnavailable(t *testing.T) {
 
 	rcRollingUpgrade := &RollingUpgradeReconciler{
 		Client:          mgr.GetClient(),
+		ASGClient:       MockAutoscalingGroup{},
+		EC2Client:       MockEC2{},
 		generatedClient: kubernetes.NewForConfigOrDie(mgr.GetConfig()),
 		admissionMap:    sync.Map{},
 		ruObjNameToASG:  sync.Map{},
@@ -2006,7 +2118,7 @@ func TestRunRestackWithNodesLessThanMaxUnavailable(t *testing.T) {
 	ctx := context.TODO()
 
 	// This execution should not perform drain or termination, but should pass
-	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, mockAutoscalingGroup, KubeCtlBinary)
+	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, KubeCtlBinary)
 	g.Expect(err).To(gomega.BeNil())
 	g.Expect(nodesProcessed).To(gomega.Equal(1))
 }