WIP

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

pkg/controller/v1alpha2/trial/trial_controller.go

@@ -144,21 +144,22 @@ func (r *ReconcileTrial) Reconcile(request reconcile.Request) (reconcile.Result,
 	instance := original.DeepCopy()
 
 	if !instance.IsCreated() {
-		//Trial not created in DB
 		if instance.Status.StartTime == nil {
 			now := metav1.Now()
 			instance.Status.StartTime = &now
 		}
 		if instance.Status.CompletionTime == nil {
 			instance.Status.CompletionTime = &metav1.Time{}
 		}
-		msg := "Trial is created"
-		instance.MarkTrialStatusCreated(TrialCreatedReason, msg)
 		err = r.CreateTrialInDB(instance)
 		if err != nil {
 			logger.Error(err, "Create trial in DB error")
-			return reconcile.Result{}, err
+			return reconcile.Result{
+				Requeue: true,
+			}, err
 		}
+		msg := "Trial is created"
+		instance.MarkTrialStatusCreated(TrialCreatedReason, msg)
 	} else {
 		// Trial already created in DB
 		err := r.reconcileTrial(instance)

pkg/controller/v1alpha2/trial/trial_controller_suite_test.go

@@ -66,6 +66,18 @@ func SetupTestReconcile(inner reconcile.Reconciler) (reconcile.Reconciler, chan
 	return fn, requests
 }
 
+func SetupTestReconcileWithResult(inner reconcile.Reconciler) (reconcile.Reconciler, chan reconcile.Request, chan reconcile.Result) {
+	requests := make(chan reconcile.Request)
+	results := make(chan reconcile.Result)
+	fn := reconcile.Func(func(req reconcile.Request) (reconcile.Result, error) {
+		result, err := inner.Reconcile(req)
+		requests <- req
+		results <- result
+		return result, err
+	})
+	return fn, requests, results
+}
+
 // StartTestManager adds recFn
 func StartTestManager(mgr manager.Manager, g *gomega.GomegaWithT) (chan struct{}, *sync.WaitGroup) {
 	stop := make(chan struct{})

pkg/controller/v1alpha2/trial/trial_controller_test.go

@@ -2,6 +2,7 @@ package trial
 
 import (
 	"bytes"
+	"fmt"
 	"testing"
 	"time"
 
@@ -31,6 +32,7 @@ const (
 )
 
 var expectedRequest = reconcile.Request{NamespacedName: types.NamespacedName{Name: trialName, Namespace: namespace}}
+var expectedResult = reconcile.Result{Requeue: true}
 var tfJobKey = types.NamespacedName{Name: "test", Namespace: namespace}
 
 func init() {
@@ -162,6 +164,129 @@ func TestReconcileTFJobTrial(t *testing.T) {
 		Should(gomega.MatchError("tfjobs.kubeflow.org \"test\" not found"))
 }
 
+func TestReconcileComplectedTFJobTrial(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+	instance := newFakeTrialWithTFJob()
+
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	mc := managerclientmock.NewMockManagerClient(mockCtrl)
+	mc.EXPECT().CreateTrialInDB(gomock.Any()).Return(nil).AnyTimes()
+	mc.EXPECT().UpdateTrialStatusInDB(gomock.Any()).Return(nil).AnyTimes()
+
+	// Setup the Manager and Controller.  Wrap the Controller Reconcile function so it writes each request to a
+	// channel when it is finished.
+	mgr, err := manager.New(cfg, manager.Options{})
+	g.Expect(err).NotTo(gomega.HaveOccurred())
+	c := mgr.GetClient()
+
+	r := &ReconcileTrial{
+		Client:        mgr.GetClient(),
+		scheme:        mgr.GetScheme(),
+		ManagerClient: mc,
+	}
+
+	r.updateStatusHandler = func(instance *trialsv1alpha2.Trial) 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()
+	}()
+
+	// 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())
+	defer c.Delete(context.TODO(), instance)
+	g.Eventually(requests, timeout).Should(gomega.Receive(gomega.Equal(expectedRequest)))
+
+	g.Eventually(func() error {
+		return c.Get(context.TODO(), expectedRequest.NamespacedName, instance)
+	}, timeout).
+		Should(gomega.Succeed())
+	instance.MarkTrialStatusSucceeded("", "")
+	g.Expect(c.Status().Update(context.TODO(), instance)).NotTo(gomega.HaveOccurred())
+	g.Eventually(func() bool {
+		err := c.Get(context.TODO(), expectedRequest.NamespacedName, instance)
+		if err == nil && instance.IsCompleted() {
+			return true
+		}
+		return false
+	}, timeout).
+		Should(gomega.BeTrue())
+}
+
+func TestFailedToCreateTrialInDB(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+	instance := newFakeTrialWithTFJob()
+
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	mc := managerclientmock.NewMockManagerClient(mockCtrl)
+	err := fmt.Errorf("test")
+	mc.EXPECT().CreateTrialInDB(gomock.Any()).Return(err).AnyTimes()
+	mc.EXPECT().UpdateTrialStatusInDB(gomock.Any()).Return(nil).AnyTimes()
+
+	// Setup the Manager and Controller.  Wrap the Controller Reconcile function so it writes each request to a
+	// channel when it is finished.
+	mgr, err := manager.New(cfg, manager.Options{})
+	g.Expect(err).NotTo(gomega.HaveOccurred())
+	c := mgr.GetClient()
+
+	r := &ReconcileTrial{
+		Client:        mgr.GetClient(),
+		scheme:        mgr.GetScheme(),
+		ManagerClient: mc,
+	}
+
+	r.updateStatusHandler = func(instance *trialsv1alpha2.Trial) error {
+		if !instance.IsCreated() {
+			t.Errorf("Expected got condition created")
+		}
+		return r.updateStatus(instance)
+	}
+
+	recFn, requests, results := SetupTestReconcileWithResult(r)
+	g.Expect(add(mgr, recFn)).NotTo(gomega.HaveOccurred())
+
+	stopMgr, mgrStopped := StartTestManager(mgr, g)
+
+	defer func() {
+		close(stopMgr)
+		mgrStopped.Wait()
+	}()
+
+	// 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())
+	defer c.Delete(context.TODO(), instance)
+	g.Eventually(requests, timeout).Should(gomega.Receive(gomega.Equal(expectedRequest)))
+	g.Eventually(results, timeout).Should(gomega.Receive(gomega.Equal(expectedResult)))
+}
+
 func newFakeTrialWithTFJob() *trialsv1alpha2.Trial {
 	objectiveSpec := commonv1alpha2.ObjectiveSpec{ObjectiveMetricName: "test"}
 	t := &trialsv1alpha2.Trial{