add ut and refactor for pkg/scheduler/plugins/util/k8s package

pkg/scheduler/plugins/util/k8s/snapshot_test.go

@@ -0,0 +1,147 @@
+package k8s
+
+import (
+	"fmt"
+	"reflect"
+	"testing"
+
+	v1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/labels"
+	"k8s.io/kubernetes/pkg/scheduler/framework"
+
+	"volcano.sh/volcano/pkg/scheduler/api"
+	"volcano.sh/volcano/pkg/scheduler/util"
+)
+
+func TestSnapshot(t *testing.T) {
+	var (
+		nodeName = "test-node"
+		pod1     = util.BuildPod("node_info_cache_test", "test-1", nodeName,
+			v1.PodRunning, api.BuildResourceList("200m", "1Ki"), "pg2",
+			make(map[string]string), make(map[string]string))
+		pod2 = util.BuildPod("node_info_cache_test", "test-2", nodeName,
+			v1.PodRunning, api.BuildResourceList("200m", "1Ki"), "pg2",
+			map[string]string{"test": "test"}, make(map[string]string))
+	)
+
+	tests := []struct {
+		name              string
+		nodeInfoMap       map[string]*framework.NodeInfo
+		expectedNodeInfos []*framework.NodeInfo
+		expectedPods      []*v1.Pod
+		expectErr         error
+	}{
+		{
+			name: "test snapshot operation",
+			nodeInfoMap: map[string]*framework.NodeInfo{nodeName: {
+				Requested:        &framework.Resource{},
+				NonZeroRequested: &framework.Resource{},
+				Allocatable:      &framework.Resource{},
+				Generation:       2,
+				ImageStates:      map[string]*framework.ImageStateSummary{},
+				PVCRefCounts:     map[string]int{},
+				Pods: []*framework.PodInfo{
+					{
+						Pod: pod1,
+					},
+					{
+						Pod: pod2,
+					},
+				},
+				PodsWithAffinity: []*framework.PodInfo{
+					{
+						Pod: pod1,
+					},
+					{
+						Pod: pod2,
+					},
+				},
+				PodsWithRequiredAntiAffinity: []*framework.PodInfo{
+					{
+						Pod: pod1,
+					},
+					{
+						Pod: pod2,
+					},
+				},
+			}},
+			expectedNodeInfos: []*framework.NodeInfo{{
+				Requested:        &framework.Resource{},
+				NonZeroRequested: &framework.Resource{},
+				Allocatable:      &framework.Resource{},
+				Generation:       2,
+				ImageStates:      map[string]*framework.ImageStateSummary{},
+				PVCRefCounts:     map[string]int{},
+				Pods: []*framework.PodInfo{
+					{
+						Pod: pod1,
+					},
+					{
+						Pod: pod2,
+					},
+				},
+				PodsWithAffinity: []*framework.PodInfo{
+					{
+						Pod: pod1,
+					},
+					{
+						Pod: pod2,
+					},
+				},
+				PodsWithRequiredAntiAffinity: []*framework.PodInfo{
+					{
+						Pod: pod1,
+					},
+					{
+						Pod: pod2,
+					},
+				},
+			}},
+			expectedPods: []*v1.Pod{pod2},
+			expectErr:    fmt.Errorf("nodeinfo not found for node name %q", nodeName),
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			snapshot := NewSnapshot(tc.nodeInfoMap)
+			nodeInfoList, err := snapshot.List()
+			if !reflect.DeepEqual(nodeInfoList, tc.expectedNodeInfos) || err != nil {
+				t.Fatalf("unexpected list nodeInfos value (+got: %s/-want: %s), err: %s", tc.expectedNodeInfos, nodeInfoList, err)
+			}
+
+			_, err = snapshot.Get(nodeName)
+			if !reflect.DeepEqual(tc.expectErr, err) {
+				t.Fatalf("unexpected get nodeInfos by nodeName value (+got: %T/-want: %T)", err, tc.expectErr)
+			}
+
+			nodeInfoList, err = snapshot.HavePodsWithAffinityList()
+			if !reflect.DeepEqual(tc.expectedNodeInfos, nodeInfoList) || err != nil {
+				t.Fatalf("unexpected list HavePodsWithAffinity nodeInfos value (+got: %s/-want: %s), err: %s", nodeInfoList, tc.expectedNodeInfos, err)
+			}
+
+			nodeInfoList, err = snapshot.HavePodsWithRequiredAntiAffinityList()
+			if !reflect.DeepEqual(tc.expectedNodeInfos, nodeInfoList) || err != nil {
+				t.Fatalf("unexpected list PodsWithRequiredAntiAffinity nodeInfos value (+got: %s/-want: %s), err: %s", nodeInfoList, tc.expectedNodeInfos, err)
+			}
+
+			sel, err := labels.Parse("test==test")
+			pods, err := snapshot.Pods().List(sel)
+			if !reflect.DeepEqual(tc.expectedPods, pods) || err != nil {
+				t.Fatalf("unexpected list pods value (+got: %s/-want: %s), err: %s", pods, tc.expectedNodeInfos, err)
+			}
+
+			pods, err = snapshot.Pods().FilteredList(func(pod *v1.Pod) bool {
+				return true
+			}, sel)
+			if !reflect.DeepEqual(tc.expectedPods, pods) || err != nil {
+				t.Fatalf("unexpected list filtered pods value (+got: %s/-want: %s), err: %s", pods, tc.expectedPods, err)
+			}
+
+			nodeInfos, err := snapshot.NodeInfos().List()
+			if !reflect.DeepEqual(tc.expectedNodeInfos, nodeInfos) || err != nil {
+				t.Fatalf("unexpected list nodeInfos value (+got: %s/-want: %s), err: %s", nodeInfos, tc.expectedNodeInfos, err)
+			}
+		})
+	}
+}

