pods: Add cluster spec test (kubeflow#649)

Signed-off-by: Ce Gao <gaoce@caicloud.io>
jetmuffin · Jun 13, 2018 · 3acc05b · 3acc05b
1 parent 157d0f2
commit 3acc05b
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 17 deletions.
diff --git a/pkg/controller.v2/controller_pod.go b/pkg/controller.v2/controller_pod.go
@@ -148,26 +148,10 @@ func (tc *TFJobController) createNewPod(tfjob *tfv1alpha2.TFJob, rt, index strin
 		podTemplate.Labels[key] = value
 	}
 
-	// Generate TF_CONFIG JSON string.
-	tfConfigStr, err := genTFConfigJSONStr(tfjob, rt, index)
-	if err != nil {
+	if err := setClusterSpec(podTemplate, tfjob, rt, index); err != nil {
 		return err
 	}
 
-	if tfConfigStr == "" {
-		return nil
-	}
-	// Add TF_CONFIG environment variable.
-	for i := range podTemplate.Spec.Containers {
-		if len(podTemplate.Spec.Containers[i].Env) == 0 {
-			podTemplate.Spec.Containers[i].Env = make([]v1.EnvVar, 0)
-		}
-		podTemplate.Spec.Containers[i].Env = append(podTemplate.Spec.Containers[i].Env, v1.EnvVar{
-			Name:  tfConfig,
-			Value: tfConfigStr,
-		})
-	}
-
 	if spec.RestartPolicy == tfv1alpha2.RestartPolicyExitCode {
 		podTemplate.Spec.RestartPolicy = v1.RestartPolicyNever
 	} else if spec.RestartPolicy == tfv1alpha2.RestartPolicy("") {
@@ -193,6 +177,29 @@ func (tc *TFJobController) createNewPod(tfjob *tfv1alpha2.TFJob, rt, index strin
 	return nil
 }
 
+func setClusterSpec(podTemplateSpec *v1.PodTemplateSpec, tfjob *tfv1alpha2.TFJob, rt, index string) error {
+	// Generate TF_CONFIG JSON string.
+	tfConfigStr, err := genTFConfigJSONStr(tfjob, rt, index)
+	if err != nil {
+		return err
+	}
+
+	if tfConfigStr == "" {
+		return nil
+	}
+	// Add TF_CONFIG environment variable.
+	for i := range podTemplateSpec.Spec.Containers {
+		if len(podTemplateSpec.Spec.Containers[i].Env) == 0 {
+			podTemplateSpec.Spec.Containers[i].Env = make([]v1.EnvVar, 0)
+		}
+		podTemplateSpec.Spec.Containers[i].Env = append(podTemplateSpec.Spec.Containers[i].Env, v1.EnvVar{
+			Name:  tfConfig,
+			Value: tfConfigStr,
+		})
+	}
+	return nil
+}
+
 // getPodsForTFJob returns the set of pods that this tfjob should manage.
 // It also reconciles ControllerRef by adopting/orphaning.
 // Note that the returned Pods are pointers into the cache.

diff --git a/pkg/controller.v2/controller_pod_test.go b/pkg/controller.v2/controller_pod_test.go
@@ -132,3 +132,39 @@ func TestAddPod(t *testing.T) {
 	}
 	close(stopCh)
 }
+
+func TestClusterSpec(t *testing.T) {
+	type tc struct {
+		tfJob               *tfv1alpha2.TFJob
+		rt                  string
+		index               string
+		expectedClusterSpec string
+	}
+	testCase := []tc{
+		tc{
+			tfJob: newTFJob(1, 0),
+			rt:    "worker",
+			index: "0",
+			expectedClusterSpec: `{"cluster":{"worker":["default-` + testTFJobName +
+				`-worker-0.default.svc.cluster.local:2222"]},"task":{"type":"worker","index":0}}`,
+		},
+		tc{
+			tfJob: newTFJob(1, 1),
+			rt:    "worker",
+			index: "0",
+			expectedClusterSpec: `{"cluster":{"ps":["default-` + testTFJobName +
+				`-ps-0.default.svc.cluster.local:2222"],"worker":["default-` + testTFJobName +
+				`-worker-0.default.svc.cluster.local:2222"]},"task":{"type":"worker","index":0}}`,
+		},
+	}
+	for _, c := range testCase {
+		demoTemplateSpec := c.tfJob.Spec.TFReplicaSpecs[tfv1alpha2.TFReplicaTypeWorker].Template
+		if err := setClusterSpec(&demoTemplateSpec, c.tfJob, c.rt, c.index); err != nil {
+			t.Errorf("Failed to set cluster spec: %v", err)
+		}
+		actual := demoTemplateSpec.Spec.Containers[0].Env[0].Value
+		if c.expectedClusterSpec != actual {
+			t.Errorf("Expected %s, got %s", c.expectedClusterSpec, actual)
+		}
+	}
+}
diff --git a/pkg/controller.v2/controller_tfjob_test.go b/pkg/controller.v2/controller_tfjob_test.go
@@ -143,6 +143,12 @@ func newTFReplicaSpecTemplate() v1.PodTemplateSpec {
 					Name:  tfv1alpha2.DefaultContainerName,
 					Image: testImageName,
 					Args:  []string{"Fake", "Fake"},
+					Ports: []v1.ContainerPort{
+						v1.ContainerPort{
+							Name:          tfv1alpha2.DefaultPortName,
+							ContainerPort: tfv1alpha2.DefaultPort,
+						},
+					},
 				},
 			},
 		},