handle uncordon scenario properly

keikoproj · Dec 4, 2023 · 5a2a243 · 5a2a243
1 parent 717b481
commit 5a2a243
Show file tree

Hide file tree

Showing 6 changed files with 371 additions and 279 deletions.
diff --git a/Makefile b/Makefile
@@ -14,6 +14,7 @@ all: manager
 # Run tests
 ENVTEST_ASSETS_DIR=$(shell pwd)/testbin
 test: manifests generate fmt vet envtest
+	go test ./controllers/... ./api/...
 	go tool cover -html=./coverage.txt -o cover.html
 
 # Build manager binary

diff --git a/controllers/helpers_test.go b/controllers/helpers_test.go
@@ -426,3 +426,7 @@ func (mockAutoscalingGroup MockAutoscalingGroup) EnterStandby(_ *autoscaling.Ent
 	output := &autoscaling.EnterStandbyOutput{}
 	return output, nil
 }
+
+func (m *MockEC2) DeleteTags(input *ec2.DeleteTagsInput) (*ec2.DeleteTagsOutput, error) {
+	return &ec2.DeleteTagsOutput{}, nil
+}
diff --git a/controllers/metrics_test.go b/controllers/metrics_test.go
@@ -11,8 +11,8 @@ import (
 func TestNodeTurnsOntoStep(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 
-	reconsiler := createRollingUpgradeReconciler(t)
-	r := createRollingUpgradeContext(reconsiler)
+	reconciler := createRollingUpgradeReconciler(t)
+	r := createRollingUpgradeContext(reconciler)
 
 	//A map to retain the steps for multiple nodes
 	nodeSteps := make(map[string][]v1alpha1.NodeStepDuration)

diff --git a/controllers/upgrade.go b/controllers/upgrade.go
@@ -802,12 +802,12 @@ func (r *RollingUpgradeContext) CordonUncordonAllNodes(cordonNode bool) (bool, e
 					r.Error(err, "failed to early cordon the nodes", "instanceID", instance.InstanceId, "name", r.RollingUpgrade.NamespacedName())
 					return false, err
 				}
-			}
-			// Set instance-state to early-cordoned tag
-			r.Info("tagging instances with cordoned=true", "instanceID", instance.InstanceId, "name", r.RollingUpgrade.NamespacedName())
-			if err := r.Auth.TagEC2instances([]string{*instance.InstanceId}, instanceStateTagKey, earlyCordonedTagValue); err != nil {
-				r.Error(err, "failed to tag instances with cordoned=true", "instanceID", instance.InstanceId, "name", r.RollingUpgrade.NamespacedName())
-				return true, err
+				// Set instance-state to early-cordoned tag
+				r.Info("tagging instances with cordoned=true", "instanceID", instance.InstanceId, "name", r.RollingUpgrade.NamespacedName())
+				if err := r.Auth.TagEC2instances([]string{*instance.InstanceId}, instanceStateTagKey, earlyCordonedTagValue); err != nil {
+					r.Error(err, "failed to tag instances with cordoned=true", "instanceID", instance.InstanceId, "name", r.RollingUpgrade.NamespacedName())
+					return true, err
+				}
 			}
 		}
 	}

diff --git a/controllers/upgrade_test.go b/controllers/upgrade_test.go
@@ -1,6 +1,7 @@
 package controllers
 
 import (
+	"context"
 	"os"
 	"testing"
 
@@ -105,6 +106,7 @@ func TestRunCordonOrUncordon(t *testing.T) {
 	for _, test := range tests {
 		rollupCtx := createRollingUpgradeContext(test.Reconciler)
 		helper := &drain.Helper{
+			Ctx:                 context.Background(),
 			Client:              rollupCtx.Auth.Kubernetes,
 			Force:               true,
 			GracePeriodSeconds:  -1,
@@ -621,50 +623,58 @@ func TestClusterBallooning(t *testing.T) {
 	}
 }
 
-func TestEarlyCordonFunction(t *testing.T) {
+func TestCordoningAndUncordoningOfNodes(t *testing.T) {
 	var tests = []struct {
 		TestDescription            string
 		Reconciler                 *RollingUpgradeReconciler
-		RollingUpgrade             *v1alpha1.RollingUpgrade
-		AsgClient                  *MockAutoscalingGroup
-		ClusterNodes               []*corev1.Node
+		Node                       *corev1.Node
 		CordonNodeFlag             bool
 		ExpectedUnschdeulableValue bool
+		ExpectedError              bool
 	}{
 		{
-			"Test if all the nodes are cordoned by default.",
+			"Test if all the nodes are cordoned.",
 			createRollingUpgradeReconciler(t),
-			createRollingUpgrade(),
-			createASGClient(),
-			createNodeSlice(),
+			createNode("mock-node-1"),
 			true,
 			true,
+			false,
 		},
 		{
-			"Test if all the nodes are cordoned by default.",
+			"Test if all the nodes are uncordoned",
 			createRollingUpgradeReconciler(t),
-			createRollingUpgrade(),
-			createASGClient(),
-			createNodeSlice(),
+			createNode("mock-node-1"),
+			false,
 			false,
 			false,
 		},
+		{
+			"Try to cordon an unknown node.",
+			createRollingUpgradeReconciler(t),
+			createNode("mock-node-4"),
+			true,
+			true,
+			true,
+		},
 	}
 	for _, test := range tests {
 		rollupCtx := createRollingUpgradeContext(test.Reconciler)
-		rollupCtx.RollingUpgrade = test.RollingUpgrade
-		rollupCtx.Cloud.ScalingGroups = test.AsgClient.autoScalingGroups
-		rollupCtx.Cloud.ClusterNodes = test.ClusterNodes
-		rollupCtx.Auth.AmazonClientSet.AsgClient = test.AsgClient
 
-		_, err := rollupCtx.CordonUncordonAllNodes(test.CordonNodeFlag)
-		if err != nil {
+		if err := rollupCtx.Auth.CordonUncordonNode(test.Node, rollupCtx.Auth.Kubernetes, test.CordonNodeFlag); err != nil && test.ExpectedError {
+			continue
+		}
+
+		// By default, nodes are uncordoned. Therefore, before testing uncordoning the node, first cordon it.
+		if !test.CordonNodeFlag {
+			rollupCtx.Auth.CordonUncordonNode(test.Node, rollupCtx.Auth.Kubernetes, true)
+		}
+
+		if err := rollupCtx.Auth.CordonUncordonNode(test.Node, rollupCtx.Auth.Kubernetes, test.CordonNodeFlag); err != nil {
 			t.Errorf("Test Description: %s \n error: %v", test.TestDescription, err)
 		}
-		for _, node := range rollupCtx.Cloud.ClusterNodes {
-			if test.ExpectedUnschdeulableValue != node.Spec.Unschedulable {
-				t.Errorf("Test Description: %s \n expectedValue: %v, actualValue: %v", test.TestDescription, test.ExpectedUnschdeulableValue, node.Spec.Unschedulable)
-			}
+
+		if test.ExpectedUnschdeulableValue != test.Node.Spec.Unschedulable {
+			t.Errorf("Test Description: %s \n expectedValue: %v, actualValue: %v", test.TestDescription, test.ExpectedUnschdeulableValue, test.Node.Spec.Unschedulable)
 		}
 	}
 }