Skip to content

Commit

Permalink
Add more coverage for controllerrevision unit test
Browse files Browse the repository at this point in the history 
Signed-off-by: Bangqi Zhu <bangqizhu@microsoft.com>
  • Loading branch information
Bangqi Zhu committed Jul 22, 2024
1 parent c70a53a commit d24fc36
Show file tree
Hide file tree
Showing 4 changed files with 128 additions and 20 deletions.
5 changes: 2 additions & 3 deletions pkg/controllers/workspace_controller.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down Expand Up @@ -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)
}
Expand Down
82 changes: 68 additions & 14 deletions pkg/controllers/workspace_controller_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ package controllers
import (
"context"
"errors"
"fmt"
"os"
"reflect"
"sort"
Expand Down Expand Up @@ -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 {
Expand All @@ -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)
}
})
}
}
1 change: 1 addition & 0 deletions pkg/utils/consts/consts.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,4 +12,5 @@ const (
AWSCloudName = "aws"
GPUString = "gpu"
SKUString = "sku"
MaxRevisionHistoryLimit = 10
)
60 changes: 57 additions & 3 deletions pkg/utils/test/testUtils.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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{},
},
Expand All @@ -99,7 +153,7 @@ var (
)

var (
MockWorkspaceWithComputeHash = &v1alpha1.Workspace{
MockWorkspaceWithComputeHash = v1alpha1.Workspace{
ObjectMeta: metav1.ObjectMeta{
Name: "testWorkspace",
Namespace: "kaito",
Expand Down

0 comments on commit d24fc36

Please sign in to comment.