Repository: spidernet-io/spiderpool
Pull request #1189 (2022-12-18 08:04:30)

test: Add unit tests for pkg statefulsetmanager
Signed-off-by: iiiceoo <ziqian.xue@daocloud.io>


Patch:
diff --git a/cmd/spiderpool-agent/cmd/daemon.go b/cmd/spiderpool-agent/cmd/daemon.go
index c8c345de6..0b89aee65 100644
--- a/cmd/spiderpool-agent/cmd/daemon.go
+++ b/cmd/spiderpool-agent/cmd/daemon.go
@@ -286,7 +286,7 @@ func initAgentServiceManagers(ctx context.Context) {
 	agentContext.PodManager = podManager
 
 	logger.Debug("Begin to initialize StatefulSet Manager")
-	statefulSetManager, err := statefulsetmanager.NewStatefulSetManager(agentContext.CRDManager)
+	statefulSetManager, err := statefulsetmanager.NewStatefulSetManager(agentContext.CRDManager.GetClient())
 	if err != nil {
 		logger.Fatal(err.Error())
 	}
diff --git a/cmd/spiderpool-controller/cmd/daemon.go b/cmd/spiderpool-controller/cmd/daemon.go
index c27d67f2c..3f9583f48 100644
--- a/cmd/spiderpool-controller/cmd/daemon.go
+++ b/cmd/spiderpool-controller/cmd/daemon.go
@@ -243,7 +243,7 @@ func initControllerServiceManagers(ctx context.Context) {
 	controllerContext.PodManager = podManager
 
 	logger.Info("Begin to initialize StatefulSet Manager")
-	statefulSetManager, err := statefulsetmanager.NewStatefulSetManager(controllerContext.CRDManager)
+	statefulSetManager, err := statefulsetmanager.NewStatefulSetManager(controllerContext.CRDManager.GetClient())
 	if err != nil {
 		logger.Fatal(err.Error())
 	}
diff --git a/pkg/statefulsetmanager/sts_manager.go b/pkg/statefulsetmanager/sts_manager.go
index ed7bbb2a8..7f887c245 100644
--- a/pkg/statefulsetmanager/sts_manager.go
+++ b/pkg/statefulsetmanager/sts_manager.go
@@ -5,13 +5,10 @@ package statefulsetmanager
 
 import (
 	"context"
-	"errors"
 	"fmt"
 
 	appsv1 "k8s.io/api/apps/v1"
-	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	apitypes "k8s.io/apimachinery/pkg/types"
-	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
 	"github.com/spidernet-io/spiderpool/pkg/constant"
@@ -20,42 +17,35 @@ import (
 type StatefulSetManager interface {
 	GetStatefulSetByName(ctx context.Context, namespace, name string) (*appsv1.StatefulSet, error)
 	ListStatefulSets(ctx context.Context, opts ...client.ListOption) (*appsv1.StatefulSetList, error)
-	IsValidStatefulSetPod(ctx context.Context, podNS, podName, podControllerType string) (bool, error)
+	IsValidStatefulSetPod(ctx context.Context, namespace, podName, podControllerType string) (bool, error)
 }
 
-type statefulSetMgr struct {
-	client     client.Client
-	runtimeMgr ctrl.Manager
+type statefulSetManager struct {
+	client client.Client
 }
 
-func NewStatefulSetManager(mgr ctrl.Manager) (StatefulSetManager, error) {
-	if mgr == nil {
-		return nil, errors.New("runtime manager must be specified")
+func NewStatefulSetManager(client client.Client) (StatefulSetManager, error) {
+	if client == nil {
+		return nil, fmt.Errorf("k8s clinet %w", constant.ErrMissingRequiredParam)
 	}
 
-	stsMgr := statefulSetMgr{
-		client:     mgr.GetClient(),
-		runtimeMgr: mgr,
-	}
-
-	return &stsMgr, nil
+	return &statefulSetManager{
+		client: client,
+	}, nil
 }
 
-func (sm *statefulSetMgr) GetStatefulSetByName(ctx context.Context, namespace, name string) (*appsv1.StatefulSet, error) {
+func (sm *statefulSetManager) GetStatefulSetByName(ctx context.Context, namespace, name string) (*appsv1.StatefulSet, error) {
 	var sts appsv1.StatefulSet
-
-	err := sm.client.Get(ctx, apitypes.NamespacedName{Namespace: namespace, Name: name}, &sts)
-	if nil != err {
+	if err := sm.client.Get(ctx, apitypes.NamespacedName{Namespace: namespace, Name: name}, &sts); err != nil {
 		return nil, err
 	}
 
 	return &sts, nil
 }
 
-func (sm *statefulSetMgr) ListStatefulSets(ctx context.Context, opts ...client.ListOption) (*appsv1.StatefulSetList, error) {
+func (sm *statefulSetManager) ListStatefulSets(ctx context.Context, opts ...client.ListOption) (*appsv1.StatefulSetList, error) {
 	var stsList appsv1.StatefulSetList
-	err := sm.client.List(ctx, &stsList, opts...)
-	if nil != err {
+	if err := sm.client.List(ctx, &stsList, opts...); err != nil {
 		return nil, err
 	}
 
@@ -65,38 +55,26 @@ func (sm *statefulSetMgr) ListStatefulSets(ctx context.Context, opts ...client.L
 // IsValidStatefulSetPod only serves for StatefulSet pod, it will check the pod whether need to be cleaned up with the given params podNS, podName.
 // Once the pod's controller StatefulSet was deleted, the pod's corresponding IPPool IP and Endpoint need to be cleaned up.
 // Or the pod's controller StatefulSet decreased its replicas and the pod's index is out of replicas, it needs to be cleaned up too.
-func (sm *statefulSetMgr) IsValidStatefulSetPod(ctx context.Context, podNS, podName, podControllerType string) (bool, error) {
+func (sm *statefulSetManager) IsValidStatefulSetPod(ctx context.Context, namespace, podName, podControllerType string) (bool, error) {
 	if podControllerType != constant.OwnerStatefulSet {
-		return false, fmt.Errorf("pod '%s/%s' owner controller type is '%s', not match StatefulSet type", podNS, podName, podControllerType)
+		return false, fmt.Errorf("pod '%s/%s' is controlled by '%s' instead of StatefulSet", namespace, podName, podControllerType)
 	}
 
-	statefulSetName, replicasIndex, found := getStatefulSetNameAndOrdinal(podName)
+	stsName, replicas, found := getStatefulSetNameAndOrdinal(podName)
 	if !found {
-		return false, fmt.Errorf("failed to get pod '%s/%s' controller StatefulSet name and pod replicas index", podNS, podName)
+		return false, fmt.Errorf("failed to parse the name and replica of its StatefulSet controller from the name of Pod '%s/%s'", namespace, podName)
 	}
 
-	isValid := true
-
-	statefulSet, err := sm.GetStatefulSetByName(ctx, podNS, statefulSetName)
-	if nil != err {
-		if !apierrors.IsNotFound(err) {
-			return false, err
-		}
-
-		// StatefulSet was deleted, just clean up IP and Endpoint
-		isValid = false
-	} else {
-		switch {
-		// pod restart
-		case replicasIndex == int(*statefulSet.Spec.Replicas)-1:
-
-		// StatefulSet decreased its replicas
-		case replicasIndex > int(*statefulSet.Spec.Replicas)-1:
-			isValid = false
+	sts, err := sm.GetStatefulSetByName(ctx, namespace, stsName)
+	if err != nil {
+		return false, client.IgnoreNotFound(err)
+	}
 
-		default:
-		}
+	// Pod controlled by StatefulSet is created or recreated.
+	if replicas <= int(*sts.Spec.Replicas)-1 {
+		return true, nil
 	}
 
-	return isValid, nil
+	// StatefulSet scaled down.
+	return false, nil
 }


Test patch:
diff --git a/pkg/namespacemanager/namespace_manager_test.go b/pkg/namespacemanager/namespace_manager_test.go
index 57398c2c1..e5eb329dc 100644
--- a/pkg/namespacemanager/namespace_manager_test.go
+++ b/pkg/namespacemanager/namespace_manager_test.go
@@ -43,7 +43,7 @@ var _ = Describe("NamespaceManager", Label("namespace_manager_test"), func() {
 			nsT = &corev1.Namespace{
 				TypeMeta: metav1.TypeMeta{
 					Kind:       "Namespace",
-					APIVersion: "v1",
+					APIVersion: corev1.SchemeGroupVersion.String(),
 				},
 				ObjectMeta: metav1.ObjectMeta{
 					Name:   nsName,
diff --git a/pkg/namespacemanager/utils_test.go b/pkg/namespacemanager/utils_test.go
index e49f77654..73768a37e 100644
--- a/pkg/namespacemanager/utils_test.go
+++ b/pkg/namespacemanager/utils_test.go
@@ -28,7 +28,7 @@ var _ = Describe("NamespaceManager utils", Label("namespace_manager_utils_test")
 			nsT = &corev1.Namespace{
 				TypeMeta: metav1.TypeMeta{
 					Kind:       "Namespace",
-					APIVersion: "v1",
+					APIVersion: corev1.SchemeGroupVersion.String(),
 				},
 				ObjectMeta: metav1.ObjectMeta{
 					Name: "namespace",
diff --git a/pkg/nodemanager/node_manager_test.go b/pkg/nodemanager/node_manager_test.go
index e913fae70..8698e6d16 100644
--- a/pkg/nodemanager/node_manager_test.go
+++ b/pkg/nodemanager/node_manager_test.go
@@ -43,7 +43,7 @@ var _ = Describe("NodeManager", Label("node_manager_test"), func() {
 			nodeT = &corev1.Node{
 				TypeMeta: metav1.TypeMeta{
 					Kind:       "Node",
-					APIVersion: "v1",
+					APIVersion: corev1.SchemeGroupVersion.String(),
 				},
 				ObjectMeta: metav1.ObjectMeta{
 					Name:   nodeName,
diff --git a/pkg/podmanager/pod_manager_test.go b/pkg/podmanager/pod_manager_test.go
index be8415e72..f8eceffdb 100644
--- a/pkg/podmanager/pod_manager_test.go
+++ b/pkg/podmanager/pod_manager_test.go
@@ -53,7 +53,7 @@ var _ = Describe("PodManager", Label("pod_manager_test"), func() {
 			podT = &corev1.Pod{
 				TypeMeta: metav1.TypeMeta{
 					Kind:       "Pod",
-					APIVersion: "v1",
+					APIVersion: corev1.SchemeGroupVersion.String(),
 				},
 				ObjectMeta: metav1.ObjectMeta{
 					Name:      podName,
diff --git a/pkg/podmanager/utils_test.go b/pkg/podmanager/utils_test.go
index c796f728a..59d1f03f4 100644
--- a/pkg/podmanager/utils_test.go
+++ b/pkg/podmanager/utils_test.go
@@ -22,7 +22,7 @@ var _ = Describe("PodManager utils", Label("pod_manager_utils_test"), func() {
 			podT = &corev1.Pod{
 				TypeMeta: metav1.TypeMeta{
 					Kind:       "Pod",
-					APIVersion: "v1",
+					APIVersion: corev1.SchemeGroupVersion.String(),
 				},
 				ObjectMeta: metav1.ObjectMeta{
 					Name:      "pod",
diff --git a/pkg/statefulsetmanager/sts_manager_suite_test.go b/pkg/statefulsetmanager/sts_manager_suite_test.go
new file mode 100644
index 000000000..a4e1f7715
--- /dev/null
+++ b/pkg/statefulsetmanager/sts_manager_suite_test.go
@@ -0,0 +1,39 @@
+// Copyright 2022 Authors of spidernet-io
+// SPDX-License-Identifier: Apache-2.0
+
+package statefulsetmanager_test
+
+import (
+	"testing"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+	appsv1 "k8s.io/api/apps/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/client/fake"
+
+	"github.com/spidernet-io/spiderpool/pkg/statefulsetmanager"
+)
+
+var scheme *runtime.Scheme
+var fakeClient client.Client
+var stsManager statefulsetmanager.StatefulSetManager
+
+func TestStatefulSetManager(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "StatefulSetManager Suite", Label("statefulsetmanager", "unitest"))
+}
+
+var _ = BeforeSuite(func() {
+	scheme = runtime.NewScheme()
+	err := appsv1.AddToScheme(scheme)
+	Expect(err).NotTo(HaveOccurred())
+
+	fakeClient = fake.NewClientBuilder().
+		WithScheme(scheme).
+		Build()
+
+	stsManager, err = statefulsetmanager.NewStatefulSetManager(fakeClient)
+	Expect(err).NotTo(HaveOccurred())
+})
diff --git a/pkg/statefulsetmanager/sts_manager_test.go b/pkg/statefulsetmanager/sts_manager_test.go
new file mode 100644
index 000000000..380824f65
--- /dev/null
+++ b/pkg/statefulsetmanager/sts_manager_test.go
@@ -0,0 +1,254 @@
+// Copyright 2022 Authors of spidernet-io
+// SPDX-License-Identifier: Apache-2.0
+
+package statefulsetmanager_test
+
+import (
+	"context"
+	"fmt"
+	"strconv"
+	"sync/atomic"
+
+	"github.com/agiledragon/gomonkey/v2"
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+	appsv1 "k8s.io/api/apps/v1"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/utils/pointer"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+
+	"github.com/spidernet-io/spiderpool/pkg/constant"
+	"github.com/spidernet-io/spiderpool/pkg/statefulsetmanager"
+)
+
+var _ = Describe("StatefulSetManager", Label("sts_manager_test"), func() {
+	Describe("New StatefulSetManager", func() {
+		It("inputs nil client", func() {
+			manager, err := statefulsetmanager.NewStatefulSetManager(nil)
+			Expect(err).To(MatchError(constant.ErrMissingRequiredParam))
+			Expect(manager).To(BeNil())
+		})
+	})
+
+	Describe("Test StatefulSetManager's method", func() {
+		var count uint64
+		var namespace string
+		var stsName string
+		var labels map[string]string
+		var stsT *appsv1.StatefulSet
+
+		BeforeEach(func() {
+			atomic.AddUint64(&count, 1)
+			namespace = "default"
+			stsName = fmt.Sprintf("sts-%v", count)
+			labels = map[string]string{"foo": fmt.Sprintf("bar-%v", count)}
+			stsT = &appsv1.StatefulSet{
+				TypeMeta: metav1.TypeMeta{
+					Kind:       "StatefulSet",
+					APIVersion: appsv1.SchemeGroupVersion.String(),
+				},
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      stsName,
+					Namespace: namespace,
+					Labels:    labels,
+				},
+				Spec: appsv1.StatefulSetSpec{},
+			}
+		})
+
+		var deleteOption *client.DeleteOptions
+
+		AfterEach(func() {
+			policy := metav1.DeletePropagationForeground
+			deleteOption = &client.DeleteOptions{
+				GracePeriodSeconds: pointer.Int64(0),
+				PropagationPolicy:  &policy,
+			}
+
+			ctx := context.TODO()
+			err := fakeClient.Delete(ctx, stsT, deleteOption)
+			Expect(client.IgnoreNotFound(err)).NotTo(HaveOccurred())
+		})
+
+		Describe("GetStatefulSetByName", func() {
+			It("gets non-existent StatefulSet", func() {
+				ctx := context.TODO()
+				sts, err := stsManager.GetStatefulSetByName(ctx, namespace, stsName)
+				Expect(apierrors.IsNotFound(err)).To(BeTrue())
+				Expect(sts).To(BeNil())
+			})
+
+			It("gets an existing StatefulSet", func() {
+				ctx := context.TODO()
+				err := fakeClient.Create(ctx, stsT)
+				Expect(err).NotTo(HaveOccurred())
+
+				sts, err := stsManager.GetStatefulSetByName(ctx, namespace, stsName)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(sts).NotTo(BeNil())
+
+				Expect(sts).To(Equal(stsT))
+			})
+		})
+
+		Describe("ListStatefulSets", func() {
+			It("failed to list StatefulSets due to some unknown errors", func() {
+				patches := gomonkey.ApplyMethodReturn(fakeClient, "List", constant.ErrUnknown)
+				defer patches.Reset()
+
+				ctx := context.TODO()
+				err := fakeClient.Create(ctx, stsT)
+				Expect(err).NotTo(HaveOccurred())
+
+				stsList, err := stsManager.ListStatefulSets(ctx)
+				Expect(err).To(MatchError(constant.ErrUnknown))
+				Expect(stsList).To(BeNil())
+			})
+
+			It("lists all StatefulSets", func() {
+				ctx := context.TODO()
+				err := fakeClient.Create(ctx, stsT)
+				Expect(err).NotTo(HaveOccurred())
+
+				stsList, err := stsManager.ListStatefulSets(ctx)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(stsList.Items).NotTo(BeEmpty())
+
+				hasSts := false
+				for _, sts := range stsList.Items {
+					if sts.Name == stsName {
+						hasSts = true
+						break
+					}
+				}
+				Expect(hasSts).To(BeTrue())
+			})
+
+			It("filters results by Namespace", func() {
+				ctx := context.TODO()
+				err := fakeClient.Create(ctx, stsT)
+				Expect(err).NotTo(HaveOccurred())
+
+				stsList, err := stsManager.ListStatefulSets(ctx, client.InNamespace(namespace))
+				Expect(err).NotTo(HaveOccurred())
+				Expect(stsList.Items).NotTo(BeEmpty())
+
+				hasSts := false
+				for _, sts := range stsList.Items {
+					if sts.Name == stsName {
+						hasSts = true
+						break
+					}
+				}
+				Expect(hasSts).To(BeTrue())
+			})
+
+			It("filters results by label selector", func() {
+				ctx := context.TODO()
+				err := fakeClient.Create(ctx, stsT)
+				Expect(err).NotTo(HaveOccurred())
+
+				stsList, err := stsManager.ListStatefulSets(ctx, client.MatchingLabels(labels))
+				Expect(err).NotTo(HaveOccurred())
+				Expect(stsList.Items).NotTo(BeEmpty())
+
+				hasSts := false
+				for _, sts := range stsList.Items {
+					if sts.Name == stsName {
+						hasSts = true
+						break
+					}
+				}
+				Expect(hasSts).To(BeTrue())
+			})
+
+			It("filters results by field selector", func() {
+				ctx := context.TODO()
+				err := fakeClient.Create(ctx, stsT)
+				Expect(err).NotTo(HaveOccurred())
+
+				stsList, err := stsManager.ListStatefulSets(ctx, client.MatchingFields{metav1.ObjectNameField: stsName})
+				Expect(err).NotTo(HaveOccurred())
+				Expect(stsList.Items).NotTo(BeEmpty())
+
+				hasSts := false
+				for _, sts := range stsList.Items {
+					if sts.Name == stsName {
+						hasSts = true
+						break
+					}
+				}
+				Expect(hasSts).To(BeTrue())
+			})
+		})
+
+		Describe("IsValidStatefulSetPod", func() {
+			var index int32
+			var stsPodName string
+			var nonStsPodName string
+
+			BeforeEach(func() {
+				index = 1
+				stsPodName = fmt.Sprintf("%s-%d", stsName, index)
+				nonStsPodName = "other"
+			})
+
+			It("is not a Pod controlled by StatefulSet", func() {
+				ctx := context.TODO()
+				valid, err := stsManager.IsValidStatefulSetPod(ctx, namespace, nonStsPodName, constant.OwnerUnknown)
+				Expect(err).To(HaveOccurred())
+				Expect(valid).To(BeFalse())
+			})
+
+			It("is a Pod controlled by StatefulSet, but the Pod name is invalid", func() {
+				ctx := context.TODO()
+				valid, err := stsManager.IsValidStatefulSetPod(ctx, namespace, nonStsPodName, constant.OwnerStatefulSet)
+				Expect(err).To(HaveOccurred())
+				Expect(valid).To(BeFalse())
+			})
+
+			It("failed to parse replica string to int due to some unknown errors", func() {
+				patches := gomonkey.ApplyFuncReturn(strconv.ParseInt, int64(0), constant.ErrUnknown)
+				defer patches.Reset()
+
+				ctx := context.TODO()
+				valid, err := stsManager.IsValidStatefulSetPod(ctx, namespace, stsPodName, constant.OwnerStatefulSet)
+				Expect(err).To(HaveOccurred())
+				Expect(valid).To(BeFalse())
+			})
+
+			It("is a valid Pod controlled by StatefulSet, but the StatefulSet no longer exists", func() {
+				ctx := context.TODO()
+				valid, err := stsManager.IsValidStatefulSetPod(ctx, namespace, stsPodName, constant.OwnerStatefulSet)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(valid).To(BeFalse())
+			})
+
+			It("used to be a Pod controlled by StatefulSet, but the StatefulSet scaled down", func() {
+				stsT.Spec.Replicas = &index
+
+				ctx := context.TODO()
+				err := fakeClient.Create(ctx, stsT)
+				Expect(err).NotTo(HaveOccurred())
+
+				valid, err := stsManager.IsValidStatefulSetPod(ctx, namespace, stsPodName, constant.OwnerStatefulSet)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(valid).To(BeFalse())
+			})
+
+			It("is a valid Pod controlled by StatefulSet", func() {
+				replicas := index + 1
+				stsT.Spec.Replicas = &replicas
+
+				ctx := context.TODO()
+				err := fakeClient.Create(ctx, stsT)
+				Expect(err).NotTo(HaveOccurred())
+
+				valid, err := stsManager.IsValidStatefulSetPod(ctx, namespace, stsPodName, constant.OwnerStatefulSet)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(valid).To(BeTrue())
+			})
+		})
+	})
+})