From d24fc368256959d5c5248abe0db45aa84fb094ba Mon Sep 17 00:00:00 2001 From: Bangqi Zhu Date: Mon, 22 Jul 2024 02:30:26 -0700 Subject: [PATCH] Add more coverage for controllerrevision unit test Signed-off-by: Bangqi Zhu --- pkg/controllers/workspace_controller.go | 5 +- pkg/controllers/workspace_controller_test.go | 82 ++++++++++++++++---- pkg/utils/consts/consts.go | 1 + pkg/utils/test/testUtils.go | 60 +++++++++++++- 4 files changed, 128 insertions(+), 20 deletions(-) diff --git a/pkg/controllers/workspace_controller.go b/pkg/controllers/workspace_controller.go index 84db27004..f55b4d857 100644 --- a/pkg/controllers/workspace_controller.go +++ b/pkg/controllers/workspace_controller.go @@ -204,7 +204,7 @@ func (c *WorkspaceReconciler) updateControllerRevision(ctx context.Context, wObj revisionNum := int64(1) - if exists { + if len(revisions.Items) > 0 { revisionNum = revisions.Items[len(revisions.Items)-1].Revision + 1 } @@ -236,8 +236,7 @@ func (c *WorkspaceReconciler) updateControllerRevision(ctx context.Context, wObj return fmt.Errorf("failed to create new ControllerRevision: %w", err) } - const maxRevisionHistoryLimit = 10 - if len(revisions.Items) > maxRevisionHistoryLimit { + if len(revisions.Items) > consts.MaxRevisionHistoryLimit { if err := c.Delete(ctx, &revisions.Items[0]); err != nil { return fmt.Errorf("failed to delete old revision: %w", err) } diff --git a/pkg/controllers/workspace_controller_test.go b/pkg/controllers/workspace_controller_test.go index e205a31ba..a86078a06 100644 --- a/pkg/controllers/workspace_controller_test.go +++ b/pkg/controllers/workspace_controller_test.go @@ -6,6 +6,7 @@ package controllers import ( "context" "errors" + "fmt" "os" "reflect" "sort" @@ -958,36 +959,86 @@ func TestApplyWorkspaceResource(t *testing.T) { func TestUpdateControllerRevision(t *testing.T) { testcases := map[string]struct { callMocks func(c *test.MockClient) - workspace *v1alpha1.Workspace + workspace v1alpha1.Workspace expectedError error + verifyCalls func(c *test.MockClient) }{ - "No new revision needed": { callMocks: func(c *test.MockClient) { - c.On("List", mock.Anything, mock.IsType(&appsv1.ControllerRevisionList{}), mock.Anything).Return(errors.New("should not be called")).Once() + c.On("List", mock.IsType(context.Background()), mock.Anything, mock.IsType(&appsv1.ControllerRevisionList{}), mock.Anything).Return(errors.New("should not be called")) }, workspace: test.MockWorkspaceWithComputeHash, expectedError: nil, + verifyCalls: func(c *test.MockClient) { + c.AssertNumberOfCalls(t, "List", 0) + c.AssertNumberOfCalls(t, "Create", 0) + c.AssertNumberOfCalls(t, "Update", 0) + c.AssertNumberOfCalls(t, "Delete", 0) + }, }, - "Fail to create ControllerRevision": { callMocks: func(c *test.MockClient) { - c.On("List", mock.Anything, mock.IsType(&appsv1.ControllerRevisionList{}), mock.Anything).Return(nil).Once() - c.On("Create", mock.Anything, mock.IsType(&appsv1.ControllerRevision{}), mock.Anything).Return(errors.New("failed to create ControllerRevision")).Once() - c.On("Update", mock.Anything, mock.IsType(&v1alpha1.Workspace{}), mock.Anything).Return(nil).Once() + c.On("List", mock.IsType(context.Background()), mock.IsType(&appsv1.ControllerRevisionList{}), mock.Anything).Return(nil) + c.On("Create", mock.IsType(context.Background()), mock.IsType(&appsv1.ControllerRevision{}), mock.Anything).Return(errors.New("failed to create ControllerRevision")) + c.On("Update", mock.IsType(context.Background()), mock.IsType(&v1alpha1.Workspace{}), mock.Anything).Return(nil) }, - workspace: test.MockWorkspaceWithMatchingLabel, + workspace: test.MockWorkspaceFailToCreateCR, expectedError: errors.New("failed to create new ControllerRevision: failed to create ControllerRevision"), + verifyCalls: func(c *test.MockClient) { + c.AssertNumberOfCalls(t, "List", 1) + c.AssertNumberOfCalls(t, "Create", 1) + c.AssertNumberOfCalls(t, "Update", 1) + c.AssertNumberOfCalls(t, "Delete", 0) + }, }, - "Successfully create new ControllerRevision": { callMocks: func(c *test.MockClient) { - c.On("List", mock.Anything, mock.IsType(&appsv1.ControllerRevisionList{}), mock.Anything).Return(nil).Once() - c.On("Create", mock.Anything, mock.IsType(&appsv1.ControllerRevision{}), mock.Anything).Return(nil).Once() - c.On("Update", mock.Anything, mock.IsType(&v1alpha1.Workspace{}), mock.Anything).Return(nil).Once() + c.On("List", mock.IsType(context.Background()), mock.IsType(&appsv1.ControllerRevisionList{}), mock.Anything).Return(nil) + c.On("Create", mock.IsType(context.Background()), mock.IsType(&appsv1.ControllerRevision{}), mock.Anything).Return(nil) + c.On("Update", mock.IsType(context.Background()), mock.IsType(&v1alpha1.Workspace{}), mock.Anything).Return(nil) }, - workspace: test.MockWorkspaceWithMatchingLabel, + workspace: test.MockWorkspaceSuccessful, expectedError: nil, + verifyCalls: func(c *test.MockClient) { + c.AssertNumberOfCalls(t, "List", 1) + c.AssertNumberOfCalls(t, "Create", 1) + c.AssertNumberOfCalls(t, "Update", 1) + c.AssertNumberOfCalls(t, "Delete", 0) + }, + }, + "Successfully delete old ControllerRevision": { + callMocks: func(c *test.MockClient) { + revisions := &appsv1.ControllerRevisionList{} + for i := 0; i <= consts.MaxRevisionHistoryLimit; i++ { + revision := &appsv1.ControllerRevision{ + ObjectMeta: v1.ObjectMeta{ + Name: fmt.Sprintf("revision-%d", i), + }, + Revision: int64(i), + } + revisions.Items = append(revisions.Items, *revision) + } + relevantMap := c.CreateMapWithType(revisions) + + for _, obj := range revisions.Items { + m := obj + objKey := client.ObjectKeyFromObject(&m) + relevantMap[objKey] = &m + } + + c.On("List", mock.IsType(context.Background()), mock.IsType(&appsv1.ControllerRevisionList{}), mock.Anything).Return(nil) + c.On("Create", mock.IsType(context.Background()), mock.IsType(&appsv1.ControllerRevision{}), mock.Anything).Return(nil) + c.On("Update", mock.IsType(context.Background()), mock.IsType(&v1alpha1.Workspace{}), mock.Anything).Return(nil) + c.On("Delete", mock.IsType(context.Background()), mock.IsType(&appsv1.ControllerRevision{}), mock.Anything).Return(nil) + }, + workspace: test.MockWorkspaceWithDeleteOldCR, + expectedError: nil, + verifyCalls: func(c *test.MockClient) { + c.AssertNumberOfCalls(t, "List", 1) + c.AssertNumberOfCalls(t, "Create", 1) + c.AssertNumberOfCalls(t, "Update", 1) + c.AssertNumberOfCalls(t, "Delete", 1) + }, }, } for k, tc := range testcases { @@ -1001,12 +1052,15 @@ func TestUpdateControllerRevision(t *testing.T) { } ctx := context.Background() - err := reconciler.updateControllerRevision(ctx, tc.workspace) + err := reconciler.updateControllerRevision(ctx, &tc.workspace) if tc.expectedError == nil { assert.Check(t, err == nil, "Not expected to return error") } else { assert.Equal(t, tc.expectedError.Error(), err.Error()) } + if tc.verifyCalls != nil { + tc.verifyCalls(mockClient) + } }) } } diff --git a/pkg/utils/consts/consts.go b/pkg/utils/consts/consts.go index d5b15c1ca..31ba780b7 100644 --- a/pkg/utils/consts/consts.go +++ b/pkg/utils/consts/consts.go @@ -12,4 +12,5 @@ const ( AWSCloudName = "aws" GPUString = "gpu" SKUString = "sku" + MaxRevisionHistoryLimit = 10 ) diff --git a/pkg/utils/test/testUtils.go b/pkg/utils/test/testUtils.go index 435c8bacb..52d52503e 100644 --- a/pkg/utils/test/testUtils.go +++ b/pkg/utils/test/testUtils.go @@ -73,9 +73,63 @@ var ( ) var ( - MockWorkspaceWithMatchingLabel = &v1alpha1.Workspace{ + MockWorkspaceWithDeleteOldCR = v1alpha1.Workspace{ ObjectMeta: metav1.ObjectMeta{ - Name: "testWorkspace", + Name: "testWorkspace", + Namespace: "kaito", + Annotations: map[string]string{ + "workspace.kaito.io/revision": "DeleteOldCR", + }, + }, + Resource: v1alpha1.ResourceSpec{ + Count: &gpuNodeCount, + InstanceType: "Standard_NC12s_v3", + LabelSelector: &metav1.LabelSelector{ + MatchLabels: map[string]string{ + "workspace.kaito.io/name": "testWorkspace", + }, + }, + }, + Inference: &v1alpha1.InferenceSpec{ + Preset: &v1alpha1.PresetSpec{ + PresetMeta: v1alpha1.PresetMeta{ + Name: "test-model", + }, + }, + }, + } +) + +var ( + MockWorkspaceFailToCreateCR = v1alpha1.Workspace{ + ObjectMeta: metav1.ObjectMeta{ + Name: "testWorkspace-failedtocreateCR", + Namespace: "kaito", + Annotations: map[string]string{}, + }, + Resource: v1alpha1.ResourceSpec{ + Count: &gpuNodeCount, + InstanceType: "Standard_NC12s_v3", + LabelSelector: &metav1.LabelSelector{ + MatchLabels: map[string]string{ + "workspace.kaito.io/name": "testWorkspace", + }, + }, + }, + Inference: &v1alpha1.InferenceSpec{ + Preset: &v1alpha1.PresetSpec{ + PresetMeta: v1alpha1.PresetMeta{ + Name: "test-model", + }, + }, + }, + } +) + +var ( + MockWorkspaceSuccessful = v1alpha1.Workspace{ + ObjectMeta: metav1.ObjectMeta{ + Name: "testWorkspace-successful", Namespace: "kaito", Annotations: map[string]string{}, }, @@ -99,7 +153,7 @@ var ( ) var ( - MockWorkspaceWithComputeHash = &v1alpha1.Workspace{ + MockWorkspaceWithComputeHash = v1alpha1.Workspace{ ObjectMeta: metav1.ObjectMeta{ Name: "testWorkspace", Namespace: "kaito",