Improve test

Signed-off-by: GonzaloSaez <[email protected]>

test/integration/mpi_job_controller_test.go

@@ -259,9 +259,37 @@ func testMpiJobWaitWorkers(t *testing.T, startSuspended bool) {
 	}
 	s.events.verify(t)
 
-	workerPods, err := getPodsForJob(ctx, s.kClient, mpiJob)
+	// The launcher job should not be created until all workers are ready even when we start in suspended mode.
+	job, err := getLauncherJobForMPIJob(ctx, s.kClient, mpiJob)
 	if err != nil {
-		t.Fatalf("Cannot get worker pods from job: %v", err)
+		t.Fatalf("Cannot get launcher job from job: %v", err)
+	}
+	if job != nil {
+		t.Fatalf("Launcher is created before workers")
+	}
+
+	if startSuspended {
+		// Resume the MPIJob so that the test can follow the normal path.
+		mpiJob.Spec.RunPolicy.Suspend = ptr.To(false)
+		mpiJob, err = s.mpiClient.KubeflowV2beta1().MPIJobs(mpiJob.Namespace).Update(ctx, mpiJob, metav1.UpdateOptions{})
+		if err != nil {
+			t.Fatalf("Error Updating MPIJob: %v", err)
+		}
+	}
+
+	var workerPods []corev1.Pod
+	if err = wait.PollUntilContextTimeout(ctx, util.WaitInterval, wait.ForeverTestTimeout, false, func(ctx context.Context) (bool, error) {
+		var err error
+		workerPods, err = getPodsForJob(ctx, s.kClient, mpiJob)
+		if err != nil {
+			return false, err
+		}
+		if len(workerPods) != 2 {
+			return false, nil
+		}
+		return true, nil
+	}); err != nil {
+		t.Errorf("Failed updating scheduler-plugins PodGroup: %v", err)
 	}
 
 	err = updatePodsToPhase(ctx, s.kClient, workerPods, corev1.PodRunning)