fix: Add more test cases

Signed-off-by: Ce Gao <gaoce@caicloud.io>
kubeflow · May 14, 2019 · c446136 · c446136
1 parent 491ca0c
commit c446136
Show file tree

Hide file tree

Showing 6 changed files with 382 additions and 75 deletions.
diff --git a/pkg/controller/v1alpha2/experiment/manifest/producer.go b/pkg/controller/v1alpha2/experiment/manifest/producer.go
@@ -16,6 +16,10 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client"
 )
 
+const (
+	defaultMetricsCollectorTemplateName = "defaultMetricsCollectorTemplate.yaml"
+)
+
 // Producer is the type for manifests producer.
 type Producer interface {
 	GetRunSpec(e *experimentsv1alpha2.Experiment, experiment, trial, namespace string) (string, error)
@@ -25,7 +29,7 @@ type Producer interface {
 
 // General is the default implementation of Producer.
 type General struct {
-	client *katibclient.KatibClient
+	client katibclient.Client
 }
 
 // New creates a new Producer.
@@ -53,7 +57,7 @@ func (g *General) GetMetricsCollectorManifest(experimentName string, trialName s
 	if mcs != nil && mcs.GoTemplate.RawTemplate != "" {
 		mtp, err = template.New("MetricsCollector").Parse(mcs.GoTemplate.RawTemplate)
 	} else {
-		mctp := "defaultMetricsCollectorTemplate.yaml"
+		mctp := defaultMetricsCollectorTemplateName
 		if mcs != nil && mcs.GoTemplate.TemplateSpec != nil {
 			mctp = mcs.GoTemplate.TemplateSpec.TemplatePath
 		}

diff --git a/pkg/controller/v1alpha2/experiment/manifest/producer_suite_test.go b/pkg/controller/v1alpha2/experiment/manifest/producer_suite_test.go
diff --git a/pkg/controller/v1alpha2/experiment/manifest/producer_test.go b/pkg/controller/v1alpha2/experiment/manifest/producer_test.go
@@ -5,14 +5,182 @@ import (
 	"testing"
 	"text/template"
 
+	"github.com/golang/mock/gomock"
+
 	experimentsv1alpha2 "github.com/kubeflow/katib/pkg/api/operators/apis/experiment/v1alpha2"
 	apiv1alpha2 "github.com/kubeflow/katib/pkg/api/v1alpha2"
-	"github.com/kubeflow/katib/pkg/util/v1alpha2/katibclient"
+	katibclientmock "github.com/kubeflow/katib/pkg/mock/v1alpha2/util/katibclient"
+)
+
+const (
+	rawTemplate = `apiVersion: batch/v1
+kind: Job
+metadata:
+name: {{.Trial}}
+namespace: {{.NameSpace}}
+spec:
+	template:
+		spec:
+			containers:
+				- name: {{.Trial}}
+				  image: katib/mxnet-mnist-example
+				  command:
+					- "python"
+					- "/mxnet/example/image-classification/train_mnist.py"
+					- "--batch-size=64"
+					{{- with .HyperParameters}}
+					{{- range .}}
+					- "{{.Name}}={{.Value}}"
+					{{- end}}
+					{{- end}}`
+	rawMetricsCollector = `apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: metrics-collector-template
+  namespace: kubeflow
+data:
+  defaultMetricsCollectorTemplate.yaml : |-
+	apiVersion: batch/v1beta1
+	kind: CronJob
+	metadata:
+	  name: {{.Trial}}
+	  namespace: {{.NameSpace}}
+	spec:
+	  schedule: "*/1 * * * *"
+	  successfulJobsHistoryLimit: 0
+	  failedJobsHistoryLimit: 1
+	  jobTemplate:
+		spec:
+		  backoffLimit: 0
+		  template:
+			spec:
+			  serviceAccountName: metrics-collector
+			  containers:
+			  - name: {{.Trial}}
+				image: katib/metrics-collector
+				args:
+				- "./metricscollector.v1alpha2"
+				- "-e"
+				- "{{.Experiment}}"
+				- "-t"
+				- "{{.Trial}}"
+				- "-k"
+				- "{{.JobKind}}"
+				- "-n"
+				- "{{.NameSpace}}"
+				- "-m"
+				- "{{.ManagerService}}"
+				- "-mn"
+				- "{{.MetricNames}}"`
 )
 
+func TestGetMetricsCollectorManifest(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	c := katibclientmock.NewMockClient(mockCtrl)
+
+	p := &General{
+		client: c,
+	}
+
+	c.EXPECT().GetMetricsCollectorTemplates(gomock.Any()).Return(map[string]string{
+		defaultMetricsCollectorTemplateName: rawMetricsCollector,
+	}, nil)
+
+	instance := newFakeInstance()
+
+	actual, err := p.GetMetricsCollectorManifest("test", "test", "Job", "fakens", []string{"test"}, instance.Spec.MetricsCollectorSpec)
+	if err != nil {
+		t.Errorf("Expected nil, got %v", err)
+	}
+	expected := `apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: metrics-collector-template
+  namespace: kubeflow
+data:
+  defaultMetricsCollectorTemplate.yaml : |-
+	apiVersion: batch/v1beta1
+	kind: CronJob
+	metadata:
+	  name: test
+	  namespace: fakens
+	spec:
+	  schedule: "*/1 * * * *"
+	  successfulJobsHistoryLimit: 0
+	  failedJobsHistoryLimit: 1
+	  jobTemplate:
+		spec:
+		  backoffLimit: 0
+		  template:
+			spec:
+			  serviceAccountName: metrics-collector
+			  containers:
+			  - name: test
+				image: katib/metrics-collector
+				args:
+				- "./metricscollector.v1alpha2"
+				- "-e"
+				- "test"
+				- "-t"
+				- "test"
+				- "-k"
+				- "Job"
+				- "-n"
+				- "fakens"
+				- "-m"
+				- "katib-manager.default:6789"
+				- "-mn"
+				- "test"`
+
+	if expected != actual.String() {
+		t.Errorf("Expected %s, got %s", expected, actual)
+	}
+}
+
+func TestGetTrialTemplateConfigMap(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	c := katibclientmock.NewMockClient(mockCtrl)
+
+	p := &General{
+		client: c,
+	}
+
+	templatePath := "test.yaml"
+
+	c.EXPECT().GetConfigMap(gomock.Any(), gomock.Any()).Return(map[string]string{
+		templatePath: rawTemplate,
+	}, nil)
+
+	instance := newFakeInstance()
+	instance.Spec.TrialTemplate.GoTemplate.TemplateSpec = &experimentsv1alpha2.TemplateSpec{
+		TemplatePath: templatePath,
+	}
+	instance.Spec.TrialTemplate.GoTemplate.RawTemplate = ""
+	actual, err := p.getTrialTemplate(instance)
+	if err != nil {
+		t.Errorf("Expected nil, got %v", err)
+	}
+	expected, err := template.New("Trial").Parse(rawTemplate)
+	if err != nil {
+		t.Errorf("Expected nil, got %v", err)
+	}
+	if !reflect.DeepEqual(expected, actual) {
+		t.Errorf("Expected %v, got %v", *expected, *actual)
+	}
+}
+
 func TestGetTrialTemplate(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	c := katibclientmock.NewMockClient(mockCtrl)
+
 	p := &General{
-		client: katibclient.MustNewTestClient(cfg),
+		client: c,
 	}
 
 	tc := newFakeInstance()
@@ -24,6 +192,9 @@ func TestGetTrialTemplate(t *testing.T) {
 	}
 
 	actual, err := p.getTrialTemplate(tc)
+	if err != nil {
+		t.Errorf("Expected nil, got %v", err)
+	}
 
 	if !reflect.DeepEqual(expected, actual) {
 		t.Errorf("Expected %v, got %v", *expected, *actual)
@@ -33,8 +204,13 @@ func TestGetTrialTemplate(t *testing.T) {
 func TestGetRunSpec(t *testing.T) {
 	tc := newFakeInstance()
 
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	c := katibclientmock.NewMockClient(mockCtrl)
+
 	p := &General{
-		client: katibclient.MustNewTestClient(cfg),
+		client: c,
 	}
 
 	actual, err := p.GetRunSpec(tc, "", "test", "testns")
@@ -64,8 +240,13 @@ spec:
 func TestGetRunSpecWithHP(t *testing.T) {
 	tc := newFakeInstance()
 
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	c := katibclientmock.NewMockClient(mockCtrl)
+
 	p := &General{
-		client: katibclient.MustNewTestClient(cfg),
+		client: c,
 	}
 
 	actual, err := p.GetRunSpecWithHyperParameters(tc, "", "test", "testns", []*apiv1alpha2.ParameterAssignment{
@@ -103,26 +284,14 @@ func newFakeInstance() *experimentsv1alpha2.Experiment {
 		Spec: experimentsv1alpha2.ExperimentSpec{
 			TrialTemplate: &experimentsv1alpha2.TrialTemplate{
 				GoTemplate: &experimentsv1alpha2.GoTemplate{
-					RawTemplate: `apiVersion: batch/v1
-kind: Job
-metadata:
-name: {{.Trial}}
-namespace: {{.NameSpace}}
-spec:
-	template:
-		spec:
-			containers:
-				- name: {{.Trial}}
-				  image: katib/mxnet-mnist-example
-				  command:
-					- "python"
-					- "/mxnet/example/image-classification/train_mnist.py"
-					- "--batch-size=64"
-					{{- with .HyperParameters}}
-					{{- range .}}
-					- "{{.Name}}={{.Value}}"
-					{{- end}}
-					{{- end}}`,
+					RawTemplate: rawTemplate,
+				},
+			},
+			MetricsCollectorSpec: &experimentsv1alpha2.MetricsCollectorSpec{
+				GoTemplate: experimentsv1alpha2.GoTemplate{
+					TemplateSpec: &experimentsv1alpha2.TemplateSpec{
+						TemplatePath: defaultMetricsCollectorTemplateName,
+					},
 				},
 			},
 		},