Repository: keikoproj/upgrade-manager
Pull request #27 (2019-11-29 07:03:54)

Added uniformAcrossAzUpdate strategy
- Externalized the node selectors behind an interface
- Refactored the restack logic to leverage node selectors to get a list of instances to restack on each iteration
- Added a strategy for `RandomUpdate`
- Added a new strategy to perform updates across AZ in same proportion. This new strategy picks either same number of nodes or same percentage of nodes from each AZ for updating in each iteration.

- Added unit tests - 80% coverage for controllers
```
go test ./api/... ./controllers/... -coverprofile cover.out
ok      github.com/keikoproj/upgrade-manager/api/v1alpha1       5.145s  coverage: 1.4% of statements
ok      github.com/keikoproj/upgrade-manager/controllers        16.179s coverage: 80.9% of statements
```
- Tested the strategy by performing Random & rolling update in a k8s cluster
  - Logs: [update-manager-logs.txt](https://github.com/keikoproj/upgrade-manager/files/3909772/update-manager-logs.txt)


- Sample `RollingUpgrade` object
```yaml
apiVersion: upgrademgr.keikoproj.io/v1alpha1
kind: RollingUpgrade
metadata:
  name: rollingupgrade-nodes
  namespace: kube-system
spec:
  asgName: nodes
  nodeIntervalSeconds: 300
  postDrain: {}
  postDrainDelaySeconds: 90
  postTerminate: {}
  preDrain: {}
  region: us-west-2
  strategy:
    drainTimeout: 180
    maxUnavailable: 100%
    type: uniformAcrossAzUpdate
```

claassistantio: [![CLA assistant check](https://cla-assistant.io/pull/badge/not_signed)](https://cla-assistant.io/keikoproj/upgrade-manager?pullRequest=27) <br/>Thank you for your submission, we really appreciate it. Like many open source projects, we ask that you sign our [Contributor License Agreement](https://cla-assistant.io/keikoproj/upgrade-manager?pullRequest=27) before we can accept your contribution.<br/><sub>You have signed the CLA already but the status is still pending? Let us [recheck](https://cla-assistant.io/check/keikoproj/upgrade-manager?pullRequest=27) it.</sub>
shrinandj: @narayanan:  This is awesome! Thanks a lot for implementing this feature. Certainly makes upgrades more flexible.

Patch:
diff --git a/README.md b/README.md
index 6f8eef3..7cb9e20 100644
--- a/README.md
+++ b/README.md
@@ -19,7 +19,9 @@ RollingUpgrade provides a Kubernetes native mechanism for doing rolling-updates
   - `postDrain.postWaitScript`: The script to run after the node is drained and the waitSeconds have passed. This can be used for ensuring that the drained pods actually were able to start elsewhere.
   - `nodeIntervalSeconds`: The amount of time in seconds to wait after each node in the ASG is terminated.
   - `postTerminate.script`: Optional bash script to execute after the node has terminated.
-  - `strategy.type`: This field is optional and currently only random update strategy is supported. Refer to [random_update_strategy.yaml](examples/random_update_strategy.yaml) for sample custom resource definition.
+  - `strategy.type`: This field is optional and currently two strategies are supported
+    - `randomUpdate` - Default is type is not specified. Picks nodes randomly for updating. Refer to [random_update_strategy.yaml](examples/random_update_strategy.yaml) for sample custom resource definition.
+    - `uniformAcrossAzUpdate` - Picks same number of nodes or same percentage of nodes from each AZ for update. Refer to [uniform_across_az_update_strategy.yaml](examples/uniform_across_az_update_strategy.yaml) for sample custom resource definition.
   - `strategy.maxUnavailable`: Optional field. The number of nodes that can be unavailable during rolling upgrade, can be specified as number of nodes or the percent of total number of nodes. Default is "1".
   - `strategy.drainTimeout`: Optional field. Node will be terminated after drain timeout even if `kubectl drain` has not been completed and value has to be specified in seconds. Default is -1.
 
diff --git a/api/v1alpha1/rollingupgrade_types.go b/api/v1alpha1/rollingupgrade_types.go
index c85c572..4fb4468 100644
--- a/api/v1alpha1/rollingupgrade_types.go
+++ b/api/v1alpha1/rollingupgrade_types.go
@@ -95,6 +95,9 @@ const (
 	// RandomUpdate strategy treats all the Azs as a single unit and picks random nodes for update
 	RandomUpdateStrategy UpdateStrategyType = "randomUpdate"
 
+	// RandomUpdate strategy treats all the Azs as a single unit and picks random nodes for update
+	UniformAcrossAzUpdateStrategy UpdateStrategyType = "uniformAcrossAzUpdate"
+
 	// Other update strategies such as rolling update by Az or rolling update with a predifined instance list
 	// can be implemented in future by adding more update strategy types
 )
diff --git a/controllers/helpers.go b/controllers/helpers.go
new file mode 100644
index 0000000..2e45dae
--- /dev/null
+++ b/controllers/helpers.go
@@ -0,0 +1,72 @@
+package controllers
+
+import (
+	"github.com/aws/aws-sdk-go/service/autoscaling"
+	upgrademgrv1alpha1 "github.com/keikoproj/upgrade-manager/api/v1alpha1"
+	"log"
+	"strconv"
+	"strings"
+)
+
+// getMaxUnavailable calculates and returns the maximum unavailable nodes
+// takes an update strategy and total number of nodes as input
+func getMaxUnavailable(strategy upgrademgrv1alpha1.UpdateStrategy, totalNodes int) int {
+	maxUnavailable := 1
+	if strategy.MaxUnavailable.Type == 0 {
+		maxUnavailable = int(strategy.MaxUnavailable.IntVal)
+	} else if strategy.MaxUnavailable.Type == 1 {
+		strVallue := strategy.MaxUnavailable.StrVal
+		intValue, _ := strconv.Atoi(strings.Trim(strVallue, "%"))
+		maxUnavailable = int(float32(intValue) / float32(100) * float32(totalNodes))
+	}
+	// setting maxUnavailable to total number of nodes when maxUnavailable is greater than total node count
+	if totalNodes < maxUnavailable {
+		log.Printf("Reducing maxUnavailable count from %d to %d as total nodes count is %d",
+			maxUnavailable, totalNodes, totalNodes)
+		maxUnavailable = totalNodes
+	}
+	// maxUnavailable has to be atleast 1 when there are nodes in the ASG
+	if totalNodes > 0 && maxUnavailable < 1 {
+		maxUnavailable = 1
+	}
+	return maxUnavailable
+}
+
+// getNextAvailableInstances checks the cluster state store for the instance state
+// and returns the next set of instances available for update
+func getNextAvailableInstances(
+	asgName string,
+	numberOfInstances int,
+	instances []*autoscaling.Instance,
+	state ClusterState) []*autoscaling.Instance {
+	return getNextSetOfAvailableInstancesInAz(asgName, "", numberOfInstances, instances, state)
+}
+
+// getNextSetOfAvailableInstancesInAz checks the cluster state store for the instance state
+// and returns the next set of instances available for update in the given AX
+func getNextSetOfAvailableInstancesInAz(
+	asgName string,
+	azName string,
+	numberOfInstances int,
+	instances []*autoscaling.Instance,
+	state ClusterState,
+) []*autoscaling.Instance {
+
+	var instancesForUpdate []*autoscaling.Instance
+	for instancesFound := 0; instancesFound < numberOfInstances; {
+		instanceId := state.getNextAvailableInstanceIdInAz(asgName, azName)
+		if len(instanceId) == 0 {
+			// All instances are updated, no more instance to update in this AZ
+			break
+		}
+
+		// check if the instance picked is part of ASG
+		for _, instance := range instances {
+			if *instance.InstanceId == instanceId {
+				instancesForUpdate = append(instancesForUpdate, instance)
+				instancesFound++
+			}
+		}
+	}
+	return instancesForUpdate
+}
diff --git a/controllers/node_selector.go b/controllers/node_selector.go
new file mode 100644
index 0000000..8e4b93b
--- /dev/null
+++ b/controllers/node_selector.go
@@ -0,0 +1,22 @@
+package controllers
+
+import (
+	"github.com/aws/aws-sdk-go/service/autoscaling"
+	upgrademgrv1alpha1 "github.com/keikoproj/upgrade-manager/api/v1alpha1"
+)
+
+type NodeSelector interface {
+	SelectNodesForRestack(state ClusterState) []*autoscaling.Instance
+}
+
+func getNodeSelector(
+	asg *autoscaling.Group,
+	ruObj *upgrademgrv1alpha1.RollingUpgrade,
+) NodeSelector {
+	switch ruObj.Spec.Strategy.Type {
+	case upgrademgrv1alpha1.UniformAcrossAzUpdateStrategy:
+		return NewUniformAcrossAzNodeSelector(asg, ruObj)
+	default:
+		return NewRandomNodeSelector(asg, ruObj)
+	}
+}
diff --git a/controllers/random_node_selector.go b/controllers/random_node_selector.go
new file mode 100644
index 0000000..e89cd7f
--- /dev/null
+++ b/controllers/random_node_selector.go
@@ -0,0 +1,33 @@
+package controllers
+
+import (
+	"github.com/aws/aws-sdk-go/service/autoscaling"
+	upgrademgrv1alpha1 "github.com/keikoproj/upgrade-manager/api/v1alpha1"
+	"log"
+)
+
+type RandomNodeSelector struct {
+	maxUnavailable int
+	ruObj          *upgrademgrv1alpha1.RollingUpgrade
+	asg            *autoscaling.Group
+}
+
+func NewRandomNodeSelector(
+	asg *autoscaling.Group,
+	ruObj *upgrademgrv1alpha1.RollingUpgrade,
+) *RandomNodeSelector {
+	maxUnavailable := getMaxUnavailable(ruObj.Spec.Strategy, len(asg.Instances))
+	log.Printf("Max unavailable calculated for %s is %d", ruObj.Name, maxUnavailable)
+	return &RandomNodeSelector{
+		maxUnavailable: maxUnavailable,
+		ruObj:          ruObj,
+		asg:            asg,
+	}
+}
+
+func (selector *RandomNodeSelector) SelectNodesForRestack(
+	state ClusterState,
+) []*autoscaling.Instance {
+	return getNextAvailableInstances(selector.ruObj.Spec.AsgName,
+		selector.maxUnavailable, selector.asg.Instances, state)
+}
diff --git a/controllers/rollingupgrade_controller.go b/controllers/rollingupgrade_controller.go
index 7f333dc..c49eaf9 100644
--- a/controllers/rollingupgrade_controller.go
+++ b/controllers/rollingupgrade_controller.go
@@ -159,7 +159,10 @@ func (r *RollingUpgradeReconciler) postDrainHelper(ruObj *upgrademgrv1alpha1.Rol
 
 // DrainNode runs "kubectl drain" on the given node
 // kubeCtlCall is provided as an argument to decouple the method from the actual kubectl call
-func (r *RollingUpgradeReconciler) DrainNode(ruObj *upgrademgrv1alpha1.RollingUpgrade, nodeName string, kubeCtlCall string, drainTimeout int) error {
+func (r *RollingUpgradeReconciler) DrainNode(ruObj *upgrademgrv1alpha1.RollingUpgrade,
+	nodeName string,
+	kubeCtlCall string,
+	drainTimeout int) error {
 	// Running kubectl drain node.
 	err := r.preDrainHelper(ruObj)
 	if err != nil {
@@ -362,14 +365,56 @@ func (r *RollingUpgradeReconciler) runRestack(ctx *context.Context, ruObj *upgra
 	// Setting default values for the Strategy in rollup object
 	r.setDefaultsForRollingUpdateStrategy(ruObj)
 
-	switch ruObj.Spec.Strategy.Type {
-	case upgrademgrv1alpha1.RandomUpdateStrategy:
-		log.Printf("Random update triggered for %s with strategy spec %+v", ruObj.Name, ruObj.Spec.Strategy)
-		return r.RandomUpdate(ctx, ruObj, svc, KubeCtlCall)
-	default:
-		error := errors.New(fmt.Sprintf("%s is not one of the predefined update strategies!", ruObj.Spec.Strategy.Type))
-		return 0, error
+	value, ok := r.ruObjNameToASG.Load(ruObj.Name)
+	if !ok {
+		msg := "Failed to find rollup name in map."
+		log.Printf(msg)
+		return 0, errors.New(msg)
+	}
+
+	asg := value.(*autoscaling.Group)
+	log.Printf("Nodes in ASG %s that *might* need to be updated: %d\n", *asg.AutoScalingGroupName, len(asg.Instances))
+
+	// No further processing is required if ASG doesn't have an instance running
+	totalNodes := len(asg.Instances)
+	// No further processing is required if ASG doesn't have an instance running
+	if totalNodes == 0 {
+		log.Printf("Total nodes found for %s is 0", ruObj.Name)
+		return 0, nil
 	}
+
+	nodeSelector := getNodeSelector(asg, ruObj)
+
+	// set the state of instances in the ASG to new in the cluster store
+	r.ClusterState.initializeAsg(*asg.AutoScalingGroupName, asg.Instances)
+
+	currentLaunchConfigName := aws.StringValue(asg.LaunchConfigurationName)
+
+	processedInstances := 0
+	for processedInstances < totalNodes {
+		// Fetch instances to update from node selector
+		instances := nodeSelector.SelectNodesForRestack(r.ClusterState)
+
+		if instances == nil {
+			errorMessage := fmt.Sprintf(
+				"No instances available for update across all AZ's for %s. Processed %d of total %d instances",
+				ruObj.Name, processedInstances, totalNodes)
+			// No instances fetched from any AZ, stop processing
+			log.Print(errorMessage)
+
+			// this should never be case, return error
+			return processedInstances, errors.New(errorMessage)
+		}
+
+		// update the instances
+		err := r.UpdateInstances(ctx, ruObj, instances, currentLaunchConfigName, KubeCtlCall, svc)
+		processedInstances += len(instances)
+		if err != nil {
+			return processedInstances, err
+		}
+	}
+	r.ClusterState.deleteEntryOfAsg(*asg.AutoScalingGroupName)
+	return processedInstances, nil
 }
 
 func (r *RollingUpgradeReconciler) finishExecution(finalStatus string, nodesProcessed int, ctx *context.Context, ruObj *upgrademgrv1alpha1.RollingUpgrade) (reconcile.Result, error) {
@@ -398,7 +443,8 @@ func (r *RollingUpgradeReconciler) finishExecution(finalStatus string, nodesProc
 }
 
 // Process actually performs the ec2-instance restacking.
-func (r *RollingUpgradeReconciler) Process(ctx *context.Context, ruObj *upgrademgrv1alpha1.RollingUpgrade) (reconcile.Result, error) {
+func (r *RollingUpgradeReconciler) Process(ctx *context.Context,
+	ruObj *upgrademgrv1alpha1.RollingUpgrade) (reconcile.Result, error) {
 	logr := r.Log.WithValues("rollingupgrade", ruObj.Name)
 
 	// If the object is being deleted, nothing to do.
@@ -542,34 +588,6 @@ func (r *RollingUpgradeReconciler) SetupWithManager(mgr ctrl.Manager) error {
 		Complete(r)
 }
 
-// getMaxUnavailable calculates and returns the maximum unavailable nodes
-// takes an update strategy and total number of nodes as input
-func getMaxUnavailable(strategy upgrademgrv1alpha1.UpdateStrategy, totalNodes int) int {
-	// Below are the constants set in intstr package
-	// const (
-	//	Int    Type = iota // The IntOrString holds an int.
-	//	String             // The IntOrString holds a string.
-	//)
-	maxUnavailable := 1
-	if strategy.MaxUnavailable.Type == 0 {
-		maxUnavailable = int(strategy.MaxUnavailable.IntVal)
-	} else if strategy.MaxUnavailable.Type == 1 {
-		strVallue := strategy.MaxUnavailable.StrVal
-		intValue, _ := strconv.Atoi(strings.Trim(strVallue, "%"))
-		maxUnavailable = int(float32(intValue) / float32(100) * float32(totalNodes))
-	}
-	// setting maxUnavailable to total number of nodes when maxUnavailable is greater than total node count
-	if totalNodes < maxUnavailable {
-		log.Printf("Reducing maxUnavailable count from %d to %d as total nodes count is %d", maxUnavailable, totalNodes, totalNodes)
-		maxUnavailable = totalNodes
-	}
-	// maxUnavailable has to be atleast 1 when there are nodes in the ASG
-	if totalNodes > 0 && maxUnavailable < 1 {
-		maxUnavailable = 1
-	}
-	return maxUnavailable
-}
-
 // validateRollingUpgradeObj validates rollup object for the type, maxUnavailable and drainTimeout
 func (r *RollingUpgradeReconciler) validateRollingUpgradeObj(ruObj *upgrademgrv1alpha1.RollingUpgrade) error {
 	strategy := ruObj.Spec.Strategy
@@ -582,7 +600,8 @@ func (r *RollingUpgradeReconciler) validateRollingUpgradeObj(ruObj *upgrademgrv1
 	// validating the maxUnavailable value
 	if strategy.MaxUnavailable.Type == 0 {
 		if strategy.MaxUnavailable.IntVal <= 0 {
-			err := errors.New(fmt.Sprintf("%s: Invalid value for maxUnavailable - %d", ruObj.Name, strategy.MaxUnavailable.IntVal))
+			err := errors.New(fmt.Sprintf("%s: Invalid value for maxUnavailable - %d",
+				ruObj.Name, strategy.MaxUnavailable.IntVal))
 			log.Print(err)
 			return err
 		}
@@ -590,14 +609,16 @@ func (r *RollingUpgradeReconciler) validateRollingUpgradeObj(ruObj *upgrademgrv1
 		strVallue := strategy.MaxUnavailable.StrVal
 		intValue, _ := strconv.Atoi(strings.Trim(strVallue, "%"))
 		if intValue <= 0 || intValue > 100 {
-			err := errors.New(fmt.Sprintf("%s: Invalid value for maxUnavailable - %s", ruObj.Name, strategy.MaxUnavailable.StrVal))
+			err := errors.New(fmt.Sprintf("%s: Invalid value for maxUnavailable - %s",
+				ruObj.Name, strategy.MaxUnavailable.StrVal))
 			log.Print(err)
 			return err
 		}
 	}
 
 	// validating the strategy type
-	if strategy.Type != upgrademgrv1alpha1.RandomUpdateStrategy {
+	if strategy.Type != upgrademgrv1alpha1.RandomUpdateStrategy &&
+		strategy.Type != upgrademgrv1alpha1.UniformAcrossAzUpdateStrategy {
 		err := errors.New(fmt.Sprintf("%s: Invalid value for strategy type - %s", ruObj.Name, strategy.Type))
 		log.Print(err)
 		return err
@@ -633,80 +654,69 @@ func (r *RollingUpgradeReconciler) setDefaultsForRollingUpdateStrategy(ruObj *up
 	}
 }
 
-// RandomUpdate treats all the azs as a single unit and picks random nodes for update
-// and rolls out the update based on the input parameters
-func (r *RollingUpgradeReconciler) RandomUpdate(ctx *context.Context, ruObj *upgrademgrv1alpha1.RollingUpgrade, svc autoscalingiface.AutoScalingAPI, KubeCtlCall string) (int, error) {
-
-	value, ok := r.ruObjNameToASG.Load(ruObj.Name)
-	if !ok {
-		msg := "Failed to find rollingUpgrade name in map."
-		log.Printf(msg)
-		return 0, errors.New(msg)
-	}
+type UpdateInstancesError struct {
+	InstanceUpdateErrors []error
+}
 
-	asg := value.(*autoscaling.Group)
-	log.Printf("Nodes in ASG %s that *might* need to be updated: %d\n", *asg.AutoScalingGroupName, len(asg.Instances))
+func (error UpdateInstancesError) Error() string {
+	return fmt.Sprintf("Error updating instances, ErrorCount: %d, Errors: %v",
+		len(error.InstanceUpdateErrors), error.InstanceUpdateErrors)
+}
 
-	// set the state of instances in the ASG to new in the cluster store
-	r.ClusterState.initializeAsg(*asg.AutoScalingGroupName, asg.Instances)
+func NewUpdateInstancesError(instanceUpdateErrors []error) *UpdateInstancesError {
+	return &UpdateInstancesError{InstanceUpdateErrors: instanceUpdateErrors}
+}
 
-	currentLaunchConfigName := aws.StringValue(asg.LaunchConfigurationName)
-	nodesProcessed := 0
-	nodesTriggered := 0
+func (r *RollingUpgradeReconciler) UpdateInstances(ctx *context.Context,
+	ruObj *upgrademgrv1alpha1.RollingUpgrade,
+	instances []*autoscaling.Instance,
+	currentLaunchConfigName string,
+	KubeCtlCall string,
+	svc autoscalingiface.AutoScalingAPI) error {
 
-	totalNodes := len(asg.Instances)
-	// No further processing is required if ASG doesn't have an instance running
+	totalNodes := len(instances)
 	if totalNodes == 0 {
-		log.Printf("Total nodes found for %s is 0", ruObj.Name)
-		return 0, nil
+		return nil
 	}
 
 	ch := make(chan error)
-	maxUnavailable := getMaxUnavailable(ruObj.Spec.Strategy, totalNodes)
-	log.Printf("Max unavailable calculated for %s is %d", ruObj.Name, maxUnavailable)
-
-	for nodesTriggered < maxUnavailable {
-		instance, available := r.getNextAvailableInstance(ruObj.Spec.AsgName, asg.Instances)
-		if !available {
-			err := errors.New("Instances are not available for update")
-			log.Printf("error: %s occurred for %s", err.Error(), ruObj.Name)
-			return nodesProcessed, err
-		}
-		go r.UpdateInstance(ctx, ruObj, instance, currentLaunchConfigName, KubeCtlCall, svc, ruObj.Spec.Strategy.DrainTimeout, ch)
-		nodesTriggered++
+
+	for _, instance := range instances {
+		go r.UpdateInstance(ctx, ruObj, instance, currentLaunchConfigName, KubeCtlCall, svc, ch)
 	}
 
+	// wait for upgrades to complete
+	nodesProcessed := 0
+	var instanceUpdateErrors []error
+Loop:
 	for err := range ch {
+		nodesProcessed++
 		switch err {
 		case nil:
-			nodesProcessed++
 			if nodesProcessed == totalNodes {
-				r.ClusterState.deleteEntryOfAsg(*asg.AutoScalingGroupName)
-				return nodesProcessed, nil
-			}
-			if nodesTriggered >= totalNodes {
-				continue
+				break Loop
 			}
-
 		default:
-			return nodesProcessed, err
+			instanceUpdateErrors = append(instanceUpdateErrors, err)
+			if nodesProcessed == totalNodes {
+				break Loop
+			}
 		}
+	}
 
-		instance, available := r.getNextAvailableInstance(ruObj.Spec.AsgName, asg.Instances)
-		if !available {
-			err := errors.New("Instances are not available for update")
-			log.Printf("error: %s occurred for %s", err.Error(), ruObj.Name)
-			return nodesProcessed, err
-		}
-		go r.UpdateInstance(ctx, ruObj, instance, currentLaunchConfigName, KubeCtlCall, svc, ruObj.Spec.Strategy.DrainTimeout, ch)
-		nodesTriggered++
+	if instanceUpdateErrors != nil && len(instanceUpdateErrors) > 0 {
+		return NewUpdateInstancesError(instanceUpdateErrors)
 	}
-	log.Printf("Deleting the entry of ASG - %s for %s", *asg.AutoScalingGroupName, ruObj.Name)
-	r.ClusterState.deleteEntryOfAsg(*asg.AutoScalingGroupName)
-	return nodesProcessed, nil
+	return nil
 }
 
-func (r *RollingUpgradeReconciler) UpdateInstance(ctx *context.Context, ruObj *upgrademgrv1alpha1.RollingUpgrade, i *autoscaling.Instance, currentLaunchConfigName string, KubeCtlCall string, svc autoscalingiface.AutoScalingAPI, drainTimeout int, ch chan error) {
+func (r *RollingUpgradeReconciler) UpdateInstance(ctx *context.Context,
+	ruObj *upgrademgrv1alpha1.RollingUpgrade,
+	i *autoscaling.Instance,
+	currentLaunchConfigName string,
+	KubeCtlCall string,
+	svc autoscalingiface.AutoScalingAPI,
+	ch chan error) {
 
 	targetLaunchConfigName := aws.StringValue(i.LaunchConfigurationName)
 	targetInstanceID := aws.StringValue(i.InstanceId)
@@ -755,18 +765,7 @@ func (r *RollingUpgradeReconciler) UpdateInstance(ctx *context.Context, ruObj *u
 	r.Update(*ctx, ruObj)
 
 	// TODO(shri): Run validate. How?
-	r.ClusterState.markUpdateCompleted(ruObj.Spec.AsgName, *i.InstanceId)
+	r.ClusterState.markUpdateCompleted(*i.InstanceId)
 	ch <- nil
 	return
 }
-
-// getNextAvailableInstance checks the cluster state store for the instance state and returns the next instance available for update
-func (r *RollingUpgradeReconciler) getNextAvailableInstance(asgName string, instances []*autoscaling.Instance) (*autoscaling.Instance, bool) {
-	instanceId := r.ClusterState.getNextAvailableInstanceId(asgName)
-	for _, instance := range instances {
-		if *instance.InstanceId == instanceId {
-			return instance, true
-		}
-	}
-	return nil, false
-}
diff --git a/controllers/rollup_cluster_state.go b/controllers/rollup_cluster_state.go
index 01e0f94..6b0feb5 100644
--- a/controllers/rollup_cluster_state.go
+++ b/controllers/rollup_cluster_state.go
@@ -17,9 +17,7 @@ limitations under the License.
 package controllers
 
 import (
-	"fmt"
 	"github.com/aws/aws-sdk-go/service/autoscaling"
-	"strings"
 	"sync"
 )
 
@@ -34,17 +32,26 @@ const (
 
 // ClusterState contains the methods to store the instance states during a cluster update
 type ClusterState interface {
-	markUpdateInitialized(asgName, instanceId string)
-	markUpdateInProgress(asgName, instanceId string)
-	markUpdateCompleted(asgName, instanceId string)
-	instanceUpdateInitialized(asgName, instanceId string) bool
-	instanceUpdateInProgress(asgName, instanceId string) bool
-	instanceUpdateCompleted(asgName, instanceId string) bool
+	markUpdateInitialized(instanceId string)
+	markUpdateInProgress(instanceId string)
+	markUpdateCompleted(instanceId string)
+	instanceUpdateInitialized(instanceId string) bool
+	instanceUpdateInProgress(instanceId string) bool
+	instanceUpdateCompleted(instanceId string) bool
 	deleteEntryOfAsg(asgName string) bool
-	getNextAvailableInstanceId(asgName string) string
+	getNextAvailableInstanceIdInAz(asgName string, azName string) string
 	initializeAsg(asgName string, instances []*autoscaling.Instance)
-	updateInstanceState(asgName, instanceId, instanceState string)
-	getInstanceState(asgName, instanceId string) string
+	addInstanceState(instanceData *InstanceData)
+	updateInstanceState(instanceId, instanceState string)
+	getInstanceState(instanceId string) string
+}
+
+type InstanceData struct {
+	Id            string
+	AmiId         string
+	AsgName       string
+	AzName        string
+	InstanceState string
 }
 
 // ClusterStateImpl implements the ClusterState interface
@@ -61,23 +68,23 @@ func NewClusterState() ClusterState {
 }
 
 // markUpdateInProgress updates the instance state to in-progresss
-func (c *ClusterStateImpl) markUpdateInProgress(asgName, instanceId string) {
-	c.updateInstanceState(asgName, instanceId, updateInProgress)
+func (c *ClusterStateImpl) markUpdateInProgress(instanceId string) {
+	c.updateInstanceState(instanceId, updateInProgress)
 }
 
 // markUpdateCompleted updates the instance state to completed
-func (c *ClusterStateImpl) markUpdateCompleted(asgName, instanceId string) {
-	c.updateInstanceState(asgName, instanceId, updateCompleted)
+func (c *ClusterStateImpl) markUpdateCompleted(instanceId string) {
+	c.updateInstanceState(instanceId, updateCompleted)
 }
 
 // instanceUpdateInProgress returns true if the instance update is in progress
-func (c *ClusterStateImpl) instanceUpdateInProgress(asgName, instanceId string) bool {
-	return c.getInstanceState(asgName, instanceId) == updateInProgress
+func (c *ClusterStateImpl) instanceUpdateInProgress(instanceId string) bool {
+	return c.getInstanceState(instanceId) == updateInProgress
 }
 
 // instanceUpdateCompleted returns true if the instance update is completed
-func (c *ClusterStateImpl) instanceUpdateCompleted(asgName, instanceId string) bool {
-	return c.getInstanceState(asgName, instanceId) == updateCompleted
+func (c *ClusterStateImpl) instanceUpdateCompleted(instanceId string) bool {
+	return c.getInstanceState(instanceId) == updateCompleted
 }
 
 // deleteEntryOfAsg deletes the entry for an ASG in the cluster state map
@@ -85,7 +92,8 @@ func (c *ClusterStateImpl) deleteEntryOfAsg(asgName string) bool {
 	deleted := false
 	ClusterStateStore.Range(func(key interface{}, value interface{}) bool {
 		keyName, _ := key.(string)
-		if strings.Contains(keyName, asgName) {
+		instanceData, _ := value.(*InstanceData)
+		if instanceData.AsgName == asgName {
 			ClusterStateStore.Delete(keyName)
 			deleted = true
 		}
@@ -95,35 +103,42 @@ func (c *ClusterStateImpl) deleteEntryOfAsg(asgName string) bool {
 }
 
 // markUpdateInitialized updates the instance state to in-progresss
-func (c *ClusterStateImpl) markUpdateInitialized(asgName, instanceId string) {
-	c.updateInstanceState(asgName, instanceId, updateInitialized)
+func (c *ClusterStateImpl) markUpdateInitialized(instanceId string) {
+	c.updateInstanceState(instanceId, updateInitialized)
 }
 
 // instanceUpdateInitialized returns true if the instance update is in progress
-func (c *ClusterStateImpl) instanceUpdateInitialized(asgName, instanceId string) bool {
-	return c.getInstanceState(asgName, instanceId) == updateInitialized
+func (c *ClusterStateImpl) instanceUpdateInitialized(instanceId string) bool {
+	return c.getInstanceState(instanceId) == updateInitialized
 }
 
 // initializeAsg adds an entry for all the instances in an ASG with updateInitialized state
 func (c *ClusterStateImpl) initializeAsg(asgName string, instances []*autoscaling.Instance) {
 	for _, instance := range instances {
-		c.updateInstanceState(asgName, *instance.InstanceId, updateInitialized)
+		instanceData := &InstanceData{
+			Id:            *instance.InstanceId,
+			AzName:        *instance.AvailabilityZone,
+			AsgName:       asgName,
+			InstanceState: updateInitialized,
+		}
+		c.addInstanceState(instanceData)
 	}
 }
 
 // getNextAvailableInstanceId returns the id of the next instance available for update in an ASG
 // adding a mutex to avoid the race conditions and same instance returned for 2 go-routines
-func (c *ClusterStateImpl) getNextAvailableInstanceId(asgName string) string {
+func (c *ClusterStateImpl) getNextAvailableInstanceIdInAz(asgName string, azName string) string {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
 	instanceId := ""
 	ClusterStateStore.Range(func(key interface{}, value interface{}) bool {
-		keyName, _ := key.(string)
-		state, _ := value.(string)
-		if strings.HasPrefix(keyName, asgName) && strings.EqualFold(state, updateInitialized) {
-			instanceId = strings.TrimPrefix(keyName, fmt.Sprintf("%s-", asgName))
-			c.markUpdateInProgress(asgName, instanceId)
+		state, _ := value.(*InstanceData)
+		if state.AsgName == asgName &&
+			(azName == "" || state.AzName == azName) &&
+			state.InstanceState == updateInitialized {
+			c.markUpdateInProgress(state.Id)
+			instanceId = state.Id
 			return false
 		}
 		return true
@@ -132,18 +147,26 @@ func (c *ClusterStateImpl) getNextAvailableInstanceId(asgName string) string {
 }
 
 // updateInstanceState updates the state of the instance in cluster store
-func (c *ClusterStateImpl) updateInstanceState(asgName, instanceId, instanceState string) {
-	key := fmt.Sprintf("%s-%s", asgName, instanceId)
-	ClusterStateStore.Store(key, instanceState)
+func (c *ClusterStateImpl) addInstanceState(instanceData *InstanceData) {
+	ClusterStateStore.Store(instanceData.Id, instanceData)
 }
 
-// getInstanceState returns the state of the instance from cluster store
-func (c *ClusterStateImpl) getInstanceState(asgName, instanceId string) string {
-	key := fmt.Sprintf("%s-%s", asgName, instanceId)
-	val, ok := ClusterStateStore.Load(key)
-	if !ok {
+// updateInstanceState updates the state of the instance in cluster store
+func (c *ClusterStateImpl) updateInstanceState(instanceId, instanceState string) {
+	val, ok := ClusterStateStore.Load(instanceId)
+	if ok {
+		instanceData, _ := val.(*InstanceData)
+		instanceData.InstanceState = instanceState
+		ClusterStateStore.Store(instanceId, instanceData)
+	}
+}
 
+// getInstanceState returns the state of the instance from cluster store
+func (c *ClusterStateImpl) getInstanceState(instanceId string) string {
+	val, ok := ClusterStateStore.Load(instanceId)
+	if ok {
+		state, _ := val.(*InstanceData)
+		return state.InstanceState
 	}
-	state, _ := val.(string)
-	return state
+	return ""
 }
diff --git a/controllers/uniform_across_az_node_selector.go b/controllers/uniform_across_az_node_selector.go
new file mode 100644
index 0000000..6b8c2b5
--- /dev/null
+++ b/controllers/uniform_across_az_node_selector.go
@@ -0,0 +1,66 @@
+package controllers
+
+import (
+	"github.com/aws/aws-sdk-go/service/autoscaling"
+	upgrademgrv1alpha1 "github.com/keikoproj/upgrade-manager/api/v1alpha1"
+	"log"
+)
+
+type azNodesCountState struct {
+	TotalNodes          int
+	MaxUnavailableNodes int
+}
+
+type UniformAcrossAzNodeSelector struct {
+	azNodeCounts map[string]*azNodesCountState
+	ruObj        *upgrademgrv1alpha1.RollingUpgrade
+	asg          *autoscaling.Group
+}
+
+func NewUniformAcrossAzNodeSelector(
+	asg *autoscaling.Group,
+	ruObj *upgrademgrv1alpha1.RollingUpgrade,
+) *UniformAcrossAzNodeSelector {
+
+	// find total number of nodes in each AZ
+	azNodeCounts := make(map[string]*azNodesCountState)
+	for _, instance := range asg.Instances {
+		if _, ok := azNodeCounts[*instance.AvailabilityZone]; ok {
+			azNodeCounts[*instance.AvailabilityZone].TotalNodes += 1
+		} else {
+			azNodeCounts[*instance.AvailabilityZone] = &azNodesCountState{TotalNodes: 1}
+		}
+	}
+
+	// find max unavailable for each az
+	for az, azNodeCount := range azNodeCounts {
+		azNodeCount.MaxUnavailableNodes = getMaxUnavailable(ruObj.Spec.Strategy, azNodeCount.TotalNodes)
+		log.Printf("Max unavailable calculated for %s, AZ %s is %d", ruObj.Name, az, azNodeCount.MaxUnavailableNodes)
+	}
+
+	return &UniformAcrossAzNodeSelector{
+		azNodeCounts: azNodeCounts,
+		ruObj:        ruObj,
+		asg:          asg,
+	}
+}
+
+func (selector *UniformAcrossAzNodeSelector) SelectNodesForRestack(
+	state ClusterState,
+) []*autoscaling.Instance {
+	var instances []*autoscaling.Instance
+
+	// Fetch instances to update from each instance group
+	for az, processedState := range selector.azNodeCounts {
+		// Collect the needed number of instances to update
+		instancesForUpdate := getNextSetOfAvailableInstancesInAz(selector.ruObj.Spec.AsgName,
+			az, processedState.MaxUnavailableNodes, selector.asg.Instances, state)
+		if instancesForUpdate == nil {
+			log.Printf("No instances available for update in AZ: %s for %s", az, selector.ruObj.Name)
+		} else {
+			instances = append(instances, instancesForUpdate...)
+		}
+	}
+
+	return instances
+}
diff --git a/examples/uniform_across_az_update_strategy.yaml b/examples/uniform_across_az_update_strategy.yaml
new file mode 100644
index 0000000..04f9d2f
--- /dev/null
+++ b/examples/uniform_across_az_update_strategy.yaml
@@ -0,0 +1,26 @@
+apiVersion: upgrademgr.keikoproj.io/v1alpha1
+kind: RollingUpgrade
+metadata:
+  generateName: rollingupgrade-sample-
+spec:
+  asgName: my-asg-1
+  nodeIntervalSeconds: 300
+  preDrain:
+    script: |
+      kubectl get pods --all-namespaces --field-selector spec.nodeName=${INSTANCE_NAME}
+  postDrain:
+    script: |
+      echo "Pods at PostDrain:"
+      kubectl get pods --all-namespaces --field-selector spec.nodeName=${INSTANCE_NAME}
+    waitSeconds: 90
+    postWaitScript: |
+      echo "Pods at postWait:"
+      kubectl get pods --all-namespaces --field-selector spec.nodeName=${INSTANCE_NAME}
+  postTerminate:
+    script: |
+      echo "PostTerminate:"
+      kubectl get pods --all-namespaces
+  strategy:
+    type: "uniformAcrossAzUpdate"
+    maxUnavailable: "25%"
+    drainTimeout: 120


Test patch:
diff --git a/controllers/helpers_test.go b/controllers/helpers_test.go
new file mode 100644
index 0000000..35177ac
--- /dev/null
+++ b/controllers/helpers_test.go
@@ -0,0 +1,174 @@
+package controllers
+
+import (
+	"encoding/json"
+	"fmt"
+	"github.com/aws/aws-sdk-go/service/autoscaling"
+	upgrademgrv1alpha1 "github.com/keikoproj/upgrade-manager/api/v1alpha1"
+	"github.com/onsi/gomega"
+	"testing"
+)
+
+func TestGetMaxUnavailableWithPercentageValue(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+
+	jsonString := "{ \"type\": \"randomUpdate\", \"maxUnavailable\": \"75%\", \"drainTimeout\": 15 }"
+	strategy := upgrademgrv1alpha1.UpdateStrategy{}
+	err := json.Unmarshal([]byte(jsonString), &strategy)
+	if err != nil {
+		fmt.Printf("Error occurred while unmarshalling strategy object, error: %s", err.Error())
+	}
+
+	g.Expect(getMaxUnavailable(strategy, 200)).To(gomega.Equal(150))
+}
+
+func TestGetMaxUnavailableWithPercentageValue33(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+
+	jsonString := "{ \"type\": \"randomUpdate\", \"maxUnavailable\": \"67%\", \"drainTimeout\": 15 }"
+	strategy := upgrademgrv1alpha1.UpdateStrategy{}
+	err := json.Unmarshal([]byte(jsonString), &strategy)
+	if err != nil {
+		fmt.Printf("Error occurred while unmarshalling strategy object, error: %s", err.Error())
+	}
+
+	g.Expect(getMaxUnavailable(strategy, 3)).To(gomega.Equal(2))
+}
+
+func TestGetMaxUnavailableWithPercentageAndSingleInstance(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+
+	totalNodes := 1
+	jsonString := "{ \"type\": \"randomUpdate\", \"maxUnavailable\": \"67%\", \"drainTimeout\": 15 }"
+	strategy := upgrademgrv1alpha1.UpdateStrategy{}
+	err := json.Unmarshal([]byte(jsonString), &strategy)
+	if err != nil {
+		fmt.Printf("Error occurred while unmarshalling strategy object, error: %s", err.Error())
+	}
+
+	g.Expect(getMaxUnavailable(strategy, totalNodes)).To(gomega.Equal(1))
+}
+
+func TestGetMaxUnavailableWithPercentageNonIntResult(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+
+	jsonString := "{ \"type\": \"randomUpdate\", \"maxUnavailable\": \"37%\", \"drainTimeout\": 15 }"
+	strategy := upgrademgrv1alpha1.UpdateStrategy{}
+	err := json.Unmarshal([]byte(jsonString), &strategy)
+	if err != nil {
+		fmt.Printf("Error occurred while unmarshalling strategy object, error: %s", err.Error())
+	}
+
+	g.Expect(getMaxUnavailable(strategy, 50)).To(gomega.Equal(18))
+}
+
+func TestGetMaxUnavailableWithIntValue(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+
+	jsonString := "{ \"type\": \"randomUpdate\", \"maxUnavailable\": 75, \"drainTimeout\": 15 }"
+	strategy := upgrademgrv1alpha1.UpdateStrategy{}
+	err := json.Unmarshal([]byte(jsonString), &strategy)
+	if err != nil {
+		fmt.Printf("Error occurred while unmarshalling strategy object, error: %s", err.Error())
+	}
+
+	g.Expect(getMaxUnavailable(strategy, 200)).To(gomega.Equal(75))
+}
+
+func TestGetNextAvailableInstance(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+
+	mockAsgName := "some-asg"
+	mockInstanceName1 := "foo1"
+	mockInstanceName2 := "bar1"
+	az := "az-1"
+	instance1 := autoscaling.Instance{InstanceId: &mockInstanceName1, AvailabilityZone: &az}
+	instance2 := autoscaling.Instance{InstanceId: &mockInstanceName2, AvailabilityZone: &az}
+
+	instancesList := []*autoscaling.Instance{}
+	instancesList = append(instancesList, &instance1, &instance2)
+	rcRollingUpgrade := &RollingUpgradeReconciler{ClusterState: clusterState}
+	rcRollingUpgrade.ClusterState.initializeAsg(mockAsgName, instancesList)
+	available := getNextAvailableInstances(mockAsgName, 1, instancesList, rcRollingUpgrade.ClusterState)
+
+	g.Expect(1).Should(gomega.Equal(len(available)))
+	g.Expect(rcRollingUpgrade.ClusterState.deleteEntryOfAsg(mockAsgName)).To(gomega.BeTrue())
+
+}
+
+func TestGetNextAvailableInstanceNoInstanceFound(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+
+	mockAsgName := "some-asg"
+	mockInstanceName1 := "foo1"
+	mockInstanceName2 := "bar1"
+	az := "az-1"
+	instance1 := autoscaling.Instance{InstanceId: &mockInstanceName1, AvailabilityZone: &az}
+	instance2 := autoscaling.Instance{InstanceId: &mockInstanceName2, AvailabilityZone: &az}
+
+	instancesList := []*autoscaling.Instance{}
+	instancesList = append(instancesList, &instance1, &instance2)
+	rcRollingUpgrade := &RollingUpgradeReconciler{ClusterState: clusterState}
+	rcRollingUpgrade.ClusterState.initializeAsg(mockAsgName, instancesList)
+	available := getNextAvailableInstances("asg2", 1, instancesList, rcRollingUpgrade.ClusterState)
+
+	g.Expect(0).Should(gomega.Equal(len(available)))
+	g.Expect(rcRollingUpgrade.ClusterState.deleteEntryOfAsg(mockAsgName)).To(gomega.BeTrue())
+
+}
+
+func TestGetNextAvailableInstanceInAz(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+
+	mockAsgName := "some-asg"
+	mockInstanceName1 := "foo1"
+	mockInstanceName2 := "bar1"
+	az := "az-1"
+	az2 := "az-2"
+	instance1 := autoscaling.Instance{InstanceId: &mockInstanceName1, AvailabilityZone: &az}
+	instance2 := autoscaling.Instance{InstanceId: &mockInstanceName2, AvailabilityZone: &az2}
+
+	instancesList := []*autoscaling.Instance{}
+	instancesList = append(instancesList, &instance1, &instance2)
+	rcRollingUpgrade := &RollingUpgradeReconciler{ClusterState: clusterState}
+	rcRollingUpgrade.ClusterState.initializeAsg(mockAsgName, instancesList)
+
+	instances := getNextSetOfAvailableInstancesInAz(mockAsgName, az, 1, instancesList, rcRollingUpgrade.ClusterState)
+	g.Expect(1).Should(gomega.Equal(len(instances)))
+	g.Expect(mockInstanceName1).Should(gomega.Equal(*instances[0].InstanceId))
+
+	instances = getNextSetOfAvailableInstancesInAz(mockAsgName, az2, 1, instancesList, rcRollingUpgrade.ClusterState)
+	g.Expect(1).Should(gomega.Equal(len(instances)))
+	g.Expect(mockInstanceName2).Should(gomega.Equal(*instances[0].InstanceId))
+
+	instances = getNextSetOfAvailableInstancesInAz(mockAsgName, "az3", 1, instancesList, rcRollingUpgrade.ClusterState)
+	g.Expect(0).Should(gomega.Equal(len(instances)))
+
+	g.Expect(rcRollingUpgrade.ClusterState.deleteEntryOfAsg(mockAsgName)).To(gomega.BeTrue())
+
+}
+
+func TestGetNextAvailableInstanceInAzGetMultipleInstances(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+
+	mockAsgName := "some-asg"
+	mockInstanceName1 := "foo1"
+	mockInstanceName2 := "bar1"
+	az := "az-1"
+	instance1 := autoscaling.Instance{InstanceId: &mockInstanceName1, AvailabilityZone: &az}
+	instance2 := autoscaling.Instance{InstanceId: &mockInstanceName2, AvailabilityZone: &az}
+
+	instancesList := []*autoscaling.Instance{}
+	instancesList = append(instancesList, &instance1, &instance2)
+	rcRollingUpgrade := &RollingUpgradeReconciler{ClusterState: clusterState}
+	rcRollingUpgrade.ClusterState.initializeAsg(mockAsgName, instancesList)
+
+	instances := getNextSetOfAvailableInstancesInAz(mockAsgName, az, 3, instancesList, rcRollingUpgrade.ClusterState)
+
+	// Even though the request is for 3 instances, only 2 should be returned as there are only 2 nodes in the ASG
+	g.Expect(2).Should(gomega.Equal(len(instances)))
+	instanceIds := []string{*instances[0].InstanceId, *instances[1].InstanceId}
+	g.Expect(instanceIds).Should(gomega.ConsistOf(mockInstanceName1, mockInstanceName2))
+
+	g.Expect(rcRollingUpgrade.ClusterState.deleteEntryOfAsg(mockAsgName)).To(gomega.BeTrue())
+}
diff --git a/controllers/node_selector_test.go b/controllers/node_selector_test.go
new file mode 100644
index 0000000..4eeae7a
--- /dev/null
+++ b/controllers/node_selector_test.go
@@ -0,0 +1,88 @@
+package controllers
+
+import (
+	"github.com/aws/aws-sdk-go/service/autoscaling"
+	upgrademgrv1alpha1 "github.com/keikoproj/upgrade-manager/api/v1alpha1"
+	"github.com/onsi/gomega"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"testing"
+)
+
+func TestGetRandomNodeSelector(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+
+	someAsg := "some-asg"
+	mockID := "some-id"
+	someLaunchConfig := "some-launch-config"
+	diffLaunchConfig := "different-launch-config"
+	az := "az-1"
+	mockInstance := autoscaling.Instance{InstanceId: &mockID, LaunchConfigurationName: &diffLaunchConfig, AvailabilityZone: &az}
+	mockAsg := autoscaling.Group{AutoScalingGroupName: &someAsg,
+		LaunchConfigurationName: &someLaunchConfig,
+		Instances:               []*autoscaling.Instance{&mockInstance}}
+
+	ruObj := &upgrademgrv1alpha1.RollingUpgrade{
+		ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"},
+		Spec:       upgrademgrv1alpha1.RollingUpgradeSpec{AsgName: someAsg},
+	}
+
+	nodeSelector := getNodeSelector(&mockAsg, ruObj)
+
+	g.Expect(nodeSelector).Should(gomega.BeAssignableToTypeOf(&RandomNodeSelector{}))
+}
+
+func TestGetUniformAcrossAzNodeSelector(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+
+	someAsg := "some-asg"
+	mockID := "some-id"
+	someLaunchConfig := "some-launch-config"
+	diffLaunchConfig := "different-launch-config"
+	az := "az-1"
+	mockInstance := autoscaling.Instance{InstanceId: &mockID, LaunchConfigurationName: &diffLaunchConfig, AvailabilityZone: &az}
+	mockAsg := autoscaling.Group{AutoScalingGroupName: &someAsg,
+		LaunchConfigurationName: &someLaunchConfig,
+		Instances:               []*autoscaling.Instance{&mockInstance}}
+
+	ruObj := &upgrademgrv1alpha1.RollingUpgrade{
+		ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"},
+		Spec: upgrademgrv1alpha1.RollingUpgradeSpec{
+			AsgName: someAsg,
+			Strategy: upgrademgrv1alpha1.UpdateStrategy{
+				Type: upgrademgrv1alpha1.UniformAcrossAzUpdateStrategy,
+			},
+		},
+	}
+
+	nodeSelector := getNodeSelector(&mockAsg, ruObj)
+
+	g.Expect(nodeSelector).Should(gomega.BeAssignableToTypeOf(&UniformAcrossAzNodeSelector{}))
+}
+
+func TestGetNodeSelectorWithInvalidStrategy(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+
+	someAsg := "some-asg"
+	mockID := "some-id"
+	someLaunchConfig := "some-launch-config"
+	diffLaunchConfig := "different-launch-config"
+	az := "az-1"
+	mockInstance := autoscaling.Instance{InstanceId: &mockID, LaunchConfigurationName: &diffLaunchConfig, AvailabilityZone: &az}
+	mockAsg := autoscaling.Group{AutoScalingGroupName: &someAsg,
+		LaunchConfigurationName: &someLaunchConfig,
+		Instances:               []*autoscaling.Instance{&mockInstance}}
+
+	ruObj := &upgrademgrv1alpha1.RollingUpgrade{
+		ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"},
+		Spec: upgrademgrv1alpha1.RollingUpgradeSpec{
+			AsgName: someAsg,
+			Strategy: upgrademgrv1alpha1.UpdateStrategy{
+				Type: "invalid",
+			},
+		},
+	}
+
+	nodeSelector := getNodeSelector(&mockAsg, ruObj)
+
+	g.Expect(nodeSelector).Should(gomega.BeAssignableToTypeOf(&RandomNodeSelector{}))
+}
diff --git a/controllers/rollingupgrade_controller_test.go b/controllers/rollingupgrade_controller_test.go
index 9990308..08bb31a 100644
--- a/controllers/rollingupgrade_controller_test.go
+++ b/controllers/rollingupgrade_controller_test.go
@@ -279,7 +279,12 @@ func TestDrainNodePostDrainFailureToDrain(t *testing.T) {
 	// Force quit from the rest of the command
 	mockKubeCtlCall := "exit 1;"
 
-	ruObj := &upgrademgrv1alpha1.RollingUpgrade{ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"}}
+	ruObj := &upgrademgrv1alpha1.RollingUpgrade{
+		ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"},
+		Spec: upgrademgrv1alpha1.RollingUpgradeSpec{
+			Strategy: upgrademgrv1alpha1.UpdateStrategy{DrainTimeout: -1},
+		},
+	}
 	rcRollingUpgrade := &RollingUpgradeReconciler{ClusterState: NewClusterState()}
 
 	err := rcRollingUpgrade.DrainNode(ruObj, mockNode, mockKubeCtlCall, ruObj.Spec.Strategy.DrainTimeout)
@@ -288,15 +293,19 @@ func TestDrainNodePostDrainFailureToDrain(t *testing.T) {
 
 type MockAutoscalingGroup struct {
 	autoscalingiface.AutoScalingAPI
-	errorFlag bool
-	awsErr    awserr.Error
+	errorFlag       bool
+	awsErr          awserr.Error
+	errorInstanceId string
 }
 
 func (mockAutoscalingGroup MockAutoscalingGroup) TerminateInstanceInAutoScalingGroup(input *autoscaling.TerminateInstanceInAutoScalingGroupInput) (*autoscaling.TerminateInstanceInAutoScalingGroupOutput, error) {
 	output := &autoscaling.TerminateInstanceInAutoScalingGroupOutput{}
 	if mockAutoscalingGroup.errorFlag {
 		if mockAutoscalingGroup.awsErr != nil {
-			return output, mockAutoscalingGroup.awsErr
+			if len(mockAutoscalingGroup.errorInstanceId) <= 0 ||
+				mockAutoscalingGroup.errorInstanceId == *input.InstanceId {
+				return output, mockAutoscalingGroup.awsErr
+			}
 		}
 	}
 	asgChange := autoscaling.Activity{ActivityId: aws.String("xxx"), AutoScalingGroupName: aws.String("sss"), Cause: aws.String("xxx"), StartTime: aws.Time(time.Now()), StatusCode: aws.String("200"), StatusMessage: aws.String("success")}
@@ -695,7 +704,8 @@ func TestRunRestackSuccessOneNode(t *testing.T) {
 	mockID := "some-id"
 	someLaunchConfig := "some-launch-config"
 	diffLaunchConfig := "different-launch-config"
-	mockInstance := autoscaling.Instance{InstanceId: &mockID, LaunchConfigurationName: &diffLaunchConfig}
+	az := "az-1"
+	mockInstance := autoscaling.Instance{InstanceId: &mockID, LaunchConfigurationName: &diffLaunchConfig, AvailabilityZone: &az}
 	mockAsg := autoscaling.Group{AutoScalingGroupName: &someAsg,
 		LaunchConfigurationName: &someLaunchConfig,
 		Instances:               []*autoscaling.Instance{&mockInstance}}
@@ -742,8 +752,9 @@ func TestRunRestackSuccessMultipleNodes(t *testing.T) {
 	mockID2 := "some-id-2"
 	someLaunchConfig := "some-launch-config"
 	diffLaunchConfig := "different-launch-config"
-	mockInstance := autoscaling.Instance{InstanceId: &mockID, LaunchConfigurationName: &diffLaunchConfig}
-	mockInstance2 := autoscaling.Instance{InstanceId: &mockID2, LaunchConfigurationName: &diffLaunchConfig}
+	az := "az-1"
+	mockInstance := autoscaling.Instance{InstanceId: &mockID, LaunchConfigurationName: &diffLaunchConfig, AvailabilityZone: &az}
+	mockInstance2 := autoscaling.Instance{InstanceId: &mockID2, LaunchConfigurationName: &diffLaunchConfig, AvailabilityZone: &az}
 	mockAsg := autoscaling.Group{AutoScalingGroupName: &someAsg,
 		LaunchConfigurationName: &someLaunchConfig,
 		Instances:               []*autoscaling.Instance{&mockInstance, &mockInstance2}}
@@ -787,7 +798,8 @@ func TestRunRestackSameLaunchConfig(t *testing.T) {
 	someAsg := "some-asg"
 	mockID := "some-id"
 	someLaunchConfig := "some-launch-config"
-	mockInstance := autoscaling.Instance{InstanceId: &mockID, LaunchConfigurationName: &someLaunchConfig}
+	az := "az-1"
+	mockInstance := autoscaling.Instance{InstanceId: &mockID, LaunchConfigurationName: &someLaunchConfig, AvailabilityZone: &az}
 	mockAsg := autoscaling.Group{AutoScalingGroupName: &someAsg,
 		LaunchConfigurationName: &someLaunchConfig,
 		Instances:               []*autoscaling.Instance{&mockInstance}}
@@ -829,7 +841,7 @@ func TestRunRestackRollingUpgradeNotInMap(t *testing.T) {
 	int, err := rcRollingUpgrade.runRestack(&ctx, ruObj, mockAutoscalingGroup, 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."))
+	g.Expect(err.Error()).To(gomega.HavePrefix("Failed to find rollup name in map."))
 }
 
 func TestRunRestackRollingUpgradeNodeNameNotFound(t *testing.T) {
@@ -839,7 +851,8 @@ func TestRunRestackRollingUpgradeNodeNameNotFound(t *testing.T) {
 	mockID := "some-id"
 	someLaunchConfig := "some-launch-config"
 	diffLaunchConfig := "different-launch-config"
-	mockInstance := autoscaling.Instance{InstanceId: &mockID, LaunchConfigurationName: &diffLaunchConfig}
+	az := "az-1"
+	mockInstance := autoscaling.Instance{InstanceId: &mockID, LaunchConfigurationName: &diffLaunchConfig, AvailabilityZone: &az}
 	mockAsg := autoscaling.Group{AutoScalingGroupName: &someAsg,
 		LaunchConfigurationName: &someLaunchConfig,
 		Instances:               []*autoscaling.Instance{&mockInstance}}
@@ -878,7 +891,8 @@ func TestRunRestackNoNodeName(t *testing.T) {
 	mockID := "some-id"
 	someLaunchConfig := "some-launch-config"
 	diffLaunchConfig := "different-launch-config"
-	mockInstance := autoscaling.Instance{InstanceId: &mockID, LaunchConfigurationName: &diffLaunchConfig}
+	az := "az-1"
+	mockInstance := autoscaling.Instance{InstanceId: &mockID, LaunchConfigurationName: &diffLaunchConfig, AvailabilityZone: &az}
 	mockAsg := autoscaling.Group{AutoScalingGroupName: &someAsg,
 		LaunchConfigurationName: &someLaunchConfig,
 		Instances:               []*autoscaling.Instance{&mockInstance}}
@@ -922,10 +936,15 @@ func TestRunRestackDrainNodeFail(t *testing.T) {
 	mockID := "some-id"
 	someLaunchConfig := "some-launch-config"
 	diffLaunchConfig := "different-launch-config"
-	mockInstance := autoscaling.Instance{InstanceId: &mockID, LaunchConfigurationName: &diffLaunchConfig}
+	az := "az-1"
+	mockInstance := autoscaling.Instance{InstanceId: &mockID,
+		LaunchConfigurationName: &diffLaunchConfig,
+		AvailabilityZone:        &az,
+	}
 	mockAsg := autoscaling.Group{AutoScalingGroupName: &someAsg,
 		LaunchConfigurationName: &someLaunchConfig,
-		Instances:               []*autoscaling.Instance{&mockInstance}}
+		Instances:               []*autoscaling.Instance{&mockInstance},
+	}
 
 	somePreDrain := upgrademgrv1alpha1.PreDrainSpec{
 		Script: "exit 1",
@@ -965,8 +984,8 @@ func TestRunRestackDrainNodeFail(t *testing.T) {
 
 	// 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)
-	g.Expect(nodesProcessed).To(gomega.Equal(0))
-	g.Expect(err.Error()).To(gomega.HavePrefix(ruObj.Name + ": Predrain script failed"))
+	g.Expect(nodesProcessed).To(gomega.Equal(1))
+	g.Expect(err.Error()).To(gomega.HavePrefix("Error updating instances, ErrorCount: 1, Errors: ["))
 }
 
 func TestRunRestackTerminateNodeFail(t *testing.T) {
@@ -976,7 +995,8 @@ func TestRunRestackTerminateNodeFail(t *testing.T) {
 	mockID := "some-id"
 	someLaunchConfig := "some-launch-config"
 	diffLaunchConfig := "different-launch-config"
-	mockInstance := autoscaling.Instance{InstanceId: &mockID, LaunchConfigurationName: &diffLaunchConfig}
+	az := "az-1"
+	mockInstance := autoscaling.Instance{InstanceId: &mockID, LaunchConfigurationName: &diffLaunchConfig, AvailabilityZone: &az}
 	mockAsg := autoscaling.Group{AutoScalingGroupName: &someAsg,
 		LaunchConfigurationName: &someLaunchConfig,
 		Instances:               []*autoscaling.Instance{&mockInstance}}
@@ -1013,74 +1033,283 @@ func TestRunRestackTerminateNodeFail(t *testing.T) {
 
 	// This execution gets past the different launch config check, but fails to terminate node
 	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, mockAutoscalingGroup, "exit 0;")
-	g.Expect(nodesProcessed).To(gomega.Equal(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"))
 }
 
-func TestGetMaxUnavailableWithPercentageValue(t *testing.T) {
+func constructAutoScalingInstance(instanceId string, launchConfigName string, azName string) *autoscaling.Instance {
+	return &autoscaling.Instance{InstanceId: &instanceId, LaunchConfigurationName: &launchConfigName, AvailabilityZone: &azName}
+}
+
+func TestUniformAcrossAzUpdateSuccessMultipleNodes(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 
-	jsonString := "{ \"type\": \"randomUpdate\", \"maxUnavailable\": \"75%\", \"drainTimeout\": 15 }"
-	strategy := upgrademgrv1alpha1.UpdateStrategy{}
-	err := json.Unmarshal([]byte(jsonString), &strategy)
-	if err != nil {
-		fmt.Printf("Error occurred while unmarshalling strategy object, error: %s", err.Error())
+	someAsg := "some-asg"
+	mockID := "some-id"
+	someLaunchConfig := "some-launch-config"
+	diffLaunchConfig := "different-launch-config"
+	az := "az-1"
+	az2 := "az-2"
+	az3 := "az-3"
+	mockAsg := autoscaling.Group{AutoScalingGroupName: &someAsg,
+		LaunchConfigurationName: &someLaunchConfig,
+		Instances: []*autoscaling.Instance{
+			constructAutoScalingInstance(mockID+"1"+az, diffLaunchConfig, az),
+			constructAutoScalingInstance(mockID+"2"+az, diffLaunchConfig, az),
+			constructAutoScalingInstance(mockID+"1"+az2, diffLaunchConfig, az2),
+			constructAutoScalingInstance(mockID+"2"+az2, diffLaunchConfig, az2),
+			constructAutoScalingInstance(mockID+"3"+az2, diffLaunchConfig, az2),
+			constructAutoScalingInstance(mockID+"1"+az3, diffLaunchConfig, az3),
+			constructAutoScalingInstance(mockID+"2"+az3, diffLaunchConfig, az3),
+			constructAutoScalingInstance(mockID+"3"+az3, diffLaunchConfig, az3),
+			constructAutoScalingInstance(mockID+"4"+az3, diffLaunchConfig, az3),
+		},
 	}
 
-	g.Expect(getMaxUnavailable(strategy, 200)).To(gomega.Equal(150))
+	ruObj := &upgrademgrv1alpha1.RollingUpgrade{
+		ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"},
+		Spec: upgrademgrv1alpha1.RollingUpgradeSpec{
+			AsgName: someAsg,
+			Strategy: upgrademgrv1alpha1.UpdateStrategy{
+				Type: upgrademgrv1alpha1.UniformAcrossAzUpdateStrategy,
+			},
+		},
+	}
+	mockAutoscalingGroup := MockAutoscalingGroup{}
+
+	mgr, err := manager.New(cfg, manager.Options{})
+	g.Expect(err).NotTo(gomega.HaveOccurred())
+	c = mgr.GetClient()
+
+	fooNode1 := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "foo-bar/9213851"}}
+	fooNode2 := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "foo-bar/1234501"}}
+	correctNode1az1 := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "fake-separator/" + mockID + "1" + az},
+		ObjectMeta: metav1.ObjectMeta{Name: "correct-node"}}
+	correctNode2az1 := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "fake-separator/" + mockID + "2" + az},
+		ObjectMeta: metav1.ObjectMeta{Name: "correct-node"}}
+	correctNode1az2 := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "fake-separator/" + mockID + "1" + az2},
+		ObjectMeta: metav1.ObjectMeta{Name: "correct-node"}}
+	correctNode2az2 := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "fake-separator/" + mockID + "2" + az2},
+		ObjectMeta: metav1.ObjectMeta{Name: "correct-node"}}
+	correctNode3az2 := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "fake-separator/" + mockID + "3" + az2},
+		ObjectMeta: metav1.ObjectMeta{Name: "correct-node"}}
+	correctNode1az3 := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "fake-separator/" + mockID + "1" + az3},
+		ObjectMeta: metav1.ObjectMeta{Name: "correct-node"}}
+	correctNode2az3 := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "fake-separator/" + mockID + "2" + az3},
+		ObjectMeta: metav1.ObjectMeta{Name: "correct-node"}}
+	correctNode3az3 := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "fake-separator/" + mockID + "3" + az3},
+		ObjectMeta: metav1.ObjectMeta{Name: "correct-node"}}
+	correctNode4az3 := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "fake-separator/" + mockID + "4" + az3},
+		ObjectMeta: metav1.ObjectMeta{Name: "correct-node"}}
+
+	nodeList := corev1.NodeList{Items: []corev1.Node{
+		fooNode1, fooNode2,
+		correctNode1az1, correctNode2az1,
+		correctNode1az2, correctNode2az2, correctNode3az2,
+		correctNode1az3, correctNode2az3, correctNode3az3, correctNode4az3,
+	}}
+	rcRollingUpgrade := &RollingUpgradeReconciler{
+		Client:          mgr.GetClient(),
+		generatedClient: kubernetes.NewForConfigOrDie(mgr.GetConfig()),
+		admissionMap:    sync.Map{},
+		ruObjNameToASG:  sync.Map{},
+		NodeList:        &nodeList,
+		ClusterState:    NewClusterState(),
+	}
+	rcRollingUpgrade.ruObjNameToASG.Store(ruObj.Name, &mockAsg)
+
+	ctx := context.TODO()
+
+	nodesProcessed, err := rcRollingUpgrade.runRestack(&ctx, ruObj, mockAutoscalingGroup, "exit 0;")
+	g.Expect(nodesProcessed).To(gomega.Equal(9))
+	g.Expect(err).To(gomega.BeNil())
 }
 
-func TestGetMaxUnavailableWithPercentageValue33(t *testing.T) {
+func TestUpdateInstances(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 
-	jsonString := "{ \"type\": \"randomUpdate\", \"maxUnavailable\": \"67%\", \"drainTimeout\": 15 }"
-	strategy := upgrademgrv1alpha1.UpdateStrategy{}
-	err := json.Unmarshal([]byte(jsonString), &strategy)
-	if err != nil {
-		fmt.Printf("Error occurred while unmarshalling strategy object, error: %s", err.Error())
+	someAsg := "some-asg"
+	mockID := "some-id"
+	mockID2 := "some-id-2"
+	someLaunchConfig := "some-launch-config"
+	diffLaunchConfig := "different-launch-config"
+	az := "az-1"
+	mockInstance := autoscaling.Instance{InstanceId: &mockID, LaunchConfigurationName: &diffLaunchConfig, AvailabilityZone: &az}
+	mockInstance2 := autoscaling.Instance{InstanceId: &mockID2, LaunchConfigurationName: &diffLaunchConfig, AvailabilityZone: &az}
+	mockAsg := autoscaling.Group{AutoScalingGroupName: &someAsg,
+		LaunchConfigurationName: &someLaunchConfig,
+		Instances:               []*autoscaling.Instance{&mockInstance, &mockInstance2}}
+
+	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())
+	c = mgr.GetClient()
+
+	fooNode1 := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "foo-bar/9213851"}}
+	fooNode2 := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "foo-bar/1234501"}}
+	correctNode := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "fake-separator/" + mockID},
+		ObjectMeta: metav1.ObjectMeta{Name: "correct-node"}}
+	correctNode2 := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "fake-separator/" + mockID2},
+		ObjectMeta: metav1.ObjectMeta{Name: "correct-node2"}}
+
+	nodeList := corev1.NodeList{Items: []corev1.Node{fooNode1, fooNode2, correctNode, correctNode2}}
+	rcRollingUpgrade := &RollingUpgradeReconciler{
+		Client:          mgr.GetClient(),
+		generatedClient: kubernetes.NewForConfigOrDie(mgr.GetConfig()),
+		admissionMap:    sync.Map{},
+		ruObjNameToASG:  sync.Map{},
+		NodeList:        &nodeList,
+		ClusterState:    NewClusterState(),
 	}
+	rcRollingUpgrade.ruObjNameToASG.Store(ruObj.Name, &mockAsg)
+
+	ctx := context.TODO()
 
-	g.Expect(getMaxUnavailable(strategy, 3)).To(gomega.Equal(2))
+	err = rcRollingUpgrade.UpdateInstances(&ctx,
+		ruObj, mockAsg.Instances, "A", "exit 0;", mockAutoscalingGroup)
+	g.Expect(err).ShouldNot(gomega.HaveOccurred())
 }
 
-func TestGetMaxUnavailableWithPercentageAndSingleInstance(t *testing.T) {
+func TestUpdateInstancesError(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 
-	totalNodes := 1
-	jsonString := "{ \"type\": \"randomUpdate\", \"maxUnavailable\": \"67%\", \"drainTimeout\": 15 }"
-	strategy := upgrademgrv1alpha1.UpdateStrategy{}
-	err := json.Unmarshal([]byte(jsonString), &strategy)
-	if err != nil {
-		fmt.Printf("Error occurred while unmarshalling strategy object, error: %s", err.Error())
+	someAsg := "some-asg"
+	mockID := "some-id"
+	mockID2 := "some-id-2"
+	someLaunchConfig := "some-launch-config"
+	diffLaunchConfig := "different-launch-config"
+	az := "az-1"
+	mockInstance := autoscaling.Instance{InstanceId: &mockID, LaunchConfigurationName: &diffLaunchConfig, AvailabilityZone: &az}
+	mockInstance2 := autoscaling.Instance{InstanceId: &mockID2, LaunchConfigurationName: &diffLaunchConfig, AvailabilityZone: &az}
+	mockAsg := autoscaling.Group{AutoScalingGroupName: &someAsg,
+		LaunchConfigurationName: &someLaunchConfig,
+		Instances:               []*autoscaling.Instance{&mockInstance, &mockInstance2}}
+
+	ruObj := &upgrademgrv1alpha1.RollingUpgrade{ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"},
+		Spec: upgrademgrv1alpha1.RollingUpgradeSpec{AsgName: someAsg}}
+	mockAutoScalingGroup := MockAutoscalingGroup{
+		errorFlag: true,
+		awsErr: awserr.New("UnKnownError",
+			"some message",
+			nil)}
+
+	mgr, err := manager.New(cfg, manager.Options{})
+	g.Expect(err).NotTo(gomega.HaveOccurred())
+	c = mgr.GetClient()
+
+	fooNode1 := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "foo-bar/9213851"}}
+	fooNode2 := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "foo-bar/1234501"}}
+	correctNode := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "fake-separator/" + mockID},
+		ObjectMeta: metav1.ObjectMeta{Name: "correct-node"}}
+	correctNode2 := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "fake-separator/" + mockID2},
+		ObjectMeta: metav1.ObjectMeta{Name: "correct-node2"}}
+
+	nodeList := corev1.NodeList{Items: []corev1.Node{fooNode1, fooNode2, correctNode, correctNode2}}
+	rcRollingUpgrade := &RollingUpgradeReconciler{
+		Client:          mgr.GetClient(),
+		generatedClient: kubernetes.NewForConfigOrDie(mgr.GetConfig()),
+		admissionMap:    sync.Map{},
+		ruObjNameToASG:  sync.Map{},
+		NodeList:        &nodeList,
+		ClusterState:    NewClusterState(),
 	}
+	rcRollingUpgrade.ruObjNameToASG.Store(ruObj.Name, &mockAsg)
 
-	g.Expect(getMaxUnavailable(strategy, totalNodes)).To(gomega.Equal(1))
+	ctx := context.TODO()
+
+	err = rcRollingUpgrade.UpdateInstances(&ctx,
+		ruObj, mockAsg.Instances, "A", "exit 0;", mockAutoScalingGroup)
+	g.Expect(err).Should(gomega.HaveOccurred())
+	g.Expect(err).Should(gomega.BeAssignableToTypeOf(&UpdateInstancesError{}))
+	if updateInstancesError, ok := err.(*UpdateInstancesError); ok {
+		g.Expect(len(updateInstancesError.InstanceUpdateErrors)).Should(gomega.Equal(2))
+		g.Expect(updateInstancesError.Error()).Should(gomega.ContainSubstring("Error updating instances, ErrorCount: 2"))
+	}
 }
 
-func TestGetMaxUnavailableWithPercentageNonIntResult(t *testing.T) {
+func TestUpdateInstancesPartialError(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 
-	jsonString := "{ \"type\": \"randomUpdate\", \"maxUnavailable\": \"37%\", \"drainTimeout\": 15 }"
-	strategy := upgrademgrv1alpha1.UpdateStrategy{}
-	err := json.Unmarshal([]byte(jsonString), &strategy)
-	if err != nil {
-		fmt.Printf("Error occurred while unmarshalling strategy object, error: %s", err.Error())
+	someAsg := "some-asg"
+	mockID := "some-id"
+	mockID2 := "some-id-2"
+	someLaunchConfig := "some-launch-config"
+	diffLaunchConfig := "different-launch-config"
+	az := "az-1"
+	mockInstance := autoscaling.Instance{InstanceId: &mockID, LaunchConfigurationName: &diffLaunchConfig, AvailabilityZone: &az}
+	mockInstance2 := autoscaling.Instance{InstanceId: &mockID2, LaunchConfigurationName: &diffLaunchConfig, AvailabilityZone: &az}
+	mockAsg := autoscaling.Group{AutoScalingGroupName: &someAsg,
+		LaunchConfigurationName: &someLaunchConfig,
+		Instances:               []*autoscaling.Instance{&mockInstance, &mockInstance2}}
+
+	ruObj := &upgrademgrv1alpha1.RollingUpgrade{ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"},
+		Spec: upgrademgrv1alpha1.RollingUpgradeSpec{AsgName: someAsg}}
+	mockAutoScalingGroup := MockAutoscalingGroup{
+		errorFlag: true,
+		awsErr: awserr.New("UnKnownError",
+			"some message",
+			nil),
+		errorInstanceId: mockID2,
+	}
+
+	mgr, err := manager.New(cfg, manager.Options{})
+	g.Expect(err).NotTo(gomega.HaveOccurred())
+	c = mgr.GetClient()
+
+	fooNode1 := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "foo-bar/9213851"}}
+	fooNode2 := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "foo-bar/1234501"}}
+	correctNode := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "fake-separator/" + mockID},
+		ObjectMeta: metav1.ObjectMeta{Name: "correct-node"}}
+	correctNode2 := corev1.Node{Spec: corev1.NodeSpec{ProviderID: "fake-separator/" + mockID2},
+		ObjectMeta: metav1.ObjectMeta{Name: "correct-node2"}}
+
+	nodeList := corev1.NodeList{Items: []corev1.Node{fooNode1, fooNode2, correctNode, correctNode2}}
+	rcRollingUpgrade := &RollingUpgradeReconciler{
+		Client:          mgr.GetClient(),
+		generatedClient: kubernetes.NewForConfigOrDie(mgr.GetConfig()),
+		admissionMap:    sync.Map{},
+		ruObjNameToASG:  sync.Map{},
+		NodeList:        &nodeList,
+		ClusterState:    NewClusterState(),
 	}
+	rcRollingUpgrade.ruObjNameToASG.Store(ruObj.Name, &mockAsg)
 
-	g.Expect(getMaxUnavailable(strategy, 50)).To(gomega.Equal(18))
+	ctx := context.TODO()
+
+	err = rcRollingUpgrade.UpdateInstances(&ctx,
+		ruObj, mockAsg.Instances, "A", "exit 0;", mockAutoScalingGroup)
+	g.Expect(err).Should(gomega.HaveOccurred())
+	g.Expect(err).Should(gomega.BeAssignableToTypeOf(&UpdateInstancesError{}))
+	if updateInstancesError, ok := err.(*UpdateInstancesError); ok {
+		g.Expect(len(updateInstancesError.InstanceUpdateErrors)).Should(gomega.Equal(1))
+		g.Expect(updateInstancesError.Error()).Should(gomega.Equal("Error updating instances, ErrorCount: 1, Errors: [UnKnownError: some message]"))
+	}
 }
 
-func TestGetMaxUnavailableWithIntValue(t *testing.T) {
+func TestUpdateInstancesWithZeroInstances(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 
-	jsonString := "{ \"type\": \"randomUpdate\", \"maxUnavailable\": 75, \"drainTimeout\": 15 }"
-	strategy := upgrademgrv1alpha1.UpdateStrategy{}
-	err := json.Unmarshal([]byte(jsonString), &strategy)
-	if err != nil {
-		fmt.Printf("Error occurred while unmarshalling strategy object, error: %s", err.Error())
+	mgr, err := manager.New(cfg, manager.Options{})
+	g.Expect(err).NotTo(gomega.HaveOccurred())
+	c = mgr.GetClient()
+
+	rcRollingUpgrade := &RollingUpgradeReconciler{
+		Client:          mgr.GetClient(),
+		generatedClient: kubernetes.NewForConfigOrDie(mgr.GetConfig()),
+		admissionMap:    sync.Map{},
+		ruObjNameToASG:  sync.Map{},
+		ClusterState:    NewClusterState(),
 	}
 
-	g.Expect(getMaxUnavailable(strategy, 200)).To(gomega.Equal(75))
+	ctx := context.TODO()
+
+	err = rcRollingUpgrade.UpdateInstances(&ctx,
+		nil, nil, "A", "exit 0;", nil)
+	g.Expect(err).ShouldNot(gomega.HaveOccurred())
 }
 
 func TestTestCallKubectlDrainWithoutDrainTimeout(t *testing.T) {
@@ -1273,27 +1502,7 @@ func TestTestCallKubectlDrainWithTimeoutOccurring(t *testing.T) {
 	g.Expect(output).To(gomega.ContainSubstring("timed-out"))
 }
 
-func TestGetNextAvailableInstance(t *testing.T) {
-	g := gomega.NewGomegaWithT(t)
-
-	mockAsgName := "some-asg"
-	mockInstanceName1 := "foo1"
-	mockInstanceName2 := "bar1"
-	instance1 := autoscaling.Instance{InstanceId: &mockInstanceName1}
-	instance2 := autoscaling.Instance{InstanceId: &mockInstanceName2}
-
-	instancesList := []*autoscaling.Instance{}
-	instancesList = append(instancesList, &instance1, &instance2)
-	rcRollingUpgrade := &RollingUpgradeReconciler{ClusterState: clusterState}
-	rcRollingUpgrade.ClusterState.initializeAsg(mockAsgName, instancesList)
-	_, available := rcRollingUpgrade.getNextAvailableInstance(mockAsgName, instancesList)
-
-	g.Expect(available).To(gomega.BeTrue())
-	g.Expect(rcRollingUpgrade.ClusterState.deleteEntryOfAsg(mockAsgName)).To(gomega.BeTrue())
-
-}
-
-func TestValidateruObj(t *testing.T) {
+func TestValidateRuObj(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 
 	strategyJsonString := "{ \"type\": \"randomUpdate\", \"maxUnavailable\": 75, \"drainTimeout\": 15 }"
@@ -1680,11 +1889,6 @@ func TestValidateruObjStrategyAfterSettingDefaultsWithOnlyMaxUnavailable(t *test
 	g.Expect(error).To((gomega.BeNil()))
 }
 
-func get(r *RollingUpgradeReconciler, asgName string, instanceList []*autoscaling.Instance, ch chan string) {
-	instance, _ := r.getNextAvailableInstance(asgName, instanceList)
-	ch <- *instance.InstanceId
-}
-
 func TestRunRestackNoNodeInAsg(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 
@@ -1730,7 +1934,8 @@ func TestRunRestackWithNodesLessThanMaxUnavailable(t *testing.T) {
 	someAsg := "some-asg"
 	mockID := "some-id"
 	someLaunchConfig := "some-launch-config"
-	mockInstance := autoscaling.Instance{InstanceId: &mockID, LaunchConfigurationName: &someLaunchConfig}
+	az := "az-1"
+	mockInstance := autoscaling.Instance{InstanceId: &mockID, LaunchConfigurationName: &someLaunchConfig, AvailabilityZone: &az}
 	mockAsg := autoscaling.Group{AutoScalingGroupName: &someAsg,
 		LaunchConfigurationName: &someLaunchConfig,
 		Instances:               []*autoscaling.Instance{&mockInstance}}
diff --git a/controllers/rollup_cluster_state_test.go b/controllers/rollup_cluster_state_test.go
index 42f5498..60bee4d 100644
--- a/controllers/rollup_cluster_state_test.go
+++ b/controllers/rollup_cluster_state_test.go
@@ -27,128 +27,129 @@ var clusterState = NewClusterState()
 func TestMarkUpdateInProgress(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 
-	mockAsgName := "some-asg"
-	mockNodeName := "some-node-name"
-	clusterState.markUpdateInProgress(mockAsgName, mockNodeName)
+	populateClusterState()
+	mockNodeName := "instance-1"
+	clusterState.markUpdateInProgress(mockNodeName)
 
-	g.Expect(clusterState.instanceUpdateInProgress(mockAsgName, mockNodeName)).To(gomega.BeTrue())
+	g.Expect(clusterState.instanceUpdateInProgress(mockNodeName)).To(gomega.BeTrue())
 }
 
 func TestMarkUpdateCompleted(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 
-	mockAsgName := "some-asg"
-	mockNodeName := "some-node-name"
-	clusterState.markUpdateCompleted(mockAsgName, mockNodeName)
+	populateClusterState()
+	mockNodeName := "instance-1"
+	clusterState.markUpdateCompleted(mockNodeName)
 
-	g.Expect(clusterState.instanceUpdateCompleted(mockAsgName, mockNodeName)).To(gomega.BeTrue())
+	g.Expect(clusterState.instanceUpdateCompleted(mockNodeName)).To(gomega.BeTrue())
 }
 
 func TestMarkUpdateInitialized(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 
-	mockAsgName := "some-asg"
-	mockNodeName := "some-node-name"
-	clusterState.markUpdateInitialized(mockAsgName, mockNodeName)
+	populateClusterState()
+	mockNodeName := "instance-1"
+	clusterState.markUpdateInitialized(mockNodeName)
 
-	g.Expect(clusterState.instanceUpdateInitialized(mockAsgName, mockNodeName)).To(gomega.BeTrue())
+	g.Expect(clusterState.instanceUpdateInitialized(mockNodeName)).To(gomega.BeTrue())
 }
 
 func TestInstanceUpdateInitialized(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 
-	mockAsgName := "some-asg-1"
-	mockNodeName := "some-node-name-1"
-
-	g.Expect(clusterState.instanceUpdateInitialized(mockAsgName, mockNodeName)).To(gomega.BeFalse())
+	mockNodeName := "instance-3"
+	g.Expect(clusterState.instanceUpdateInitialized(mockNodeName)).To(gomega.BeFalse())
 }
 
-func TestIinstanceUpdateInProgress(t *testing.T) {
+func TestInstanceUpdateInProgress(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 
-	mockAsgName := "some-asg-2"
-	mockNodeName := "some-node-name-2"
+	populateClusterState()
+	mockNodeName := "instance-2"
 
-	g.Expect(clusterState.instanceUpdateInProgress(mockAsgName, mockNodeName)).To(gomega.BeFalse())
+	g.Expect(clusterState.instanceUpdateInProgress(mockNodeName)).To(gomega.BeFalse())
 }
 
-func TestIinstanceUpdateCompleted(t *testing.T) {
+func TestInstanceUpdateCompleted(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 
-	mockAsgName := "some-asg-3"
-	mockNodeName := "some-node-name-3"
+	populateClusterState()
+	mockNodeName := "instance-1"
 
-	g.Expect(clusterState.instanceUpdateCompleted(mockAsgName, mockNodeName)).To(gomega.BeFalse())
+	g.Expect(clusterState.instanceUpdateCompleted(mockNodeName)).To(gomega.BeFalse())
 }
 
 func TestUpdateInstanceState(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 
-	mockAsgName := "some-asg"
-	mockNodeName := "some-node-name"
+	populateClusterState()
+	mockNodeName := "instance-1"
 	mockInstanceState := "to-be-updated"
-	clusterState.updateInstanceState(mockAsgName, mockNodeName, mockInstanceState)
+	clusterState.updateInstanceState(mockNodeName, mockInstanceState)
 
-	g.Expect(clusterState.getInstanceState(mockAsgName, mockNodeName)).To(gomega.ContainSubstring(mockInstanceState))
+	g.Expect(clusterState.getInstanceState(mockNodeName)).To(gomega.ContainSubstring(mockInstanceState))
 }
 
 func TestInitializeAsg(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 
-	mockAsgName := "some-asg"
-	mockInstance1 := "mock-instance-1"
-	mockInstance2 := "mock-instance-2"
-	instances := []*autoscaling.Instance{}
-	instances = append(instances, &autoscaling.Instance{InstanceId: &mockInstance1})
-	instances = append(instances, &autoscaling.Instance{InstanceId: &mockInstance2})
-	clusterState.initializeAsg(mockAsgName, instances)
+	populateClusterState()
 
-	for _, instance := range instances {
-		g.Expect(clusterState.instanceUpdateInitialized(mockAsgName, *instance.InstanceId)).To(gomega.BeTrue())
+	instanceIds := []string{"instance-1", "instance-2"}
+	for _, instance := range instanceIds {
+		g.Expect(clusterState.instanceUpdateInitialized(instance)).To(gomega.BeTrue())
 	}
 }
 
 func TestDeleteEntryOfAsg(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 
-	mockAsgName := "some-asg-1"
-	mockNodeName := "some-node-1"
-	clusterState.markUpdateInProgress(mockAsgName, mockNodeName)
+	populateClusterState()
+	mockAsgName := "asg-1"
 
 	g.Expect(clusterState.deleteEntryOfAsg(mockAsgName)).To(gomega.BeTrue())
 	g.Expect(clusterState.deleteEntryOfAsg(mockAsgName)).To(gomega.BeFalse())
-
 }
 
 func TestInstanceStateUpdateSequence(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 
-	mockAsgName := "some-asg-2"
-	mockNodeName := "some-node-2"
-	clusterState.markUpdateInitialized(mockAsgName, mockNodeName)
-	g.Expect(clusterState.instanceUpdateInitialized(mockAsgName, mockNodeName)).To(gomega.BeTrue())
-	clusterState.markUpdateInProgress(mockAsgName, mockNodeName)
-	g.Expect(clusterState.instanceUpdateInProgress(mockAsgName, mockNodeName)).To(gomega.BeTrue())
-	clusterState.markUpdateCompleted(mockAsgName, mockNodeName)
-	g.Expect(clusterState.instanceUpdateCompleted(mockAsgName, mockNodeName)).To(gomega.BeTrue())
+	populateClusterState()
+	mockAsgName := "asg-1"
+	mockNodeName := "instance-1"
+	clusterState.markUpdateInitialized(mockNodeName)
+	g.Expect(clusterState.instanceUpdateInitialized(mockNodeName)).To(gomega.BeTrue())
+	clusterState.markUpdateInProgress(mockNodeName)
+	g.Expect(clusterState.instanceUpdateInProgress(mockNodeName)).To(gomega.BeTrue())
+	clusterState.markUpdateCompleted(mockNodeName)
+	g.Expect(clusterState.instanceUpdateCompleted(mockNodeName)).To(gomega.BeTrue())
 
 	g.Expect(clusterState.deleteEntryOfAsg(mockAsgName)).To(gomega.BeTrue())
 	g.Expect(clusterState.deleteEntryOfAsg(mockAsgName)).To(gomega.BeFalse())
 }
 
-func TestGetNextAvailableInstanceId(t *testing.T) {
+func TestGetNextAvailableInstanceIdInAz(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 
-	mockAsgName := "some-asg-2"
-	mockInstance1 := "mock-instance-1"
-	mockInstance2 := "mock-instance-2"
-	instances := []*autoscaling.Instance{}
-	instances = append(instances, &autoscaling.Instance{InstanceId: &mockInstance1})
-	instances = append(instances, &autoscaling.Instance{InstanceId: &mockInstance2})
-	clusterState.deleteEntryOfAsg(mockAsgName)
+	populateClusterState()
+	mockAsgName := "asg-1"
+	mockInstance1 := "instance-1"
+	mockInstance2 := "instance-2"
 
-	clusterState.initializeAsg(mockAsgName, instances)
-	clusterState.markUpdateInProgress(mockAsgName, mockInstance1)
+	clusterState.markUpdateInProgress(mockInstance1)
 
-	g.Expect(clusterState.getNextAvailableInstanceId(mockAsgName)).To(gomega.ContainSubstring(mockInstance2))
+	g.Expect(clusterState.getNextAvailableInstanceIdInAz(mockAsgName, "az-1")).To(gomega.ContainSubstring(mockInstance2))
+
+	g.Expect(clusterState.getNextAvailableInstanceIdInAz(mockAsgName, "az-2")).To(gomega.ContainSubstring(""))
+}
+
+func populateClusterState() {
+	asgName := "asg-1"
+	instance1 := "instance-1"
+	instance2 := "instance-2"
+	az := "az-1"
+	instances := []*autoscaling.Instance{}
+	instances = append(instances, &autoscaling.Instance{InstanceId: &instance1, AvailabilityZone: &az})
+	instances = append(instances, &autoscaling.Instance{InstanceId: &instance2, AvailabilityZone: &az})
+	clusterState.initializeAsg(asgName, instances)
 }
diff --git a/controllers/uniform_across_az_node_selector_test.go b/controllers/uniform_across_az_node_selector_test.go
new file mode 100644
index 0000000..126e749
--- /dev/null
+++ b/controllers/uniform_across_az_node_selector_test.go
@@ -0,0 +1,107 @@
+package controllers
+
+import (
+	"github.com/aws/aws-sdk-go/service/autoscaling"
+	upgrademgrv1alpha1 "github.com/keikoproj/upgrade-manager/api/v1alpha1"
+	"github.com/onsi/gomega"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"testing"
+)
+
+func TestUniformAcrossAzNodeSelectorSelectNodes(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+
+	someAsg := "some-asg"
+	mockID := "some-id"
+	someLaunchConfig := "some-launch-config"
+	diffLaunchConfig := "different-launch-config"
+	az := "az-1"
+	az2 := "az-2"
+	az1Instance1 := constructAutoScalingInstance(mockID+"1-"+az, diffLaunchConfig, az)
+	az2Instance1 := constructAutoScalingInstance(mockID+"1-"+az2, diffLaunchConfig, az2)
+	az2Instance2 := constructAutoScalingInstance(mockID+"2-"+az2, diffLaunchConfig, az2)
+	mockAsg := autoscaling.Group{AutoScalingGroupName: &someAsg,
+		LaunchConfigurationName: &someLaunchConfig,
+		Instances: []*autoscaling.Instance{
+			az1Instance1,
+			az2Instance1,
+			az2Instance2,
+		},
+	}
+
+	ruObj := &upgrademgrv1alpha1.RollingUpgrade{
+		ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"},
+		Spec: upgrademgrv1alpha1.RollingUpgradeSpec{
+			AsgName: someAsg,
+			Strategy: upgrademgrv1alpha1.UpdateStrategy{
+				Type: upgrademgrv1alpha1.UniformAcrossAzUpdateStrategy,
+			},
+		},
+	}
+
+	clusterState := NewClusterState()
+	clusterState.initializeAsg(*mockAsg.AutoScalingGroupName, mockAsg.Instances)
+
+	nodeSelector := NewUniformAcrossAzNodeSelector(&mockAsg, ruObj)
+	instances := nodeSelector.SelectNodesForRestack(clusterState)
+
+	g.Expect(2).To(gomega.Equal(len(instances)))
+
+	// group instances by AZ
+	instancesByAz := make(map[string][]*autoscaling.Instance)
+	for _, instance := range instances {
+		az := instance.AvailabilityZone
+		if _, ok := instancesByAz[*az]; !ok {
+			instancesInAz := make([]*autoscaling.Instance, 0, len(instances))
+			instancesByAz[*az] = instancesInAz
+		}
+		instancesByAz[*az] = append(instancesByAz[*az], instance)
+	}
+
+	// assert on number of instances in each az
+	g.Expect(1).To(gomega.Equal(len(instancesByAz[az])))
+	g.Expect(1).To(gomega.Equal(len(instancesByAz[az2])))
+}
+
+func TestUniformAcrossAzNodeSelectorSelectNodesOneAzComplete(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+
+	someAsg := "some-asg"
+	mockID := "some-id"
+	someLaunchConfig := "some-launch-config"
+	diffLaunchConfig := "different-launch-config"
+	az := "az-1"
+	az2 := "az-2"
+	az1Instance1 := constructAutoScalingInstance(mockID+"1-"+az, diffLaunchConfig, az)
+	az2Instance1 := constructAutoScalingInstance(mockID+"1-"+az2, diffLaunchConfig, az2)
+	az2Instance2 := constructAutoScalingInstance(mockID+"2-"+az2, diffLaunchConfig, az2)
+	mockAsg := autoscaling.Group{AutoScalingGroupName: &someAsg,
+		LaunchConfigurationName: &someLaunchConfig,
+		Instances: []*autoscaling.Instance{
+			az1Instance1,
+			az2Instance1,
+			az2Instance2,
+		},
+	}
+
+	ruObj := &upgrademgrv1alpha1.RollingUpgrade{
+		ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"},
+		Spec: upgrademgrv1alpha1.RollingUpgradeSpec{
+			AsgName: someAsg,
+			Strategy: upgrademgrv1alpha1.UpdateStrategy{
+				Type: upgrademgrv1alpha1.UniformAcrossAzUpdateStrategy,
+			},
+		},
+	}
+
+	clusterState := NewClusterState()
+	clusterState.initializeAsg(*mockAsg.AutoScalingGroupName, mockAsg.Instances)
+	clusterState.markUpdateCompleted(mockID + "1-" + az)
+	clusterState.markUpdateCompleted(mockID + "1-" + az2)
+
+	nodeSelector := NewUniformAcrossAzNodeSelector(&mockAsg, ruObj)
+	instances := nodeSelector.SelectNodesForRestack(clusterState)
+
+	g.Expect(1).To(gomega.Equal(len(instances)))
+	g.Expect(&az2).To(gomega.Equal(instances[0].AvailabilityZone))
+}