feat: Add test cases

Signed-off-by: Ce Gao <gaoce@caicloud.io>

pkg/controller/v1alpha2/consts/const.go

@@ -1,5 +1,7 @@
 package consts
 
 const (
-	ConfigExperimentSuggestionName string = "experiment-suggestion-name"
+	ConfigExperimentSuggestionName = "experiment-suggestion-name"
+
+	LabelExperimentName = "experiment"
 )

pkg/controller/v1alpha2/experiment/experiment_controller.go

@@ -217,7 +217,9 @@ func (r *ReconcileExperiment) Reconcile(request reconcile.Request) (reconcile.Re
 	}
 	instance := original.DeepCopy()
 
+	log.Info("DEBUG", "ins", instance)
 	if needUpdate, finalizers := instance.NeedUpdateFinalizers(); needUpdate {
+		log.Info("DEBUG+2", "ins", instance, "finalizers", finalizers)
 		return r.updateFinalizers(instance, finalizers)
 	}
 

pkg/controller/v1alpha2/experiment/experiment_controller_test.go

@@ -7,14 +7,21 @@ import (
 
 	"github.com/golang/mock/gomock"
 	"github.com/onsi/gomega"
+	"k8s.io/apimachinery/pkg/api/errors"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/apimachinery/pkg/types"
+	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/manager"
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
 	logf "sigs.k8s.io/controller-runtime/pkg/runtime/log"
 
+	commonapiv1alpha2 "github.com/kubeflow/katib/pkg/api/operators/apis/common/v1alpha2"
 	experimentsv1alpha2 "github.com/kubeflow/katib/pkg/api/operators/apis/experiment/v1alpha2"
+	trialsv1alpha2 "github.com/kubeflow/katib/pkg/api/operators/apis/trial/v1alpha2"
+	api_pb "github.com/kubeflow/katib/pkg/api/v1alpha2"
+	"github.com/kubeflow/katib/pkg/controller/v1alpha2/consts"
 	managerclientmock "github.com/kubeflow/katib/pkg/mock/v1alpha2/experiment/managerclient"
 	manifestmock "github.com/kubeflow/katib/pkg/mock/v1alpha2/experiment/manifest"
 	suggestionmock "github.com/kubeflow/katib/pkg/mock/v1alpha2/experiment/suggestion"
@@ -24,7 +31,7 @@ const (
 	experimentName = "foo"
 	namespace      = "default"
 
-	timeout = time.Second * 20
+	timeout = time.Second * 5
 )
 
 var expectedRequest = reconcile.Request{NamespacedName: types.NamespacedName{Name: experimentName, Namespace: namespace}}
@@ -80,6 +87,12 @@ func TestCreateExperiment(t *testing.T) {
 		mgrStopped.Wait()
 	}()
 
+	stopCh := make(chan struct{})
+	mgr.GetCache().WaitForCacheSync(stopCh)
+	defer func() {
+		close(stopCh)
+	}()
+
 	// Create the Trial object and expect the Reconcile and Deployment to be created
 	err = c.Create(context.TODO(), instance)
 	// The instance object may not be a valid object because it might be missing some required fields.
@@ -89,15 +102,216 @@ func TestCreateExperiment(t *testing.T) {
 		return
 	}
 	g.Expect(err).NotTo(gomega.HaveOccurred())
-	defer c.Delete(context.TODO(), instance)
 	g.Eventually(requests, timeout).Should(gomega.Receive(gomega.Equal(expectedRequest)))
+
+	g.Expect(c.Delete(context.TODO(), instance)).NotTo(gomega.HaveOccurred())
+	g.Eventually(func() bool {
+		return errors.IsNotFound(c.Get(context.TODO(),
+			expectedRequest.NamespacedName, instance))
+	}, timeout).Should(gomega.BeTrue())
+}
+
+func TestReconcileExperiment(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+	instance := newFakeInstance()
+
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+	mc := managerclientmock.NewMockManagerClient(mockCtrl)
+	mc.EXPECT().CreateExperimentInDB(gomock.Any()).Return(nil).AnyTimes()
+	mc.EXPECT().UpdateExperimentStatusInDB(gomock.Any()).Return(nil).AnyTimes()
+
+	mockCtrl2 := gomock.NewController(t)
+	defer mockCtrl2.Finish()
+	suggestion := suggestionmock.NewMockSuggestion(mockCtrl)
+	suggestion.EXPECT().GetSuggestions(gomock.Any(), gomock.Any()).Return([]*api_pb.Trial{
+		{
+			Name: trialKey.Name,
+			Spec: &api_pb.TrialSpec{},
+		},
+	}, nil).AnyTimes()
+
+	mockCtrl3 := gomock.NewController(t)
+	defer mockCtrl3.Finish()
+	generator := manifestmock.NewMockGenerator(mockCtrl)
+	generator.EXPECT().GetRunSpecWithHyperParameters(gomock.Any(), gomock.Any(),
+		gomock.Any(), gomock.Any(), gomock.Any()).Return(`apiVersion: "kubeflow.org/v1beta2"
+kind: "TFJob"
+metadata:
+  name: "test"
+  namespace: "default"
+spec:
+  tfReplicaSpecs:
+    PS:
+      replicas: 2
+      restartPolicy: Never
+      template:
+        spec:
+          containers:
+            - name: tensorflow
+              image: kubeflow/tf-dist-mnist-test:1.0
+    Worker:
+      replicas: 4
+      restartPolicy: Never
+      template:
+        spec:
+          containers:
+            - name: tensorflow
+              image: kubeflow/tf-dist-mnist-test:1.0`, nil).AnyTimes()
+	generator.EXPECT().GetMetricsCollectorManifest(
+		gomock.Any(), gomock.Any(),
+		gomock.Any(), gomock.Any(),
+		gomock.Any(), gomock.Any()).Return(`apiVersion: batch/v1beta1
+kind: CronJob
+metadata:
+  name: test
+  namespace: default
+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"
+            - "teste"
+            - "-t"
+            - "test"
+            - "-k"
+            - "TFJob"
+            - "-n"
+            - "default"
+            - "-m"
+            - "test"
+            - "-mn"
+            - "test"
+          restartPolicy: Never`, nil).AnyTimes()
+
+	mgr, err := manager.New(cfg, manager.Options{})
+	g.Expect(err).NotTo(gomega.HaveOccurred())
+	c := mgr.GetClient()
+
+	r := &ReconcileExperiment{
+		Client:        mgr.GetClient(),
+		scheme:        mgr.GetScheme(),
+		ManagerClient: mc,
+		Suggestion:    suggestion,
+		Generator:     generator,
+	}
+	r.updateStatusHandler = func(instance *experimentsv1alpha2.Experiment) error {
+		if !instance.IsCreated() {
+			t.Errorf("Expected got condition created")
+		}
+		return r.updateStatus(instance)
+	}
+
+	recFn, requests := SetupTestReconcile(r)
+	g.Expect(add(mgr, recFn)).NotTo(gomega.HaveOccurred())
+
+	stopMgr, mgrStopped := StartTestManager(mgr, g)
+
+	defer func() {
+		close(stopMgr)
+		mgrStopped.Wait()
+	}()
+
+	stopCh := make(chan struct{})
+	mgr.GetCache().WaitForCacheSync(stopCh)
+	defer func() {
+		close(stopCh)
+	}()
+
+	// Create the Trial object and expect the Reconcile and Deployment to be created
+	err = c.Create(context.TODO(), instance)
+	// The instance object may not be a valid object because it might be missing some required fields.
+	// Please modify the instance object by adding required fields and then remove the following if statement.
+	if apierrors.IsInvalid(err) {
+		t.Logf("failed to create object, got an invalid object error: %v", err)
+		return
+	}
+	g.Expect(err).NotTo(gomega.HaveOccurred())
+	g.Eventually(requests, timeout).Should(gomega.Receive(gomega.Equal(expectedRequest)))
+
+	trials := &trialsv1alpha2.TrialList{}
+	g.Eventually(func() int {
+		label := labels.Set{
+			consts.LabelExperimentName: experimentName,
+		}
+		c.List(context.TODO(), &client.ListOptions{
+			LabelSelector: label.AsSelector(),
+		}, trials)
+		return len(trials.Items)
+	}, timeout).
+		Should(gomega.Equal(1))
+
+	g.Expect(c.Delete(context.TODO(), instance)).NotTo(gomega.HaveOccurred())
+	g.Eventually(func() bool {
+		return errors.IsNotFound(c.Get(context.TODO(),
+			expectedRequest.NamespacedName, instance))
+	}, timeout).Should(gomega.BeTrue())
 }
 
 func newFakeInstance() *experimentsv1alpha2.Experiment {
+	var parallelCount int32 = 1
+	var goal float64 = 99.9
 	return &experimentsv1alpha2.Experiment{
 		ObjectMeta: metav1.ObjectMeta{
 			Name:      experimentName,
 			Namespace: namespace,
 		},
+		Spec: experimentsv1alpha2.ExperimentSpec{
+			RetainHistoricalData: false,
+			ParallelTrialCount:   &parallelCount,
+			MaxTrialCount:        &parallelCount,
+			Objective: &commonapiv1alpha2.ObjectiveSpec{
+				Type:                commonapiv1alpha2.ObjectiveTypeMaximize,
+				Goal:                &goal,
+				ObjectiveMetricName: "accuracy",
+			},
+			TrialTemplate: &experimentsv1alpha2.TrialTemplate{
+				GoTemplate: &experimentsv1alpha2.GoTemplate{
+					RawTemplate: `apiVersion: "kubeflow.org/v1beta2"
+kind: TFJob
+metadata:
+  name: {{.Trial}}
+  namespace: {{.NameSpace}}
+spec:
+  tfReplicaSpecs:
+  Worker:
+    replicas: 1 
+    restartPolicy: OnFailure
+    template:
+      spec:
+        containers:
+          - name: tensorflow 
+            image: gcr.io/kubeflow-ci/tf-mnist-with-summaries:1.0
+            imagePullPolicy: Always
+            command:
+              - "python"
+              - "/var/tf_mnist/mnist_with_summaries.py"
+              - "--log_dir=/train/{{.Trial}}"
+              {{- with .HyperParameters}}
+              {{- range .}}
+              - "{{.Name}}={{.Value}}"
+              {{- end}}
+              {{- end}}
+            volumeMounts:
+              - mountPath: "/train"
+                name: "train"
+        volumes:
+          - name: "train"
+            persistentVolumeClaim:
+              claimName: "tfevent-volume"`,
+				},
+			},
+		},
 	}
 }

