Optimize the time format layout

Signed-off-by: learner0810 <[email protected]>
Project-HAMi · Jan 5, 2025 · cbccbd4 · cbccbd4
1 parent ee7447a
commit cbccbd4
Show file tree

Hide file tree

Showing 5 changed files with 200 additions and 5 deletions.
diff --git a/pkg/device/ascend/device_test.go b/pkg/device/ascend/device_test.go
@@ -510,7 +510,7 @@ func Test_CheckHealth(t *testing.T) {
 				n: corev1.Node{
 					ObjectMeta: metav1.ObjectMeta{
 						Annotations: map[string]string{
-							util.HandshakeAnnos["huawei.com/Ascend910"]: "Requesting_2128.12.02 00:00:00",
+							util.HandshakeAnnos["huawei.com/Ascend910"]: "Requesting_2128-12-02 00:00:00",
 						},
 					},
 				},
@@ -564,7 +564,7 @@ func Test_CheckHealth(t *testing.T) {
 				n: corev1.Node{
 					ObjectMeta: metav1.ObjectMeta{
 						Annotations: map[string]string{
-							util.HandshakeAnnos["huawei.com/Ascend910"]: "Requesting_2024.01.02 00:00:00",
+							util.HandshakeAnnos["huawei.com/Ascend910"]: "Requesting_2024-01-02 00:00:00",
 						},
 					},
 				},

diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go
@@ -198,7 +198,7 @@ func (s *Scheduler) RegisterFromNodeAnnotations() {
 				_, ok := util.HandshakeAnnos[devhandsk]
 				if ok {
 					tmppat := make(map[string]string)
-					tmppat[util.HandshakeAnnos[devhandsk]] = "Requesting_" + time.Now().Format("2006.01.02 15:04:05")
+					tmppat[util.HandshakeAnnos[devhandsk]] = "Requesting_" + time.Now().Format(time.DateTime)
 					klog.V(5).InfoS("New timestamp", util.HandshakeAnnos[devhandsk], tmppat[util.HandshakeAnnos[devhandsk]], "nodeName", val.Name)
 					n, err := util.GetNode(val.Name)
 					if err != nil {

diff --git a/pkg/scheduler/scheduler_test.go b/pkg/scheduler/scheduler_test.go
@@ -18,6 +18,7 @@ package scheduler
 
 import (
 	"context"
+	"strings"
 	"testing"
 	"time"
 
@@ -604,3 +605,59 @@ func Test_Filter(t *testing.T) {
 		})
 	}
 }
+
+func Test_RegisterFromNodeAnnotations(t *testing.T) {
+	tests := []struct {
+		name      string
+		Scheduler *Scheduler
+		want      func(node *corev1.Node) bool
+	}{
+		{
+			name: "test node hand shake annotations layout",
+			Scheduler: func() *Scheduler {
+				s := NewScheduler()
+				s.stopCh = make(chan struct{})
+				s.nodeNotify = make(chan struct{})
+				client.KubeClient = fake.NewSimpleClientset()
+				s.kubeClient = client.KubeClient
+				informerFactory := informers.NewSharedInformerFactoryWithOptions(client.KubeClient, time.Hour*1)
+				s.nodeLister = informerFactory.Core().V1().Nodes().Lister()
+				node := &corev1.Node{
+					ObjectMeta: metav1.ObjectMeta{
+						Name: "node",
+					},
+				}
+				client.KubeClient.CoreV1().Nodes().Create(context.TODO(), node, metav1.CreateOptions{})
+				err := informerFactory.Core().V1().Nodes().Informer().GetIndexer().Add(node)
+				if err != nil {
+					t.Errorf("failed to create node err; %v", err)
+					return nil
+				}
+				return s
+			}(),
+			want: func(node *corev1.Node) bool {
+				_, err := time.Parse(time.DateTime, strings.TrimPrefix(node.Annotations["hami.io/node-handshake"], "Requesting_"))
+				_, err2 := time.Parse(time.DateTime, strings.TrimPrefix(node.Annotations["hami.io/node-handshake-dcu"], "Requesting_"))
+				return err == nil && err2 == nil
+			},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			time.AfterFunc(5*time.Second, func() {
+				close(test.Scheduler.stopCh)
+			})
+			go func() {
+				test.Scheduler.nodeNotify <- struct{}{}
+			}()
+			test.Scheduler.RegisterFromNodeAnnotations()
+			node, err := test.Scheduler.kubeClient.CoreV1().Nodes().Get(context.TODO(), "node", metav1.GetOptions{})
+			if err != nil {
+				t.Errorf("failed to get node err; %v", err)
+				return
+			}
+			assert.Equal(t, test.want(node), true)
+		})
+	}
+}
diff --git a/pkg/util/util.go b/pkg/util/util.go
@@ -371,7 +371,7 @@ func InitKlogFlags() *flag.FlagSet {
 func CheckHealth(devType string, n *corev1.Node) (bool, bool) {
 	handshake := n.Annotations[HandshakeAnnos[devType]]
 	if strings.Contains(handshake, "Requesting") {
-		formertime, _ := time.Parse("2006.01.02 15:04:05", strings.Split(handshake, "_")[1])
+		formertime, _ := time.Parse(time.DateTime, strings.Split(handshake, "_")[1])
 		return time.Now().Before(formertime.Add(time.Second * 60)), false
 	} else if strings.Contains(handshake, "Deleted") {
 		return true, false
@@ -382,7 +382,7 @@ func CheckHealth(devType string, n *corev1.Node) (bool, bool) {
 
 func MarkAnnotationsToDelete(devType string, nn string) error {
 	tmppat := make(map[string]string)
-	tmppat[devType] = "Deleted_" + time.Now().Format("2006.01.02 15:04:05")
+	tmppat[devType] = "Deleted_" + time.Now().Format(time.DateTime)
 	n, err := GetNode(nn)
 	if err != nil {
 		klog.Errorln("get node failed", err.Error())

diff --git a/pkg/util/util_test.go b/pkg/util/util_test.go
@@ -17,12 +17,18 @@ limitations under the License.
 package util
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"testing"
 
 	"gotest.tools/v3/assert"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/client-go/kubernetes/fake"
+
+	"github.com/Project-HAMi/HAMi/pkg/util/client"
 )
 
 var inRequestDevices map[string]string
@@ -465,3 +471,135 @@ func Test_EncodeNodeDevices(t *testing.T) {
 		})
 	}
 }
+
+func Test_CheckHealth(t *testing.T) {
+	tests := []struct {
+		name string
+		args struct {
+			devType string
+			n       corev1.Node
+		}
+		want1 bool
+		want2 bool
+	}{
+		{
+			name: "Requesting state",
+			args: struct {
+				devType string
+				n       corev1.Node
+			}{
+				devType: "huawei.com/Ascend910",
+				n: corev1.Node{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{
+							HandshakeAnnos["huawei.com/Ascend910"]: "Requesting_2128-12-02 00:00:00",
+						},
+					},
+				},
+			},
+			want1: true,
+			want2: false,
+		},
+		{
+			name: "Deleted state",
+			args: struct {
+				devType string
+				n       corev1.Node
+			}{
+				devType: "huawei.com/Ascend910",
+				n: corev1.Node{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{
+							HandshakeAnnos["huawei.com/Ascend910"]: "Deleted",
+						},
+					},
+				},
+			},
+			want1: true,
+			want2: false,
+		},
+		{
+			name: "Unknown state",
+			args: struct {
+				devType string
+				n       corev1.Node
+			}{
+				devType: "huawei.com/Ascend910",
+				n: corev1.Node{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{
+							HandshakeAnnos["huawei.com/Ascend910"]: "Unknown",
+						},
+					},
+				},
+			},
+			want1: true,
+			want2: true,
+		},
+		{
+			name: "Requesting state expired",
+			args: struct {
+				devType string
+				n       corev1.Node
+			}{
+				devType: "huawei.com/Ascend910",
+				n: corev1.Node{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{
+							HandshakeAnnos["huawei.com/Ascend910"]: "Requesting_2024-01-02 00:00:00",
+						},
+					},
+				},
+			},
+			want1: false,
+			want2: false,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			result1, result2 := CheckHealth(test.args.devType, &test.args.n)
+			assert.Equal(t, result1, test.want1)
+			assert.Equal(t, result2, test.want2)
+		})
+	}
+}
+
+func TestMarkAnnotationsToDelete(t *testing.T) {
+	client.KubeClient = fake.NewSimpleClientset()
+	client.KubeClient.CoreV1().Nodes().Create(context.TODO(), &corev1.Node{
+		ObjectMeta: metav1.ObjectMeta{Name: "node-worker2"},
+	}, metav1.CreateOptions{})
+	type args struct {
+		devType string
+		nn      string
+	}
+	tests := []struct {
+		name    string
+		args    args
+		wantErr bool
+	}{
+		{
+			name: "node not found",
+			args: args{
+				devType: "huawei.com/Ascend910",
+				nn:      "node-worker1",
+			},
+			wantErr: true,
+		},
+		{
+			name: "mark annotations to delete",
+			args: args{
+				devType: "huawei.com/Ascend910",
+				nn:      "node-worker2",
+			},
+			wantErr: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if err := MarkAnnotationsToDelete(tt.args.devType, tt.args.nn); (err != nil) != tt.wantErr {
+				t.Errorf("MarkAnnotationsToDelete() error = %v, wantErr %v", err, tt.wantErr)
+			}
+		})
+	}
+}