diff --git a/go.mod b/go.mod index 52a6ca20..d25cd4c6 100644 --- a/go.mod +++ b/go.mod @@ -8,6 +8,7 @@ require ( github.com/onsi/ginkgo v1.16.5 github.com/onsi/gomega v1.27.6 github.com/spf13/pflag v1.0.5 + github.com/stretchr/testify v1.8.4 google.golang.org/grpc v1.59.0 k8s.io/api v0.28.3 k8s.io/apimachinery v0.28.3 @@ -49,6 +50,7 @@ require ( github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect github.com/nxadm/tail v1.4.8 // indirect github.com/pkg/errors v0.9.1 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect github.com/prometheus/client_golang v1.17.0 // indirect github.com/prometheus/client_model v0.5.0 // indirect github.com/prometheus/common v0.45.0 // indirect diff --git a/pkg/driver/node.go b/pkg/driver/node.go index f0d32a61..287327dc 100644 --- a/pkg/driver/node.go +++ b/pkg/driver/node.go @@ -72,7 +72,7 @@ func newNodeService(driverOptions *DriverOptions) nodeService { // Remove taint from node to indicate driver startup success // This is done in the background as a goroutine to allow for driver startup - go removeTaintInBackground(cloud.DefaultKubernetesAPIClient) + go removeTaintInBackground(cloud.DefaultKubernetesAPIClient, removeNotReadyTaint) return nodeService{ metadata: metadata, @@ -298,12 +298,12 @@ type JSONPatch struct { } // removeTaintInBackground is a goroutine that retries removeNotReadyTaint with exponential backoff -func removeTaintInBackground(k8sClient cloud.KubernetesAPIClient) { +func removeTaintInBackground(k8sClient cloud.KubernetesAPIClient, removalFunc func(cloud.KubernetesAPIClient) error) { backoffErr := wait.ExponentialBackoff(taintRemovalBackoff, func() (bool, error) { - err := removeNotReadyTaint(k8sClient) + err := removalFunc(k8sClient) if err != nil { klog.ErrorS(err, "Unexpected failure when attempting to remove node taint(s)") - return false, err + return false, nil } return true, nil }) diff --git a/pkg/driver/node_test.go b/pkg/driver/node_test.go index e031ee6c..bebbbfa0 100644 --- a/pkg/driver/node_test.go +++ b/pkg/driver/node_test.go @@ -19,18 +19,19 @@ package driver import ( "context" "fmt" + "github.com/container-storage-interface/spec/lib/go/csi" + "github.com/golang/mock/gomock" + "github.com/stretchr/testify/assert" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" corev1 "k8s.io/api/core/v1" "k8s.io/client-go/kubernetes" "reflect" + "sigs.k8s.io/aws-fsx-csi-driver/pkg/cloud" cloudMock "sigs.k8s.io/aws-fsx-csi-driver/pkg/cloud/mocks" "sigs.k8s.io/aws-fsx-csi-driver/pkg/driver/internal" - "testing" - - "github.com/container-storage-interface/spec/lib/go/csi" - "github.com/golang/mock/gomock" driverMocks "sigs.k8s.io/aws-fsx-csi-driver/pkg/driver/mocks" + "testing" ) var ( @@ -824,6 +825,21 @@ func TestRemoveNotReadyTaint(t *testing.T) { } } +func TestRemoveTaintInBackground(t *testing.T) { + mockRemovalCount := 0 + mockRemovalFunc := func(_ cloud.KubernetesAPIClient) error { + mockRemovalCount += 1 + if mockRemovalCount == 3 { + return nil + } else { + return fmt.Errorf("Taint removal failed!") + } + } + + removeTaintInBackground(nil, mockRemovalFunc) + assert.Equal(t, mockRemovalCount, 3) +} + func getNodeMock(mockCtl *gomock.Controller, nodeName string, returnNode *corev1.Node, returnError error) (kubernetes.Interface, *driverMocks.MockNodeInterface) { mockClient := driverMocks.NewMockKubernetesClient(mockCtl) mockCoreV1 := driverMocks.NewMockCoreV1Interface(mockCtl)