pkg/controller/v1alpha2/experiment/experiment_util.go

@@ -16,6 +16,7 @@ import (
 	experimentsv1alpha2 "github.com/kubeflow/katib/pkg/api/operators/apis/experiment/v1alpha2"
 	trialsv1alpha2 "github.com/kubeflow/katib/pkg/api/operators/apis/trial/v1alpha2"
 	apiv1alpha2 "github.com/kubeflow/katib/pkg/api/v1alpha2"
+	"github.com/kubeflow/katib/pkg/controller/v1alpha2/consts"
 )
 
 func (r *ReconcileExperiment) createTrialInstance(expInstance *experimentsv1alpha2.Experiment, trialInstance *apiv1alpha2.Trial) error {
@@ -25,7 +26,7 @@ func (r *ReconcileExperiment) createTrialInstance(expInstance *experimentsv1alph
 	trial := &trialsv1alpha2.Trial{}
 	trial.Name = fmt.Sprintf("%s-%s", expInstance.GetName(), utilrand.String(8))
 	trial.Namespace = expInstance.GetNamespace()
-	trial.Labels = map[string]string{"experiment": expInstance.GetName()}
+	trial.Labels = map[string]string{consts.LabelExperimentName: expInstance.GetName()}
 
 	if err := controllerutil.SetControllerReference(expInstance, trial, r.scheme); err != nil {
 		logger.Error(err, "Set controller reference error")
@@ -73,6 +74,7 @@ func (r *ReconcileExperiment) createTrialInstance(expInstance *experimentsv1alph
 	}
 	trial.Spec.MetricsCollectorSpec = mcSpec
 
+	logger.Info("DEBUG-util", "trial", trial)
 	if err := r.Create(context.TODO(), trial); err != nil {
 		logger.Error(err, "Trial create error", "Trial name", trial.Name)
 		return err