add test

Author: abbshr

pkg/driver/driver.go

@@ -130,7 +130,9 @@ func (d *Driver) Run() error {
 
 	// Remove taint from node to indicate driver startup success
 	// This is done at the last possible moment to prevent race conditions or false positive removals
-	go tryRemoveNotReadyTaintUntilSucceed(cloud.DefaultKubernetesAPIClient, time.Second)
+	go tryRemoveNotReadyTaintUntilSucceed(time.Second, func() error {
+		return removeNotReadyTaint(cloud.DefaultKubernetesAPIClient)
+	})
 
 	klog.Infof("Listening for connections on address: %#v", listener.Addr())
 	return d.srv.Serve(listener)

pkg/driver/node.go

@@ -515,9 +515,9 @@ func removeNotReadyTaint(k8sClient cloud.KubernetesAPIClient) error {
 }
 
 // remove taint may failed, this keep retring until succeed, make sure the taint will eventually being removed
-func tryRemoveNotReadyTaintUntilSucceed(k8sClient cloud.KubernetesAPIClient, interval time.Duration) {
+func tryRemoveNotReadyTaintUntilSucceed(interval time.Duration, removeFn func() error) {
 	for {
-		err := removeNotReadyTaint(k8sClient)
+		err := removeFn()
 		if err == nil {
 			return
 		}

pkg/driver/node_test.go

@@ -18,6 +18,7 @@ package driver
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"os"
 	"reflect"
@@ -988,3 +989,32 @@ func getNodeMock(mockCtl *gomock.Controller, nodeName string, returnNode *corev1
 
 	return mockClient, mockNode
 }
+
+func TestTryRemoveNotReadyTaintUntilSucceed(t *testing.T) {
+	{
+		i := 0
+		tryRemoveNotReadyTaintUntilSucceed(time.Second, func() error {
+			i++
+			if i < 3 {
+				return errors.New("test")
+			}
+
+			return nil
+		})
+
+		if i != 3 {
+			t.Fatalf("unexpected result")
+		}
+	}
+	{
+		i := 0
+		tryRemoveNotReadyTaintUntilSucceed(time.Second, func() error {
+			i++
+			return nil
+		})
+
+		if i != 1 {
+			t.Fatalf("unexpected result")
+		}
+	}
+}