refactor the unit tests for all jobs

kubeflow · May 5, 2022 · bcbf9ba · bcbf9ba
1 parent a87ed14
commit bcbf9ba
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 86 deletions.
diff --git a/pkg/common/util/v1/testutil/tfjob.go b/pkg/common/util/v1/testutil/tfjob.go
@@ -245,29 +245,3 @@ func SetTFJobCompletionTime(tfJob *tfv1.TFJob) {
 	now := metav1.Time{Time: time.Now()}
 	tfJob.Status.CompletionTime = &now
 }
-
-func NewInvalidTFJobWithNoContainerNamedTensorflow(name string) *tfv1.TFJob {
-	tfJob := &tfv1.TFJob{
-		ObjectMeta: metav1.ObjectMeta{
-			Name:      name,
-			Namespace: metav1.NamespaceDefault,
-		},
-		Spec: tfv1.TFJobSpec{
-			TFReplicaSpecs: map[commonv1.ReplicaType]*commonv1.ReplicaSpec{
-				tfv1.TFReplicaTypeWorker: {
-					Template: v1.PodTemplateSpec{
-						Spec: v1.PodSpec{
-							Containers: []v1.Container{
-								{
-									Name:  "",
-									Image: "kubeflow/tf-dist-mnist-test:1.0",
-								},
-							},
-						},
-					},
-				},
-			},
-		},
-	}
-	return tfJob
-}
diff --git a/pkg/controller.v1/mpi/mpijob_controller_test.go b/pkg/controller.v1/mpi/mpijob_controller_test.go
@@ -128,37 +128,6 @@ func newMPIJobWithLauncher(name string, replicas *int32, pusPerReplica int64, re
 	return mpiJob
 }
 
-func newInvalidMPIJobWithUndefinedContainerName(name string) *kubeflow.MPIJob {
-	cleanPodPolicyAll := common.CleanPodPolicyAll
-	mpiJob := &kubeflow.MPIJob{
-		TypeMeta: metav1.TypeMeta{APIVersion: kubeflow.SchemeGroupVersion.String()},
-		ObjectMeta: metav1.ObjectMeta{
-			Name:      name,
-			Namespace: metav1.NamespaceDefault,
-		},
-		Spec: kubeflow.MPIJobSpec{
-			RunPolicy: common.RunPolicy{
-				CleanPodPolicy: &cleanPodPolicyAll,
-			},
-			MPIReplicaSpecs: map[common.ReplicaType]*common.ReplicaSpec{
-				kubeflow.MPIReplicaTypeLauncher: &common.ReplicaSpec{
-					Template: corev1.PodTemplateSpec{
-						Spec: corev1.PodSpec{
-							Containers: []corev1.Container{
-								corev1.Container{
-									Name:  "",
-									Image: "kubeflow/tf-dist-mnist-test:1.0",
-								},
-							},
-						},
-					},
-				},
-			},
-		},
-	}
-	return mpiJob
-}
-
 var _ = Describe("MPIJob controller", func() {
 	// Define utility constants for object names and testing timeouts/durations and intervals.
 	const (
@@ -744,10 +713,14 @@ var _ = Describe("MPIJob controller", func() {
 		It("Should return error", func() {
 			By("Calling Reconcile method")
 			ctx := context.Background()
+			startTime := metav1.Now()
+			completionTime := metav1.Now()
 
 			jobName := "test-invalid-job-spec"
+			mpiJob := newMPIJobWithLauncher(jobName, int32Ptr(1), 1, gpuResourceName, &startTime, &completionTime)
+			mpiJob.Spec.MPIReplicaSpecs[kubeflow.MPIReplicaTypeLauncher].Template.Spec.Containers[0].Name = ""
+			mpiJob.Spec.MPIReplicaSpecs[kubeflow.MPIReplicaTypeWorker].Template.Spec.Containers[0].Name = ""
 
-			mpiJob := newInvalidMPIJobWithUndefinedContainerName(jobName)
 			Expect(testK8sClient.Create(ctx, mpiJob)).Should(Succeed())
 
 			req := ctrl.Request{NamespacedName: types.NamespacedName{

diff --git a/pkg/controller.v1/pytorch/pytorchjob_controller_test.go b/pkg/controller.v1/pytorch/pytorchjob_controller_test.go
@@ -304,7 +304,30 @@ var _ = Describe("PyTorchJob controller", func() {
 
 			jobName := "test-invalid-job-spec"
 
-			pytorchJob := newInvalidPyTorchJobWithNoContainerNamedForTest(jobName, metav1.NamespaceDefault)
+			pytorchJob := newPyTorchJobForTest(jobName, metav1.NamespaceDefault)
+			pytorchJob.Spec.PyTorchReplicaSpecs = map[commonv1.ReplicaType]*commonv1.ReplicaSpec{
+				pytorchv1.PyTorchReplicaTypeWorker: {
+					Replicas: int32Ptr(1),
+					Template: corev1.PodTemplateSpec{
+						Spec: corev1.PodSpec{
+							Containers: []corev1.Container{
+								{
+									Image: "test-image",
+									Name:  "",
+									Ports: []corev1.ContainerPort{
+										{
+											Name:          pytorchv1.DefaultPortName,
+											ContainerPort: expectedPort,
+											Protocol:      corev1.ProtocolTCP,
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			}
+
 			Expect(testK8sClient.Create(ctx, pytorchJob)).Should(Succeed())
 
 			req := ctrl.Request{NamespacedName: types.NamespacedName{
@@ -340,28 +363,3 @@ func getCondition(status commonv1.JobStatus, condType commonv1.JobConditionType)
 	}
 	return nil
 }
-
-func newInvalidPyTorchJobWithNoContainerNamedForTest(name, namespace string) *pytorchv1.PyTorchJob {
-	return &pytorchv1.PyTorchJob{
-		ObjectMeta: metav1.ObjectMeta{
-			Name:      name,
-			Namespace: namespace,
-		},
-		Spec: pytorchv1.PyTorchJobSpec{
-			PyTorchReplicaSpecs: map[commonv1.ReplicaType]*commonv1.ReplicaSpec{
-				pytorchv1.PyTorchReplicaTypeWorker: {
-					Template: corev1.PodTemplateSpec{
-						Spec: corev1.PodSpec{
-							Containers: []corev1.Container{
-								{
-									Name:  "",
-									Image: "gcr.io/kubeflow-ci/pytorch-dist-mnist_test:1.0",
-								},
-							},
-						},
-					},
-				},
-			},
-		},
-	}
-}
diff --git a/pkg/controller.v1/tensorflow/tfjob_controller_test.go b/pkg/controller.v1/tensorflow/tfjob_controller_test.go
@@ -328,9 +328,12 @@ var _ = Describe("TFJob controller", func() {
 			By("Calling Reconcile method")
 			ctx := context.Background()
 
-			jobName := "test-invalid-job-spec"
+			workerNum, psNum := 1, 1
+			tfJob := testutil.NewTFJob(workerNum, psNum)
+			tfJob.Name = "test-invalid-job-spec"
+			tfJob.Spec.TFReplicaSpecs[tfv1.TFReplicaTypePS].Template.Spec.Containers[0].Name = ""
+			tfJob.Spec.TFReplicaSpecs[tfv1.TFReplicaTypeWorker].Template.Spec.Containers[0].Name = ""
 
-			tfJob := testutil.NewInvalidTFJobWithNoContainerNamedTensorflow(jobName)
 			Expect(testK8sClient.Create(ctx, tfJob)).Should(Succeed())
 
 			req := ctrl.Request{NamespacedName: types.NamespacedName{