[slice] Remove duplicated test code (#1103)

* Remove duplicated test code

* Remove reduntand code

Author: pajakd

slice/internal/util/testingjobs/job/wrappers.go

@@ -23,6 +23,8 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/utils/ptr"
 	kueueconstants "sigs.k8s.io/kueue/pkg/controller/constants"
+
+	"tpu-slice-controller/internal/core"
 )
 
 // JobWrapper wraps a Job.
@@ -125,3 +127,8 @@ func (j *JobWrapper) Args(args ...string) *JobWrapper {
 	j.Spec.Template.Spec.Containers[0].Args = args
 	return j
 }
+
+func (j *JobWrapper) NodeAffinity(key string, values []string) *JobWrapper {
+	core.AddNodeAffinity(&j.Spec.Template, key, corev1.NodeSelectorOpIn, values)
+	return j
+}

slice/internal/webhooks/job_webhook_test.go

@@ -22,92 +22,14 @@ import (
 
 	"github.com/google/go-cmp/cmp"
 	batchv1 "k8s.io/api/batch/v1"
-	corev1 "k8s.io/api/core/v1"
-	"k8s.io/apimachinery/pkg/api/resource"
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-	"k8s.io/utils/ptr"
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
-	kueueconstants "sigs.k8s.io/kueue/pkg/controller/constants"
 
 	slice "tpu-slice-controller/api/v1beta1"
 	"tpu-slice-controller/internal/core"
 	utiltesting "tpu-slice-controller/internal/util/testing"
+	testingjob "tpu-slice-controller/internal/util/testingjobs/job"
 )
 
