integration tests update

kubernetes-sigs · Jul 19, 2024 · 116865c · 116865c
1 parent a074860
commit 116865c
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 3 deletions.
diff --git a/test/integration/controller/jobs/mxjob/mxjob_controller_test.go b/test/integration/controller/jobs/mxjob/mxjob_controller_test.go
@@ -129,7 +129,26 @@ var _ = ginkgo.Describe("Job controller when waitForPodsReady enabled", ginkgo.O
 
 	ginkgo.DescribeTable("Single job at different stages of progress towards completion",
 		func(podsReadyTestSpec kftesting.PodsReadyTestSpec) {
-			kfJob := kubeflowjob.KubeflowJob{KFJobControl: (*workloadmxjob.JobControl)(testingmxjob.MakeMXJob(jobName, ns.Name).Parallelism(2, 2).Obj())}
+			kfJob := kubeflowjob.KubeflowJob{
+				KFJobControl: (*workloadmxjob.JobControl)(
+					testingmxjob.MakeMXJob(jobName, ns.Name).
+						MXReplicaSpecs(
+							testingmxjob.MXReplicaSpecRequirement{
+								ReplicaType:  kftraining.MXJobReplicaTypeScheduler,
+								ReplicaCount: 1,
+							},
+							testingmxjob.MXReplicaSpecRequirement{
+								ReplicaType:  kftraining.MXJobReplicaTypeServer,
+								ReplicaCount: 1,
+							},
+							testingmxjob.MXReplicaSpecRequirement{
+								ReplicaType:  kftraining.MXJobReplicaTypeWorker,
+								ReplicaCount: 1,
+							},
+						).
+						Parallelism(2, 2).
+						Obj()),
+			}
 			createdJob := kubeflowjob.KubeflowJob{KFJobControl: (*workloadmxjob.JobControl)(&kftraining.MXJob{})}
 
 			kftesting.JobControllerWhenWaitForPodsReadyEnabled(ctx, k8sClient, kfJob, createdJob, podsReadyTestSpec, []kftesting.PodSetsResource{
@@ -287,6 +306,7 @@ var _ = ginkgo.Describe("Job controller interacting with scheduler", ginkgo.Orde
 
 		kfJob := kubeflowjob.KubeflowJob{KFJobControl: (*workloadmxjob.JobControl)(
 			testingmxjob.MakeMXJob(jobName, ns.Name).Queue(localQueue.Name).
+				MXReplicaSpecsDefault().
 				Request(kftraining.MXJobReplicaTypeScheduler, corev1.ResourceCPU, "3").
 				Request(kftraining.MXJobReplicaTypeServer, corev1.ResourceCPU, "4").
 				Request(kftraining.MXJobReplicaTypeWorker, corev1.ResourceCPU, "4").

diff --git a/test/integration/controller/jobs/paddlejob/paddlejob_controller_test.go b/test/integration/controller/jobs/paddlejob/paddlejob_controller_test.go
@@ -75,7 +75,7 @@ var _ = ginkgo.Describe("Job controller", ginkgo.Ordered, ginkgo.ContinueOnFailu
 	})
 
 	ginkgo.It("Should reconcile PaddleJobs", func() {
-		kfJob := kubeflowjob.KubeflowJob{KFJobControl: (*workloadpaddlejob.JobControl)(testingpaddlejob.MakePaddleJob(jobName, ns.Name).Obj())}
+		kfJob := kubeflowjob.KubeflowJob{KFJobControl: (*workloadpaddlejob.JobControl)(testingpaddlejob.MakePaddleJob(jobName, ns.Name).PaddleReplicaSpecsDefault().Obj())}
 		createdJob := kubeflowjob.KubeflowJob{KFJobControl: (*workloadpaddlejob.JobControl)(&kftraining.PaddleJob{})}
 
 		kftesting.ShouldReconcileJob(ctx, k8sClient, kfJob, createdJob, []kftesting.PodSetsResource{
@@ -127,7 +127,7 @@ var _ = ginkgo.Describe("Job controller when waitForPodsReady enabled", ginkgo.O
 
 	ginkgo.DescribeTable("Single job at different stages of progress towards completion",
 		func(podsReadyTestSpec kftesting.PodsReadyTestSpec) {
-			kfJob := kubeflowjob.KubeflowJob{KFJobControl: (*workloadpaddlejob.JobControl)(testingpaddlejob.MakePaddleJob(jobName, ns.Name).Parallelism(2).Obj())}
+			kfJob := kubeflowjob.KubeflowJob{KFJobControl: (*workloadpaddlejob.JobControl)(testingpaddlejob.MakePaddleJob(jobName, ns.Name).PaddleReplicaSpecsDefault().Parallelism(2).Obj())}
 			createdJob := kubeflowjob.KubeflowJob{KFJobControl: (*workloadpaddlejob.JobControl)(&kftraining.PaddleJob{})}
 
 			kftesting.JobControllerWhenWaitForPodsReadyEnabled(ctx, k8sClient, kfJob, createdJob, podsReadyTestSpec, []kftesting.PodSetsResource{
@@ -281,6 +281,7 @@ var _ = ginkgo.Describe("Job controller interacting with scheduler", ginkgo.Orde
 
 		kfJob := kubeflowjob.KubeflowJob{KFJobControl: (*workloadpaddlejob.JobControl)(
 			testingpaddlejob.MakePaddleJob(jobName, ns.Name).Queue(localQueue.Name).
+				PaddleReplicaSpecsDefault().
 				Request(kftraining.PaddleJobReplicaTypeMaster, corev1.ResourceCPU, "3").
 				Request(kftraining.PaddleJobReplicaTypeWorker, corev1.ResourceCPU, "4").
 				Obj(),