feat(trial): Refactor trial controller and add basic test cases (#528)

* feat(trial): Add test * vendor: Update Signed-off-by: Ce Gao <[email protected]> * feat(trial): Add test cases Signed-off-by: Ce Gao <[email protected]> * chore: Add mockgen and test crd Signed-off-by: Ce Gao <[email protected]> * fix: Update test Signed-off-by: Ce Gao <[email protected]> * fix: Update Signed-off-by: Ce Gao <[email protected]> * fix: Fix export path Signed-off-by: Ce Gao <[email protected]> * fix: Introduce the dependency Signed-off-by: Ce Gao <[email protected]> * feat: Add status handler Signed-off-by: Ce Gao <[email protected]> * Fix: Use 1.0.7 Signed-off-by: Ce Gao <[email protected]> * fix: Fix v1alpha1 Signed-off-by: Ce Gao <[email protected]> * fix: Update deploy scripts Signed-off-by: Ce Gao <[email protected]> * chore: Rebase master Signed-off-by: Ce Gao <[email protected]>
kubeflow · May 20, 2019 · bc343ff · bc343ff
1 parent 0712177
commit bc343ff
Show file tree

Hide file tree

Showing 31 changed files with 1,567 additions and 111 deletions.
diff --git a/Gopkg.lock b/Gopkg.lock
diff --git a/...fests/v1alpha2/katib-controller/crds.yaml → ...pha2/katib-controller/crd-experiment.yaml b/...fests/v1alpha2/katib-controller/crds.yaml → ...pha2/katib-controller/crd-experiment.yaml
@@ -19,25 +19,3 @@ spec:
     kind: Experiment
     singular: experiment
     plural: experiments
----
-apiVersion: apiextensions.k8s.io/v1beta1
-kind: CustomResourceDefinition
-metadata:
-  name: trials.kubeflow.org
-spec:
-  additionalPrinterColumns:
-  - JSONPath: .status.conditions[-1:].type
-    name: Status
-    type: string
-  - JSONPath: .metadata.creationTimestamp
-    name: Age
-    type: date
-  group: kubeflow.org
-  version: v1alpha2
-  scope: Namespaced
-  subresources:
-    status: {}
-  names:
-    kind: Trial
-    singular: trial
-    plural: trials
diff --git a/manifests/v1alpha2/katib-controller/crd-trial.yaml b/manifests/v1alpha2/katib-controller/crd-trial.yaml
@@ -0,0 +1,21 @@
+apiVersion: apiextensions.k8s.io/v1beta1
+kind: CustomResourceDefinition
+metadata:
+  name: trials.kubeflow.org
+spec:
+  additionalPrinterColumns:
+  - JSONPath: .status.conditions[-1:].type
+    name: Status
+    type: string
+  - JSONPath: .metadata.creationTimestamp
+    name: Age
+    type: date
+  group: kubeflow.org
+  version: v1alpha2
+  scope: Namespaced
+  subresources:
+    status: {}
+  names:
+    kind: Trial
+    singular: trial
+    plural: trials
diff --git a/pkg/controller/v1alpha2/experiment/validator/validator_test.go b/pkg/controller/v1alpha2/experiment/validator/validator_test.go
@@ -34,7 +34,7 @@ spec:
 	mockCtrl := gomock.NewController(t)
 	defer mockCtrl.Finish()
 
-	p := mockmanifest.NewMockProducer(mockCtrl)
+	p := mockmanifest.NewMockGenerator(mockCtrl)
 	g := New(p)
 
 	p.EXPECT().GetRunSpec(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(trialTFJobTemplate, nil)
@@ -55,7 +55,7 @@ metadata:
 	mockCtrl := gomock.NewController(t)
 	defer mockCtrl.Finish()
 
-	p := mockmanifest.NewMockProducer(mockCtrl)
+	p := mockmanifest.NewMockGenerator(mockCtrl)
 	g := New(p)
 
 	p.EXPECT().GetRunSpec(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(trialJobTemplate, nil)
@@ -70,7 +70,7 @@ func TestValidateExperiment(t *testing.T) {
 	mockCtrl := gomock.NewController(t)
 	defer mockCtrl.Finish()
 
-	p := mockmanifest.NewMockProducer(mockCtrl)
+	p := mockmanifest.NewMockGenerator(mockCtrl)
 	g := New(p)
 
 	trialJobTemplate := `apiVersion: "batch/v1"

diff --git a/...ontroller/v1alpha2/trial/util/api_util.go → ...pha2/trial/managerclient/managerclient.go b/...ontroller/v1alpha2/trial/util/api_util.go → ...pha2/trial/managerclient/managerclient.go
@@ -1,19 +1,4 @@
-/*
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-*/
-
-package util
+package managerclient
 
 import (
 	"fmt"
@@ -24,8 +9,25 @@ import (
 	common "github.com/kubeflow/katib/pkg/common/v1alpha2"
 )
 
-func CreateTrialInDB(instance *trialsv1alpha2.Trial) error {
-	trial := GetTrialConf(instance)
+// ManagerClient is the interface for katib manager client in trial controller.
+type ManagerClient interface {
+	CreateTrialInDB(instance *trialsv1alpha2.Trial) error
+	UpdateTrialStatusInDB(instance *trialsv1alpha2.Trial) error
+	GetTrialObservation(instance *trialsv1alpha2.Trial) error
+	GetTrialConf(instance *trialsv1alpha2.Trial) *api_pb.Trial
+}
+
+// DefaultClient implements the Client interface.
+type DefaultClient struct {
+}
+
+// New creates a new ManagerClient.
+func New() ManagerClient {
+	return &DefaultClient{}
+}
+
+func (d *DefaultClient) CreateTrialInDB(instance *trialsv1alpha2.Trial) error {
+	trial := d.GetTrialConf(instance)
 	request := &api_pb.RegisterTrialRequest{
 		Trial: trial,
 	}
@@ -35,7 +37,7 @@ func CreateTrialInDB(instance *trialsv1alpha2.Trial) error {
 	return nil
 }
 
-func UpdateTrialStatusInDB(instance *trialsv1alpha2.Trial) error {
+func (d *DefaultClient) UpdateTrialStatusInDB(instance *trialsv1alpha2.Trial) error {
 	newStatus := &api_pb.TrialStatus{
 		StartTime:      common.ConvertTime2RFC3339(instance.Status.StartTime),
 		CompletionTime: common.ConvertTime2RFC3339(instance.Status.CompletionTime),
@@ -55,21 +57,21 @@ func UpdateTrialStatusInDB(instance *trialsv1alpha2.Trial) error {
 		newStatus.Observation = observation
 	}
 	request := &api_pb.UpdateTrialStatusRequest{
-		NewStatus:      newStatus,
-		TrialName:      instance.Name,
+		NewStatus: newStatus,
+		TrialName: instance.Name,
 	}
 	if _, err := common.UpdateTrialStatus(request); err != nil {
 		return err
 	}
 	return nil
 }
 
-func GetTrialObservation(instance *trialsv1alpha2.Trial) error {
+func (d *DefaultClient) GetTrialObservation(instance *trialsv1alpha2.Trial) error {
 
 	return nil
 }
 
-func GetTrialConf(instance *trialsv1alpha2.Trial) *api_pb.Trial {
+func (d *DefaultClient) GetTrialConf(instance *trialsv1alpha2.Trial) *api_pb.Trial {
 	trial := &api_pb.Trial{
 		Spec: &api_pb.TrialSpec{
 			Objective: &api_pb.ObjectiveSpec{

diff --git a/pkg/controller/v1alpha2/trial/trial_controller.go b/pkg/controller/v1alpha2/trial/trial_controller.go
@@ -41,6 +41,7 @@ import (
 
 	trialsv1alpha2 "github.com/kubeflow/katib/pkg/api/operators/apis/trial/v1alpha2"
 	commonv1alpha2 "github.com/kubeflow/katib/pkg/common/v1alpha2"
+	"github.com/kubeflow/katib/pkg/controller/v1alpha2/trial/managerclient"
 	trialutil "github.com/kubeflow/katib/pkg/controller/v1alpha2/trial/util"
 )
 
@@ -61,7 +62,13 @@ func Add(mgr manager.Manager) error {
 
 // newReconciler returns a new reconcile.Reconciler
 func newReconciler(mgr manager.Manager) reconcile.Reconciler {
-	return &ReconcileTrial{Client: mgr.GetClient(), scheme: mgr.GetScheme()}
+	r := &ReconcileTrial{
+		Client: mgr.GetClient(), 
+		scheme: mgr.GetScheme(),
+		ManagerClient: managerclient.New(),
+	}
+	r.updateStatusHandler = r.updateStatus
+	return r
 }
 
 // add adds a new Controller to mgr with r as the reconcile.Reconciler
@@ -109,6 +116,10 @@ var _ reconcile.Reconciler = &ReconcileTrial{}
 type ReconcileTrial struct {
 	client.Client
 	scheme *runtime.Scheme
+
+	managerclient.ManagerClient
+	// updateStatusHandler is defined for test purpose.
+	updateStatusHandler updateStatusFunc
 }
 
 // Reconcile reads that state of the cluster for a Trial object and makes changes based on the state read
@@ -149,7 +160,7 @@ func (r *ReconcileTrial) Reconcile(request reconcile.Request) (reconcile.Result,
 		}
 		msg := "Trial is created"
 		instance.MarkTrialStatusCreated(trialutil.TrialCreatedReason, msg)
-		err = trialutil.CreateTrialInDB(instance)
+		err = r.CreateTrialInDB(instance)
 		if err != nil {
 			logger.Error(err, "Create trial in DB error")
 			return reconcile.Result{}, err
@@ -165,12 +176,12 @@ func (r *ReconcileTrial) Reconcile(request reconcile.Request) (reconcile.Result,
 
 	if !equality.Semantic.DeepEqual(original.Status, instance.Status) {
 		//assuming that only status change
-		err = trialutil.UpdateTrialStatusInDB(instance)
+		err = r.UpdateTrialStatusInDB(instance)
 		if err != nil {
 			logger.Error(err, "Update trial status in DB error")
 			return reconcile.Result{}, err
 		}
-		err = r.Status().Update(context.TODO(), instance)
+		err = r.updateStatusHandler(instance)
 		if err != nil {
 			logger.Error(err, "Update trial instance status error")
 			return reconcile.Result{}, err
@@ -301,7 +312,7 @@ func (r *ReconcileTrial) reconcileMetricsCollector(instance *trialsv1alpha2.Tria
 	if err != nil {
 		if errors.IsNotFound(err) {
 			logger.Info("Creating Metrics Collector",
-				"name", desiredMetricsCollector.GetName(), 
+				"name", desiredMetricsCollector.GetName(),
 				"namespace", desiredMetricsCollector.GetNamespace())
 			err = r.Create(context.TODO(), desiredMetricsCollector)
 			if err != nil {

diff --git a/pkg/controller/v1alpha2/trial/trial_controller_status.go b/pkg/controller/v1alpha2/trial/trial_controller_status.go
@@ -0,0 +1,17 @@
+package trial
+
+import (
+	"context"
+
+	trialsv1alpha2 "github.com/kubeflow/katib/pkg/api/operators/apis/trial/v1alpha2"
+)
+
+type updateStatusFunc func(instance *trialsv1alpha2.Trial) error
+
+func (r *ReconcileTrial) updateStatus(instance *trialsv1alpha2.Trial) error {
+	err := r.Status().Update(context.TODO(), instance)
+	if err != nil {
+		return err
+	}
+	return nil
+}
diff --git a/pkg/controller/v1alpha2/trial/trial_controller_suite_test.go b/pkg/controller/v1alpha2/trial/trial_controller_suite_test.go
@@ -0,0 +1,79 @@
+/*
+Copyright 2019 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package trial
+
+import (
+	stdlog "log"
+	"os"
+	"path/filepath"
+	"sync"
+	"testing"
+
+	"github.com/onsi/gomega"
+	"k8s.io/client-go/kubernetes/scheme"
+	"k8s.io/client-go/rest"
+	"sigs.k8s.io/controller-runtime/pkg/envtest"
+	"sigs.k8s.io/controller-runtime/pkg/manager"
+	"sigs.k8s.io/controller-runtime/pkg/reconcile"
+
+	"github.com/kubeflow/katib/pkg/api/operators/apis"
+)
+
+var cfg *rest.Config
+
+func TestMain(m *testing.M) {
+	t := &envtest.Environment{
+		CRDDirectoryPaths: []string{
+			filepath.Join("..", "..", "..", "..", "manifests", "v1alpha2", "katib-controller"),
+			filepath.Join("..", "..", "..", "..", "test", "unit", "v1alpha2", "crds"),
+		},
+	}
+	apis.AddToScheme(scheme.Scheme)
+
+	var err error
+	if cfg, err = t.Start(); err != nil {
+		stdlog.Fatal(err)
+	}
+
+	code := m.Run()
+	t.Stop()
+	os.Exit(code)
+}
+
+// SetupTestReconcile returns a reconcile.Reconcile implementation that delegates to inner and
+// writes the request to requests after Reconcile is finished.
+func SetupTestReconcile(inner reconcile.Reconciler) (reconcile.Reconciler, chan reconcile.Request) {
+	requests := make(chan reconcile.Request)
+	fn := reconcile.Func(func(req reconcile.Request) (reconcile.Result, error) {
+		result, err := inner.Reconcile(req)
+		requests <- req
+		return result, err
+	})
+	return fn, requests
+}
+
+// StartTestManager adds recFn
+func StartTestManager(mgr manager.Manager, g *gomega.GomegaWithT) (chan struct{}, *sync.WaitGroup) {
+	stop := make(chan struct{})
+	wg := &sync.WaitGroup{}
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		g.Expect(mgr.Start(stop)).NotTo(gomega.HaveOccurred())
+	}()
+	return stop, wg
+}