pkg/scheduler/plugins/util/nodelock/nodelock.go

@@ -23,6 +23,7 @@ import (
 	"path/filepath"
 	"time"
 
+	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/client-go/rest"
@@ -62,20 +63,11 @@ func UseClient(client kubernetes.Interface) error {
 	return nil
 }
 
-func setNodeLock(nodeName string, lockName string) error {
-	ctx := context.Background()
-	node, err := kubeClient.CoreV1().Nodes().Get(ctx, nodeName, metav1.GetOptions{})
-	if err != nil {
-		klog.Errorln("get node failed", err.Error())
-		return err
-	}
-	if _, ok := node.ObjectMeta.Annotations[lockName]; ok {
-		klog.V(3).Infof("node %s is locked", nodeName)
-		return fmt.Errorf("node %s is locked", nodeName)
-	}
+func updateNodeAnnotations(ctx context.Context, node *v1.Node, updateFunc func(annotations map[string]string)) error {
 	newNode := node.DeepCopy()
-	newNode.ObjectMeta.Annotations[lockName] = time.Now().Format(time.RFC3339)
-	_, err = kubeClient.CoreV1().Nodes().Update(ctx, newNode, metav1.UpdateOptions{})
+	updateFunc(newNode.ObjectMeta.Annotations)
+	nodeName := newNode.Name
+	_, err := kubeClient.CoreV1().Nodes().Update(ctx, newNode, metav1.UpdateOptions{})
 	for i := 0; i < MaxLockRetry && err != nil; i++ {
 		klog.ErrorS(err, "Failed to update node", "node", nodeName, "retry", i)
 		time.Sleep(100 * time.Millisecond)
@@ -84,18 +76,38 @@ func setNodeLock(nodeName string, lockName string) error {
 			klog.ErrorS(err, "Failed to get node when retry to update", "node", nodeName)
 			continue
 		}
-		newNode := node.DeepCopy()
-		newNode.ObjectMeta.Annotations[lockName] = time.Now().Format(time.RFC3339)
+		newNode = node.DeepCopy()
+		updateFunc(newNode.ObjectMeta.Annotations)
 		_, err = kubeClient.CoreV1().Nodes().Update(ctx, newNode, metav1.UpdateOptions{})
 	}
+	if err != nil {
+		klog.ErrorS(err, "Failed to update node", "node", nodeName)
+		return fmt.Errorf("failed to update node %s, exceeded retry count %d", nodeName, MaxLockRetry)
+	}
+	return nil
+}
+
+func setNodeLock(nodeName string, lockName string) error {
+	ctx := context.Background()
+	node, err := kubeClient.CoreV1().Nodes().Get(ctx, nodeName, metav1.GetOptions{})
+	if err != nil {
+		return err
+	}
+	if _, ok := node.ObjectMeta.Annotations[lockName]; ok {
+		klog.V(3).Infof("node %s is locked", nodeName)
+		return fmt.Errorf("node %s is locked", nodeName)
+	}
+	updateFunc := func(annotations map[string]string) {
+		annotations[lockName] = time.Now().Format(time.RFC3339)
+	}
+	err = updateNodeAnnotations(ctx, node, updateFunc)
 	if err != nil {
 		return fmt.Errorf("setNodeLock exceeds retry count %d", MaxLockRetry)
 	}
 	klog.InfoS("Node lock set", "node", nodeName)
 	return nil
 }
 
-// ReleaseNodeLock release a certain device lock on a certain node
 func ReleaseNodeLock(nodeName string, lockName string) error {
 	ctx := context.Background()
 	node, err := kubeClient.CoreV1().Nodes().Get(ctx, nodeName, metav1.GetOptions{})
@@ -106,21 +118,10 @@ func ReleaseNodeLock(nodeName string, lockName string) error {
 		klog.V(3).InfoS("Node lock not set", "node", nodeName)
 		return nil
 	}
-	newNode := node.DeepCopy()
-	delete(newNode.ObjectMeta.Annotations, lockName)
-	_, err = kubeClient.CoreV1().Nodes().Update(ctx, newNode, metav1.UpdateOptions{})
-	for i := 0; i < MaxLockRetry && err != nil; i++ {
-		klog.ErrorS(err, "Failed to update node", "node", nodeName, "retry", i)
-		time.Sleep(100 * time.Millisecond)
-		node, err = kubeClient.CoreV1().Nodes().Get(ctx, nodeName, metav1.GetOptions{})
-		if err != nil {
-			klog.ErrorS(err, "Failed to get node when retry to update", "node", nodeName)
-			continue
-		}
-		newNode := node.DeepCopy()
-		delete(newNode.ObjectMeta.Annotations, lockName)
-		_, err = kubeClient.CoreV1().Nodes().Update(ctx, newNode, metav1.UpdateOptions{})
+	updateFunc := func(annotations map[string]string) {
+		delete(annotations, lockName)
 	}
+	err = updateNodeAnnotations(ctx, node, updateFunc)
 	if err != nil {
 		return fmt.Errorf("releaseNodeLock exceeds retry count %d", MaxLockRetry)
 	}

pkg/scheduler/plugins/util/nodelock/nodelock_test.go

@@ -0,0 +1,59 @@
+package nodelock
+
+import (
+	"fmt"
+	"reflect"
+	"testing"
+	"time"
+
+	v1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/client-go/kubernetes/fake"
+)
+
+func TestLockNode(t *testing.T) {
+	var (
+		nodeName = "test-node"
+		lockName = "test-node-lock"
+	)
+
+	tests := []struct {
+		name      string
+		expectErr error
+		node      *v1.Node
+	}{
+		{
+			name:      "lock node success",
+			expectErr: nil,
+			node: buildNode(nodeName,
+				map[string]string{lockName: time.Now().Add(-time.Minute * 20).Format(time.RFC3339)}),
+		},
+		{
+			name:      "lock node failed",
+			expectErr: fmt.Errorf("node %s has been locked within 5 minutes", nodeName),
+			node: buildNode(nodeName,
+				map[string]string{lockName: time.Now().Add(time.Minute * 20).Format(time.RFC3339)}),
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			fakeClient := fake.NewSimpleClientset(tc.node)
+			_ = UseClient(fakeClient)
+			gotErr := LockNode(nodeName, lockName)
+			if !reflect.DeepEqual(tc.expectErr, gotErr) {
+				t.Errorf("LockNode error: (+got: %T/-want: %T)", gotErr, tc.expectErr)
+			}
+		})
+	}
+}
+
+// buildNode builts node
+func buildNode(name string, annotations map[string]string) *v1.Node {
+	return &v1.Node{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:        name,
+			Annotations: annotations,
+		},
+	}
+}