Refactor tests

Author: pajakd

slice/config/dev/manager_config_patch.yaml

@@ -4,7 +4,7 @@
   value: --zap-devel
 - op: add
   path: /spec/template/spec/containers/0/args/0
-  value: --zap-log-level=5
+  value: --zap-log-level=3
 - op: add
   path: /spec/template/spec/containers/0/args/0
   value: --retry-delay-on-slice-failure=0s

slice/config/rbac/role.yaml

@@ -78,6 +78,16 @@ rules:
   - patch
   - update
   - watch
+- apiGroups:
+  - leaderworkerset.x-k8s.io
+  resources:
+  - leaderworkersets
+  verbs:
+  - get
+  - list
+  - patch
+  - update
+  - watch
 - apiGroups:
   - kueue.x-k8s.io
   resources:
@@ -106,13 +116,3 @@ rules:
   - patch
   - update
   - watch
-- apiGroups:
-  - leaderworkerset.x-k8s.io
-  resources:
-  - leaderworkersets
-  verbs:
-  - get
-  - list
-  - patch
-  - update
-  - watch

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

@@ -182,3 +182,72 @@ func (w *Wrapper) LeaderNodeAffinity(key string, values []string) *Wrapper {
 	)
 	return w
 }
+
+func (w *Wrapper) StartupPolicy(policy leaderworkersetv1.StartupPolicyType) *Wrapper {
+	w.Spec.StartupPolicy = policy
+	return w
+}
+
+func (w *Wrapper) WorkerName(name string) *Wrapper {
+	if len(w.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec.Containers) == 0 {
+		w.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec.Containers = []corev1.Container{{}}
+	}
+	w.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec.Containers[0].Name = name
+	return w
+}
+
+func (w *Wrapper) LeaderName(name string) *Wrapper {
+	if w.Spec.LeaderWorkerTemplate.LeaderTemplate == nil {
+		w.Spec.LeaderWorkerTemplate.LeaderTemplate = &corev1.PodTemplateSpec{}
+	}
+	if len(w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers) == 0 {
+		w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers = []corev1.Container{{}}
+	}
+	w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].Name = name
+	return w
+}
+
+func (w *Wrapper) LeaderImage(img string) *Wrapper {
+	if w.Spec.LeaderWorkerTemplate.LeaderTemplate == nil {
+		w.Spec.LeaderWorkerTemplate.LeaderTemplate = &corev1.PodTemplateSpec{}
+	}
+	if len(w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers) == 0 {
+		w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers = []corev1.Container{{}}
+	}
+	w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].Image = img
+	return w
+}
+
+func (w *Wrapper) LeaderArgs(args ...string) *Wrapper {
+	if w.Spec.LeaderWorkerTemplate.LeaderTemplate == nil {
+		w.Spec.LeaderWorkerTemplate.LeaderTemplate = &corev1.PodTemplateSpec{}
+	}
+	if len(w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers) == 0 {
+		w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers = []corev1.Container{{}}
+	}
+	w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].Args = args
+	return w
+}
+
+func (w *Wrapper) LeaderLimit(resourceName corev1.ResourceName, quantity string) *Wrapper {
+	if w.Spec.LeaderWorkerTemplate.LeaderTemplate == nil {
+		w.Spec.LeaderWorkerTemplate.LeaderTemplate = &corev1.PodTemplateSpec{}
+	}
+	if len(w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers) == 0 {
+		w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers = []corev1.Container{{}}
+	}
+	if w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].Resources.Limits == nil {
+		w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].Resources.Limits = make(corev1.ResourceList)
+	}
+	w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].Resources.Limits[resourceName] = resource.MustParse(quantity)
+	return w
+}
+
+func (w *Wrapper) LeaderRequestAndLimit(resourceName corev1.ResourceName, quantity string) *Wrapper {
+	w.LeaderLimit(resourceName, quantity)
+	if w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].Resources.Requests == nil {
+		w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].Resources.Requests = make(corev1.ResourceList)
+	}
+	w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].Resources.Requests[resourceName] = resource.MustParse(quantity)
+	return w
+}

slice/test/e2e/subslice/lws_test.go

@@ -22,7 +22,6 @@ import (
 	"github.com/onsi/ginkgo/v2"
 	"github.com/onsi/gomega"
 	corev1 "k8s.io/api/core/v1"
-	"k8s.io/apimachinery/pkg/api/resource"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta2"
 	"sigs.k8s.io/kueue/pkg/util/tas"
@@ -109,6 +108,8 @@ var _ = ginkgo.Describe("LWS Subslicing", func() {
 		wrapper := testingjobslws.MakeLeaderWorkerSet(name, ns.Name).
 			Queue(lq.Name).
 			Size(tc.size).
+			StartupPolicy(leaderworkersetv1.LeaderCreatedStartupPolicy).
+			WorkerName("worker").
 			WorkerImage(utils.E2eTestAgnHostImage).
 			WorkerArgs(utils.BehaviorWaitForDeletion...).
 			WorkerAnnotation(core.TPUSliceTopologyAnnotation, tc.topology).
@@ -117,34 +118,18 @@ var _ = ginkgo.Describe("LWS Subslicing", func() {
 
 		if tc.withLeader {
 			wrapper = wrapper.
+				LeaderName("leader").
+				LeaderImage(utils.E2eTestAgnHostImage).
+				LeaderArgs(utils.BehaviorWaitForDeletion...).
 				LeaderAnnotation(core.TPUSliceTopologyAnnotation, tc.topology).
 				LeaderNodeSelector("cloud.google.com/gke-tpu-accelerator", string(slice.TypeTpu7x))
-		}
-
-		lws := wrapper.Obj()
-		lws.Spec.StartupPolicy = leaderworkersetv1.LeaderCreatedStartupPolicy
-		if len(lws.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec.Containers) > 0 {
-			lws.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec.Containers[0].Name = "worker"
-		}
-		if tc.withLeader {
-			container := corev1.Container{
-				Name:  "leader",
-				Image: utils.E2eTestAgnHostImage,
-				Args:  utils.BehaviorWaitForDeletion,
-			}
 			if tc.leaderRequiresTPUs {
-				container.Resources = corev1.ResourceRequirements{
-					Limits: corev1.ResourceList{
-						core.TPUResourceName: resource.MustParse("4"),
-					},
-					Requests: corev1.ResourceList{
-						core.TPUResourceName: resource.MustParse("4"),
-					},
-				}
+				wrapper = wrapper.LeaderRequestAndLimit(core.TPUResourceName, "4")
 			}
-			lws.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers = []corev1.Container{container}
 		}
 
+		lws := wrapper.Obj()
+
 		ginkgo.By("Creating a LeaderWorkerSet", func() {
 			utils.MustCreate(ctx, k8sClient, lws)
 		})