-type JobWrapper struct {
-	batchv1.Job
-}
-
-func MakeJob(name, ns string) *JobWrapper {
-	return &JobWrapper{
-		Job: batchv1.Job{
-			ObjectMeta: metav1.ObjectMeta{
-				Name:      name,
-				Namespace: ns,
-			},
-			Spec: batchv1.JobSpec{
-				Template: corev1.PodTemplateSpec{
-					Spec: corev1.PodSpec{
-						RestartPolicy: corev1.RestartPolicyNever,
-						Containers: []corev1.Container{
-							{
-								Name:  "c",
-								Image: "pause",
-							},
-						},
-					},
-				},
-			},
-		},
-	}
-}
-
-func (j *JobWrapper) Obj() *batchv1.Job {
-	return &j.Job
-}
-
-func (j *JobWrapper) Queue(q string) *JobWrapper {
-	if j.Labels == nil {
-		j.Labels = make(map[string]string)
-	}
-	j.Labels[kueueconstants.QueueLabel] = q
-	return j
-}
-
-func (j *JobWrapper) Parallelism(p int32) *JobWrapper {
-	j.Spec.Parallelism = ptr.To(p)
-	return j
-}
-
-func (j *JobWrapper) PodAnnotation(k, v string) *JobWrapper {
-	if j.Spec.Template.Annotations == nil {
-		j.Spec.Template.Annotations = make(map[string]string)
-	}
-	j.Spec.Template.Annotations[k] = v
-	return j
-}
-
-func (j *JobWrapper) NodeSelector(k, v string) *JobWrapper {
-	if j.Spec.Template.Spec.NodeSelector == nil {
-		j.Spec.Template.Spec.NodeSelector = make(map[string]string)
-	}
-	j.Spec.Template.Spec.NodeSelector[k] = v
-	return j
-}
-
-func (j *JobWrapper) Request(r corev1.ResourceName, v string) *JobWrapper {
-	if j.Spec.Template.Spec.Containers[0].Resources.Limits == nil {
-		j.Spec.Template.Spec.Containers[0].Resources.Limits = make(corev1.ResourceList)
-	}
-	j.Spec.Template.Spec.Containers[0].Resources.Limits[r] = resource.MustParse(v)
-	return j
-}
-
-func (j *JobWrapper) NodeAffinity(key string, values []string) *JobWrapper {
-	core.AddNodeAffinity(&j.Spec.Template, key, corev1.NodeSelectorOpIn, values)
-	return j
-}
-
 func TestJobDefault(t *testing.T) {
 	const (
 		baseJobName   = "job"
@@ -120,71 +42,71 @@ func TestJobDefault(t *testing.T) {
 		wantErr error
 	}{
 		"no queue label": {
-			job: MakeJob(baseJobName, baseNamespace).
+			job: testingjob.MakeJob(baseJobName, baseNamespace).
 				Parallelism(48).
 				PodAnnotation(core.TPUSliceTopologyAnnotation, "4x4x12").
 				NodeSelector("cloud.google.com/gke-tpu-accelerator", string(slice.TypeTpu7x)).
 				Obj(),
-			wantJob: MakeJob(baseJobName, baseNamespace).
+			wantJob: testingjob.MakeJob(baseJobName, baseNamespace).
 				Parallelism(48).
 				PodAnnotation(core.TPUSliceTopologyAnnotation, "4x4x12").
 				NodeSelector("cloud.google.com/gke-tpu-accelerator", string(slice.TypeTpu7x)).
 				Obj(),
 		},
 		"no tpu topology annotation": {
-			job: MakeJob(baseJobName, baseNamespace).
+			job: testingjob.MakeJob(baseJobName, baseNamespace).
 				Queue("queue-name").
 				Parallelism(48).
 				NodeSelector("cloud.google.com/gke-tpu-accelerator", string(slice.TypeTpu7x)).
 				Obj(),
-			wantJob: MakeJob(baseJobName, baseNamespace).
+			wantJob: testingjob.MakeJob(baseJobName, baseNamespace).
 				Queue("queue-name").
 				Parallelism(48).
 				NodeSelector("cloud.google.com/gke-tpu-accelerator", string(slice.TypeTpu7x)).
 				Obj(),
 		},
 		"no tpu accelerator node selector label": {
-			job: MakeJob(baseJobName, baseNamespace).
+			job: testingjob.MakeJob(baseJobName, baseNamespace).
 				Queue("queue-name").
 				Parallelism(48).
 				PodAnnotation(core.TPUSliceTopologyAnnotation, "4x4x12").
 				Obj(),
-			wantJob: MakeJob(baseJobName, baseNamespace).
+			wantJob: testingjob.MakeJob(baseJobName, baseNamespace).
 				Queue("queue-name").
 				Parallelism(48).
 				PodAnnotation(core.TPUSliceTopologyAnnotation, "4x4x12").
 				Obj(),
 		},
 		"should set default values": {
-			job: MakeJob(baseJobName, baseNamespace).
+			job: testingjob.MakeJob(baseJobName, baseNamespace).
 				Queue("queue-name").
 				Parallelism(48).
 				PodAnnotation(core.TPUSliceTopologyAnnotation, "4x4x12").
 				NodeSelector("cloud.google.com/gke-tpu-accelerator", string(slice.TypeTpu7x)).
-				Request(core.TPUResourceName, "4").
+				Limit(core.TPUResourceName, "4").
 				Obj(),
-			wantJob: MakeJob(baseJobName, baseNamespace).
+			wantJob: testingjob.MakeJob(baseJobName, baseNamespace).
 				Queue("queue-name").
 				Parallelism(48).
 				PodAnnotation(core.TPUSliceTopologyAnnotation, "4x4x12").
 				PodAnnotation(kueue.PodSetRequiredTopologyAnnotation, core.TPUBlockLabel).
 				PodAnnotation(kueue.PodSetSliceRequiredTopologyAnnotation, core.TPUSubBlockLabel).
 				PodAnnotation(kueue.PodSetSliceSizeAnnotation, "16").
 				NodeSelector("cloud.google.com/gke-tpu-accelerator", string(slice.TypeTpu7x)).
-				Request(core.TPUResourceName, "4").
+				Limit(core.TPUResourceName, "4").
 				NodeAffinity(core.TPUSliceHealthNodeSelectorKey, []string{core.TPUSliceHealthNodeSelectorHealthy}).
 				Obj(),
 		},
 		"should respect existing NodeSelector for health": {
-			job: MakeJob(baseJobName, baseNamespace).
+			job: testingjob.MakeJob(baseJobName, baseNamespace).
 				Queue("queue-name").
 				Parallelism(48).
 				PodAnnotation(core.TPUSliceTopologyAnnotation, "4x4x12").
 				NodeSelector("cloud.google.com/gke-tpu-accelerator", string(slice.TypeTpu7x)).
 				NodeSelector(core.TPUSliceHealthNodeSelectorKey, "HEALTHY").
-				Request(core.TPUResourceName, "4").
+				Limit(core.TPUResourceName, "4").
 				Obj(),
-			wantJob: MakeJob(baseJobName, baseNamespace).
+			wantJob: testingjob.MakeJob(baseJobName, baseNamespace).
 				Queue("queue-name").
 				Parallelism(48).
 				PodAnnotation(core.TPUSliceTopologyAnnotation, "4x4x12").
@@ -193,19 +115,19 @@ func TestJobDefault(t *testing.T) {
 				PodAnnotation(kueue.PodSetSliceSizeAnnotation, "16").
 				NodeSelector("cloud.google.com/gke-tpu-accelerator", string(slice.TypeTpu7x)).
 				NodeSelector(core.TPUSliceHealthNodeSelectorKey, "HEALTHY").
-				Request(core.TPUResourceName, "4").
+				Limit(core.TPUResourceName, "4").
 				Obj(),
 		},
 		"should respect existing NodeAffinity for health": {
-			job: MakeJob(baseJobName, baseNamespace).
+			job: testingjob.MakeJob(baseJobName, baseNamespace).
 				Queue("queue-name").
 				Parallelism(48).
 				PodAnnotation(core.TPUSliceTopologyAnnotation, "4x4x12").
 				NodeSelector("cloud.google.com/gke-tpu-accelerator", string(slice.TypeTpu7x)).
 				NodeAffinity(core.TPUSliceHealthNodeSelectorKey, []string{"HEALTHY", "DEGRADED"}).
-				Request(core.TPUResourceName, "4").
+				Limit(core.TPUResourceName, "4").
 				Obj(),
-			wantJob: MakeJob(baseJobName, baseNamespace).
+			wantJob: testingjob.MakeJob(baseJobName, baseNamespace).
 				Queue("queue-name").
 				Parallelism(48).
 				PodAnnotation(core.TPUSliceTopologyAnnotation, "4x4x12").
@@ -214,16 +136,16 @@ func TestJobDefault(t *testing.T) {
 				PodAnnotation(kueue.PodSetSliceSizeAnnotation, "16").
 				NodeSelector("cloud.google.com/gke-tpu-accelerator", string(slice.TypeTpu7x)).
 				NodeAffinity(core.TPUSliceHealthNodeSelectorKey, []string{"HEALTHY", "DEGRADED"}).
-				Request(core.TPUResourceName, "4").
+				Limit(core.TPUResourceName, "4").
 				Obj(),
 		},
 		"should reject incorrectly configured job not utilizing entire cube; single cube": {
-			job: MakeJob(baseJobName, baseNamespace).
+			job: testingjob.MakeJob(baseJobName, baseNamespace).
 				Queue("queue-name").
 				Parallelism(16).
 				PodAnnotation(core.TPUSliceTopologyAnnotation, "4x4x4").
 				NodeSelector("cloud.google.com/gke-tpu-accelerator", string(slice.TypeTpu7x)).
-				Request(core.TPUResourceName, "1").
+				Limit(core.TPUResourceName, "1").
 				Obj(),
 			wantErr: errors.New("invalid job \"job\": configuration results in 16 TPUs requested per cube, but must be exactly 64 TPUs (full utilization)"),
 		},

slice/test/utils/e2e.go

@@ -120,12 +120,12 @@ func waitForOperatorAvailability(ctx context.Context, k8sClient client.Client, k
 	ginkgo.GinkgoLogr.Info("Deployment is available in the cluster", "deployment", key, "waitingTime", time.Since(waitForAvailableStart))
 }
 
-func LabelNodesWithTopology(ctx context.Context, k8sClient client.Client, topology string, nodeNames []string) {
+func LabelNodes(ctx context.Context, k8sClient client.Client, key, value string, nodeNames []string) {
 	for _, nodeName := range nodeNames {
 		gomega.Eventually(func(g gomega.Gomega) {
 			n := &corev1.Node{}
 			g.Expect(k8sClient.Get(ctx, client.ObjectKey{Name: nodeName}, n)).To(gomega.Succeed())
-			metav1.SetMetaDataLabel(&n.ObjectMeta, core.TPUTopologyAnnotation, topology)
+			metav1.SetMetaDataLabel(&n.ObjectMeta, key, value)
 			g.Expect(k8sClient.Update(ctx, n)).To(gomega.Succeed())
 		}, Timeout, Interval).Should(gomega.Succeed())
 	}
@@ -134,7 +134,7 @@ func LabelNodesWithTopology(ctx context.Context, k8sClient client.Client, topolo
 			gomega.Eventually(func(g gomega.Gomega) {
 				n := &corev1.Node{}
 				g.Expect(k8sClient.Get(ctx, client.ObjectKey{Name: nodeName}, n)).To(gomega.Succeed())
-				delete(n.Labels, core.TPUTopologyAnnotation)
+				delete(n.Labels, key)
 				g.Expect(k8sClient.Update(ctx, n)).To(gomega.Succeed())
 			}, Timeout, Interval).Should(gomega.Succeed())
 		}
@@ -166,45 +166,27 @@ func SetSliceReady(ctx context.Context, k8sClient client.Client, sliceKey client
 			nodeNames = append(nodeNames, node.Name)
 		}
 	}
-	LabelNodesWithTopology(ctx, k8sClient, topology, nodeNames)
-	LabelNodesWithSlice(ctx, k8sClient, createdSlice.Name, nodeNames)
+	LabelNodes(ctx, k8sClient, core.TPUTopologyAnnotation, topology, nodeNames)
+	LabelNodes(ctx, k8sClient, core.TPUSliceNodeLabel, createdSlice.Name, nodeNames)
 }
 
-func ExpectJobSetHasAntiAffinity(js *jobset.JobSet) {
+func expectAntiAffinity(template *corev1.PodTemplateSpec, name, kind string) {
 	ginkgo.GinkgoHelper()
-	for _, replicatedJob := range js.Spec.ReplicatedJobs {
-		req := core.FindNodeAffinityRequirement(&replicatedJob.Template.Spec.Template, core.TPUSliceNodeLabel)
-		gomega.ExpectWithOffset(1, req).ShouldNot(gomega.BeNil(), "replicated job %q should have anti-affinity for key %s", replicatedJob.Name, core.TPUSliceNodeLabel)
-		gomega.ExpectWithOffset(1, req.Operator).Should(gomega.Equal(corev1.NodeSelectorOpDoesNotExist))
-	}
+	req := core.FindNodeAffinityRequirement(template, core.TPUSliceNodeLabel)
+	gomega.Expect(req).ShouldNot(gomega.BeNil(), "%s %q should have anti-affinity for key %s", kind, name, core.TPUSliceNodeLabel)
+	gomega.Expect(req.Operator).Should(gomega.Equal(corev1.NodeSelectorOpDoesNotExist))
 }
 
-func LabelNodesWithSlice(ctx context.Context, k8sClient client.Client, sliceName string, nodeNames []string) {
-	for _, nodeName := range nodeNames {
-		gomega.Eventually(func(g gomega.Gomega) {
-			n := &corev1.Node{}
-			g.Expect(k8sClient.Get(ctx, client.ObjectKey{Name: nodeName}, n)).To(gomega.Succeed())
-			metav1.SetMetaDataLabel(&n.ObjectMeta, core.TPUSliceNodeLabel, sliceName)
-			g.Expect(k8sClient.Update(ctx, n)).To(gomega.Succeed())
-		}, Timeout, Interval).Should(gomega.Succeed())
+func ExpectJobSetHasAntiAffinity(js *jobset.JobSet) {
+	ginkgo.GinkgoHelper()
+	for _, replicatedJob := range js.Spec.ReplicatedJobs {
+		expectAntiAffinity(&replicatedJob.Template.Spec.Template, replicatedJob.Name, "replicated job")
 	}
-	ginkgo.DeferCleanup(func(ctx context.Context, k8sClient client.Client, nodeNames []string) {
-		for _, nodeName := range nodeNames {
-			gomega.Eventually(func(g gomega.Gomega) {
-				n := &corev1.Node{}
-				g.Expect(k8sClient.Get(ctx, client.ObjectKey{Name: nodeName}, n)).To(gomega.Succeed())
-				delete(n.Labels, core.TPUSliceNodeLabel)
-				g.Expect(k8sClient.Update(ctx, n)).To(gomega.Succeed())
-			}, Timeout, Interval).Should(gomega.Succeed())
-		}
-	}, k8sClient, nodeNames)
 }
 
 func ExpectWorkloadHasAntiAffinity(wl *kueue.Workload) {
 	ginkgo.GinkgoHelper()
 	for _, ps := range wl.Spec.PodSets {
-		req := core.FindNodeAffinityRequirement(&ps.Template, core.TPUSliceNodeLabel)
-		gomega.ExpectWithOffset(1, req).ShouldNot(gomega.BeNil(), "podset %q should have anti-affinity for key %s", ps.Name, core.TPUSliceNodeLabel)
-		gomega.ExpectWithOffset(1, req.Operator).Should(gomega.Equal(corev1.NodeSelectorOpDoesNotExist))
+		expectAntiAffinity(&ps.Template, string(ps.Name), "podset")
 	}
 }