test(utils/drain): fix failing tests

- refactor code to add cusom controller pod test
Signed-off-by: vadasambar <[email protected]>

cluster-autoscaler/utils/drain/drain_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package drain
 
 import (
+	"fmt"
 	"testing"
 	"time"
 
@@ -35,17 +36,29 @@ import (
 
 // testOpts represents parameters required for a single unit test
 type testOpts struct {
-	description         string
-	pods                []*apiv1.Pod
-	pdbs                []*policyv1.PodDisruptionBudget
-	rcs                 []*apiv1.ReplicationController
-	replicaSets         []*appsv1.ReplicaSet
-	expectFatal         bool
-	expectPods          []*apiv1.Pod
+	description string
+	pods        []*apiv1.Pod
+	pdbs        []*policyv1.PodDisruptionBudget
+	rcs         []*apiv1.ReplicationController
+	replicaSets []*appsv1.ReplicaSet
+	expectFatal bool
+	// expectFatalFn is ORed with expectFatal to decide fatality
+	// i.e., if (expectFatal || expectFatalFn()) = true, fatal
+	// if !(expectFatal || expectFatalFn()) = false , not fatal
+	expectFatalFn func() bool
+	expectPods    []*apiv1.Pod
+	// expectPodsFn is to provide more flexibility in adding pods
+	// return pods are appended to expectPods above
+	expectPodsFn        func() []*apiv1.Pod
 	expectDaemonSetPods []*apiv1.Pod
 	expectBlockingPod   *BlockingPod
+	// expectBlockingPodFn provides more flexibility than expectBllockingPod
+	// result from this function is used only if expectBlockingPod is nil
+	expectBlockingPodFn func() *BlockingPod
 }
 
+var skipNodesWithCustomControllerPods bool = false
+
 func TestDrain(t *testing.T) {
 	testTime := time.Date(2020, time.December, 18, 17, 0, 0, 0, time.UTC)
 	replicas := int32(5)
@@ -488,11 +501,26 @@ func TestDrain(t *testing.T) {
 			expectDaemonSetPods: []*apiv1.Pod{},
 		},
 		{
-			description:         "Custom-controller-managed pod",
-			pods:                []*apiv1.Pod{customControllerPod},
-			pdbs:                []*policyv1.PodDisruptionBudget{},
-			expectFatal:         false,
-			expectPods:          []*apiv1.Pod{customControllerPod},
+			description: "Custom-controller-managed pod",
+			pods:        []*apiv1.Pod{customControllerPod},
+			pdbs:        []*policyv1.PodDisruptionBudget{},
+			expectFatal: false,
+			expectFatalFn: func() bool {
+				return !skipNodesWithCustomControllerPods
+			},
+			expectPodsFn: func() []*apiv1.Pod {
+				if skipNodesWithCustomControllerPods {
+					return []*apiv1.Pod{customControllerPod}
+				}
+
+				return []*apiv1.Pod{}
+			},
+			expectBlockingPodFn: func() *BlockingPod {
+				if skipNodesWithCustomControllerPods {
+					return nil
+				}
+				return &BlockingPod{Pod: customControllerPod, Reason: NotReplicated}
+			},
 			expectDaemonSetPods: []*apiv1.Pod{},
 		},
 		{
@@ -650,12 +678,14 @@ func TestDrain(t *testing.T) {
 	}
 
 	// run all tests for skipNodesWithCustomControllerPods=false
-	// vadasambar: TODO: remove this when we get rid of SkipNodesWithCustomControllerPods
+	// vadasambar: TODO: remove this when we get rid of skipNodesWithCustomControllerPods
+	skipNodesWithCustomControllerPods = false
 	for _, test := range tests {
 		testIteration(t, &test, &ds, &job, &statefulset, testTime, false)
 	}
 
 	// run all tests for skipNodesWithCustomControllerPods=true
+	skipNodesWithCustomControllerPods = true
 	for _, test := range tests {
 		testIteration(t, &test, &ds, &job, &statefulset, testTime, true)
 	}
@@ -683,23 +713,39 @@ func testIteration(t *testing.T, test *testOpts, ds *appsv1.DaemonSet, job *batc
 
 	registry := kube_util.NewListerRegistry(nil, nil, nil, nil, nil, dsLister, rcLister, jobLister, rsLister, ssLister)
 
+	if test.description == "Custom-controller-managed pod" {
+		fmt.Println("skipNodesWithCustomControllerPods", skipNodesWithCustomControllerPods)
+	}
+
 	pods, daemonSetPods, blockingPod, err := GetPodsForDeletionOnNodeDrain(test.pods, test.pdbs, true, true, skipNodesWithCustomControllerPods, registry, 0, testTime)
 
-	if test.expectFatal {
-		assert.Equal(t, test.expectBlockingPod, blockingPod)
+	expectFatal := test.expectFatal || (test.expectFatalFn != nil && test.expectFatalFn())
+	expectBlockingPod := test.expectBlockingPod
+	if expectBlockingPod == nil && test.expectBlockingPodFn != nil {
+		expectBlockingPod = test.expectBlockingPodFn()
+	}
+	if expectFatal {
+		assert.Equal(t, expectBlockingPod, blockingPod)
 		if err == nil {
 			t.Fatalf("%s: unexpected non-error", test.description)
 		}
 	}
 
-	if !test.expectFatal {
+	if !expectFatal {
 		assert.Nil(t, blockingPod)
 		if err != nil {
 			t.Fatalf("%s: error occurred: %v", test.description, err)
 		}
 	}
 
-	if len(pods) != len(test.expectPods) {
+	expectPods := test.expectPods
+	if test.expectPodsFn != nil {
+		expectPods = append(expectPods, test.expectPodsFn()...)
+	}
+
+	if len(pods) != len(expectPods) {
+		fmt.Println("pods", pods)
+		fmt.Println("expectPods", expectPods)
 		t.Fatalf("Wrong pod list content: %v", test.description)
 	}