From 711b5ae18f6d9cd0bc6cb626f88369ffc813f9ee Mon Sep 17 00:00:00 2001
From: jinye <jinye.shi@daocloud.io>
Date: Thu, 9 Jan 2025 10:27:41 +0800
Subject: [PATCH] add ut for device/iluvatar

Signed-off-by: jinye <jinye.shi@daocloud.io>
---
 pkg/device/iluvatar/device_test.go | 143 +++++++++++++++++++++++++++++
 1 file changed, 143 insertions(+)

diff --git a/pkg/device/iluvatar/device_test.go b/pkg/device/iluvatar/device_test.go
index df98b652e..2d88b032c 100644
--- a/pkg/device/iluvatar/device_test.go
+++ b/pkg/device/iluvatar/device_test.go
@@ -17,13 +17,16 @@ limitations under the License.
 package iluvatar
 
 import (
+	"flag"
 	"strconv"
 	"testing"
 	"time"
 
+	"gotest.tools/v3/assert"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/klog/v2"
 
 	"github.com/Project-HAMi/HAMi/pkg/util"
 )
@@ -168,3 +171,143 @@ func TestPatchAnnotations(t *testing.T) {
 		})
 	}
 }
+
+func Test_MutateAdmission(t *testing.T) {
+	tests := []struct {
+		name string
+		args struct {
+			ctr *corev1.Container
+			p   *corev1.Pod
+		}
+		want bool
+		err  error
+	}{
+		{
+			name: "IluvatarResourceCount and IluvatarResourceCores set to limits",
+			args: struct {
+				ctr *corev1.Container
+				p   *corev1.Pod
+			}{
+				ctr: &corev1.Container{
+					Resources: corev1.ResourceRequirements{
+						Limits: corev1.ResourceList{
+							"iluvatar.ai/vgpu":       *resource.NewQuantity(2, resource.DecimalSI),
+							"iluvatar.ai/vcuda-core": *resource.NewQuantity(1, resource.DecimalSI),
+						},
+					},
+				},
+				p: &corev1.Pod{},
+			},
+			want: true,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			config := IluvatarConfig{
+				ResourceCountName:  "iluvatar.ai/vgpu",
+				ResourceCoreName:   "iluvatar.ai/vcuda-core",
+				ResourceMemoryName: "iluvatar.ai/vcuda-memory",
+			}
+			InitIluvatarDevice(config)
+			dev := IluvatarDevices{}
+			result, _ := dev.MutateAdmission(test.args.ctr, test.args.p)
+			assert.Equal(t, result, test.want)
+		})
+	}
+}
+
+func Test_GetNodeDevices(t *testing.T) {
+	tests := []struct {
+		name string
+		args corev1.Node
+		want []*util.DeviceInfo
+	}{
+		{
+			name: "get node devices",
+			args: corev1.Node{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "test",
+				},
+				Status: corev1.NodeStatus{
+					Capacity: corev1.ResourceList{
+						"iluvatar.ai/vcuda-core":   *resource.NewQuantity(1, resource.DecimalSI),
+						"iluvatar.ai/vcuda-memory": *resource.NewQuantity(1000, resource.DecimalSI),
+					},
+				},
+			},
+			want: []*util.DeviceInfo{
+				{
+					Index:   uint(0),
+					ID:      "test-iluvatar-0",
+					Count:   int32(100),
+					Devmem:  int32(25600000),
+					Devcore: int32(100),
+					Type:    IluvatarGPUDevice,
+					Numa:    0,
+					Health:  true,
+				},
+			},
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			dev := IluvatarDevices{}
+			fs := flag.FlagSet{}
+			ParseConfig(&fs)
+			result, err := dev.GetNodeDevices(test.args)
+			if err != nil {
+				klog.InfoS("failed to get node devices info")
+			}
+			assert.DeepEqual(t, result, test.want)
+		})
+	}
+}
+
+func Test_PatchAnnotations(t *testing.T) {
+	tests := []struct {
+		name string
+		args struct {
+			annoinput *map[string]string
+			pd        util.PodDevices
+		}
+		want map[string]string
+	}{
+		{
+			name: "exist device",
+			args: struct {
+				annoinput *map[string]string
+				pd        util.PodDevices
+			}{
+				annoinput: &map[string]string{},
+				pd: util.PodDevices{
+					IluvatarGPUDevice: util.PodSingleDevice{
+						util.ContainerDevices{
+							{
+								Idx:       0,
+								UUID:      "test123",
+								Type:      IluvatarGPUDevice,
+								Usedcores: int32(1),
+								Usedmem:   int32(1000),
+							},
+						},
+					},
+				},
+			},
+			want: map[string]string{
+				util.InRequestDevices[IluvatarGPUDevice]: "test123,Iluvatar,1000,1:;",
+				util.SupportDevices[IluvatarGPUDevice]:   "test123,Iluvatar,1000,1:;",
+				"iluvatar.ai/gpu-assigned":               "false",
+				"iluvatar.ai/predicate-gpu-idx-0":        "0",
+			},
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			dev := IluvatarDevices{}
+			result := dev.PatchAnnotations(test.args.annoinput, test.args.pd)
+			assert.Equal(t, result[IluvatarGPUDevice], test.want[IluvatarGPUDevice])
+			assert.Equal(t, result["iluvatar.ai/gpu-assigned"], test.want["iluvatar.ai/gpu-assigned"])
+			assert.Equal(t, result["iluvatar.ai/predicate-gpu-idx-0"], test.want["iluvatar.ai/predicate-gpu-idx-0"])
+		})
+	}
+}