From 317c40d32049033379175df67d0bd445bae69f37 Mon Sep 17 00:00:00 2001 From: forsaken628 Date: Fri, 14 Jun 2024 11:43:29 +0800 Subject: [PATCH] use gomock Signed-off-by: forsaken628 --- .../trial/trial_controller_test.go | 34 +++++++++---------- 1 file changed, 16 insertions(+), 18 deletions(-) diff --git a/pkg/controller.v1beta1/trial/trial_controller_test.go b/pkg/controller.v1beta1/trial/trial_controller_test.go index 4b82c225285..085f2e7c0fb 100644 --- a/pkg/controller.v1beta1/trial/trial_controller_test.go +++ b/pkg/controller.v1beta1/trial/trial_controller_test.go @@ -21,6 +21,7 @@ import ( "testing" "time" + "github.com/golang/mock/gomock" "github.com/onsi/gomega" "github.com/prometheus/client_golang/prometheus" "github.com/spf13/viper" @@ -42,6 +43,7 @@ import ( "github.com/kubeflow/katib/pkg/controller.v1beta1/consts" trialutil "github.com/kubeflow/katib/pkg/controller.v1beta1/trial/util" "github.com/kubeflow/katib/pkg/controller.v1beta1/util" + managerclientmock "github.com/kubeflow/katib/pkg/mock/v1beta1/trial/managerclient" ) const ( @@ -87,7 +89,9 @@ func TestAdd(t *testing.T) { func TestReconcileBatchJob(t *testing.T) { g := gomega.NewGomegaWithT(t) - mockMC := &mockManagerClient{} + mockCtrl := gomock.NewController(t) + defer mockCtrl.Finish() + mockManagerClient := managerclientmock.NewMockManagerClient(mockCtrl) // Setup the Manager and Controller. Wrap the Controller Reconcile function so it writes each request to a // channel when it is finished. @@ -98,7 +102,7 @@ func TestReconcileBatchJob(t *testing.T) { r := &ReconcileTrial{ Client: mgr.GetClient(), scheme: mgr.GetScheme(), - ManagerClient: mockMC, + ManagerClient: mockManagerClient, recorder: mgr.GetEventRecorderFor(ControllerName), collector: trialutil.NewTrialsCollector(mgr.GetCache(), prometheus.NewRegistry()), } @@ -177,7 +181,8 @@ func TestReconcileBatchJob(t *testing.T) { t.Run(`Trial run with "Failed" BatchJob.`, func(t *testing.T) { g := gomega.NewGomegaWithT(t) - mockMC.msg = observationLogUnavailable + mockManagerClient.EXPECT().DeleteTrialObservationLog(gomock.Any()).Return(nil, nil) + trial := newFakeTrialBatchJob() batchJob := &batchv1.Job{} @@ -236,7 +241,10 @@ func TestReconcileBatchJob(t *testing.T) { t.Run(`Trail with "Complete" BatchJob and Available metrics.`, func(t *testing.T) { g := gomega.NewGomegaWithT(t) - mockMC.msg = observationLogAvailable + gomock.InOrder( + mockManagerClient.EXPECT().GetTrialObservationLog(gomock.Any()).Return(observationLogAvailable, nil).MinTimes(1), + mockManagerClient.EXPECT().DeleteTrialObservationLog(gomock.Any()).Return(nil, nil), + ) batchJob := &batchv1.Job{} batchJobCompleteMessage := "BatchJob completed test message" batchJobCompleteReason := "BatchJob completed test reason" @@ -284,7 +292,10 @@ func TestReconcileBatchJob(t *testing.T) { t.Run(`Trail with "Complete" BatchJob and Unavailable metrics.`, func(t *testing.T) { g := gomega.NewGomegaWithT(t) - mockMC.msg = observationLogUnavailable + gomock.InOrder( + mockManagerClient.EXPECT().GetTrialObservationLog(gomock.Any()).Return(observationLogUnavailable, nil).MinTimes(1), + mockManagerClient.EXPECT().DeleteTrialObservationLog(gomock.Any()).Return(nil, nil), + ) // Create the Trial trial := newFakeTrialBatchJob() g.Expect(c.Create(ctx, trial)).NotTo(gomega.HaveOccurred()) @@ -433,16 +444,3 @@ func newFakeTrialBatchJob() *trialsv1beta1.Trial { }, } } - -type mockManagerClient struct { - msg *api_pb.GetObservationLogReply -} - -func (c *mockManagerClient) GetTrialObservationLog(instance *trialsv1beta1.Trial) (*api_pb.GetObservationLogReply, error) { - return c.msg, nil -} - -func (c *mockManagerClient) DeleteTrialObservationLog(instance *trialsv1beta1.Trial) (*api_pb.DeleteObservationLogReply, error) { - c.msg = nil - return &api_pb.DeleteObservationLogReply{}, nil -}