add ValidateAlphaTwoTFJobSpec to check v1alpha2.TFJobSpec is valid (k…

…ubeflow#702) Signed-off-by: Pengyu Chen <[email protected]>
jetmuffin · Jun 28, 2018 · d3ecb52 · d3ecb52
1 parent aa7e9a8
commit d3ecb52
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 2 deletions.
diff --git a/pkg/apis/tensorflow/validation/validation.go b/pkg/apis/tensorflow/validation/validation.go
@@ -19,10 +19,24 @@ import (
 	"fmt"
 
 	tfv1 "github.com/kubeflow/tf-operator/pkg/apis/tensorflow/v1alpha1"
+	tfv2 "github.com/kubeflow/tf-operator/pkg/apis/tensorflow/v1alpha2"
 	"github.com/kubeflow/tf-operator/pkg/util"
 )
 
-// ValidateTFJobSpec checks that the TFJobSpec is valid.
+// ValidateAlphaTwoTFJobSpec checks that the v1alpha2.TFJobSpec is valid.
+func ValidateAlphaTwoTFJobSpec(c *tfv2.TFJobSpec) error {
+	if c.TFReplicaSpecs == nil {
+		return fmt.Errorf("TFJobSpec is not valid")
+	}
+	for _, value := range c.TFReplicaSpecs {
+		if value == nil || len(value.Template.Spec.Containers) == 0 {
+			return fmt.Errorf("TFJobSpec is not valid")
+		}
+	}
+	return nil
+}
+
+// ValidateTFJobSpec checks that the v1alpha1.TFJobSpec is valid.
 func ValidateTFJobSpec(c *tfv1.TFJobSpec) error {
 	if c.TerminationPolicy == nil || c.TerminationPolicy.Chief == nil {
 		return fmt.Errorf("invalid termination policy: %v", c.TerminationPolicy)

diff --git a/pkg/apis/tensorflow/validation/validation_test.go b/pkg/apis/tensorflow/validation/validation_test.go
@@ -18,11 +18,37 @@ import (
 	"testing"
 
 	tfv1 "github.com/kubeflow/tf-operator/pkg/apis/tensorflow/v1alpha1"
+	tfv2 "github.com/kubeflow/tf-operator/pkg/apis/tensorflow/v1alpha2"
 
 	"github.com/gogo/protobuf/proto"
 	"k8s.io/api/core/v1"
 )
 
+func TestValidateAlphaTwoTFJobSpec(t *testing.T) {
+	testCases := []tfv2.TFJobSpec{
+		{
+			TFReplicaSpecs: nil,
+		},
+		{
+			TFReplicaSpecs: map[tfv2.TFReplicaType]*tfv2.TFReplicaSpec{
+				tfv2.TFReplicaTypeWorker: &tfv2.TFReplicaSpec{
+					Template: v1.PodTemplateSpec{
+						Spec: v1.PodSpec{
+							Containers: []v1.Container{},
+						},
+					},
+				},
+			},
+		},
+	}
+	for _, c := range testCases {
+		err := ValidateAlphaTwoTFJobSpec(&c)
+		if err.Error() != "TFJobSpec is not valid" {
+			t.Error("Failed validate the alpha2.TFJobSpec")
+		}
+	}
+}
+
 func TestValidate(t *testing.T) {
 	type testCase struct {
 		in             *tfv1.TFJobSpec

diff --git a/pkg/controller.v2/informer.go b/pkg/controller.v2/informer.go
@@ -13,6 +13,7 @@ import (
 	"k8s.io/client-go/tools/cache"
 
 	tfv1alpha2 "github.com/kubeflow/tf-operator/pkg/apis/tensorflow/v1alpha2"
+	"github.com/kubeflow/tf-operator/pkg/apis/tensorflow/validation"
 	tfjobinformers "github.com/kubeflow/tf-operator/pkg/client/informers/externalversions"
 	tfjobinformersv1alpha2 "github.com/kubeflow/tf-operator/pkg/client/informers/externalversions/kubeflow/v1alpha2"
 	"github.com/kubeflow/tf-operator/pkg/util/unstructured"
@@ -92,7 +93,7 @@ func tfJobFromUnstructured(obj interface{}) (*tfv1alpha2.TFJob, error) {
 	// This is a simple validation for TFJob to close
 	// https://github.com/kubeflow/tf-operator/issues/641
 	// TODO(gaocegege): Add more validation here.
-	if err != nil || tfjob.Spec.TFReplicaSpecs == nil {
+	if err != nil || validation.ValidateAlphaTwoTFJobSpec(&tfjob.Spec) != nil {
 		return &tfjob, errFailedMarshal
 	}
 	return &tfjob, nil