Skip to content

Commit

Permalink
chore: Add utils functions and move tests into seperate package (#351)
Browse files Browse the repository at this point in the history 
**Reason for Change**:
Add some future needed util functions, and move tests to seperate
package to prevent cyclical depdencies where API imports utils then
utils imports API. In order to use utils in API need to seperate out
test logic.
  • Loading branch information
@ishaansehgal99
ishaansehgal99 authored Apr 16, 2024
1 parent f88b529 commit 9e2ac24
Show file tree
Hide file tree
Showing 14 changed files with 213 additions and 175 deletions.
2 changes: 1 addition & 1 deletion pkg/controllers/workspace_controller.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -575,4 +575,4 @@ func (c *WorkspaceReconciler) watchMachines() handler.EventHandler {
},
}
})
}
}
158 changes: 79 additions & 79 deletions pkg/controllers/workspace_controller_test.go
View file

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion pkg/inference/preset-inferences.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -120,7 +120,7 @@ func CreatePresetInference(ctx context.Context, workspaceObj *kaitov1alpha1.Work

var volumes []corev1.Volume
var volumeMounts []corev1.VolumeMount
volume, volumeMount := utils.ConfigSHMVolume(workspaceObj)
volume, volumeMount := utils.ConfigSHMVolume(*workspaceObj.Resource.Count)
if volume.Name != "" {
volumes = append(volumes, volume)
}
Expand Down
14 changes: 7 additions & 7 deletions pkg/inference/preset-inferences_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,12 @@ package inference

import (
"context"
"github.com/azure/kaito/pkg/utils/test"
"reflect"
"strings"
"testing"

"github.com/azure/kaito/pkg/model"
"github.com/azure/kaito/pkg/utils"
"github.com/azure/kaito/pkg/utils/plugin"
"github.com/stretchr/testify/mock"
appsv1 "k8s.io/api/apps/v1"
Expand All @@ -19,19 +19,19 @@ import (
)

func TestCreatePresetInference(t *testing.T) {
utils.RegisterTestModel()
test.RegisterTestModel()
testcases := map[string]struct {
nodeCount int
modelName string
callMocks func(c *utils.MockClient)
callMocks func(c *test.MockClient)
workload string
expectedCmd string
}{

"test-model": {
nodeCount: 1,
modelName: "test-model",
callMocks: func(c *utils.MockClient) {
callMocks: func(c *test.MockClient) {
c.On("Create", mock.IsType(context.Background()), mock.IsType(&appsv1.Deployment{}), mock.Anything).Return(nil)
},
workload: "Deployment",
Expand All @@ -43,7 +43,7 @@ func TestCreatePresetInference(t *testing.T) {
"test-distributed-model": {
nodeCount: 1,
modelName: "test-distributed-model",
callMocks: func(c *utils.MockClient) {
callMocks: func(c *test.MockClient) {
c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&corev1.Service{}), mock.Anything).Return(nil)
c.On("Create", mock.IsType(context.Background()), mock.IsType(&appsv1.StatefulSet{}), mock.Anything).Return(nil)
},
Expand All @@ -54,10 +54,10 @@ func TestCreatePresetInference(t *testing.T) {

for k, tc := range testcases {
t.Run(k, func(t *testing.T) {
mockClient := utils.NewClient()
mockClient := test.NewClient()
tc.callMocks(mockClient)

workspace := utils.MockWorkspaceWithPreset
workspace := test.MockWorkspaceWithPreset
workspace.Resource.Count = &tc.nodeCount

useHeadlessSvc := false
Expand Down
16 changes: 8 additions & 8 deletions pkg/inference/template_inference_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,33 +5,33 @@ package inference
import (
"context"
"errors"
"github.com/azure/kaito/pkg/utils/test"
"testing"

"github.com/azure/kaito/pkg/utils"
"github.com/stretchr/testify/mock"
"gotest.tools/assert"
v1 "k8s.io/api/apps/v1"
)

func TestCreateTemplateInference(t *testing.T) {
testcases := map[string]struct {
callMocks func(c *utils.MockClient)
callMocks func(c *test.MockClient)
expectedError error
}{
"Fail to create template inference because deployment creation fails": {
callMocks: func(c *utils.MockClient) {
callMocks: func(c *test.MockClient) {
c.On("Create", mock.IsType(context.Background()), mock.IsType(&v1.Deployment{}), mock.Anything).Return(errors.New("Failed to create resource"))
},
expectedError: errors.New("Failed to create resource"),
},
"Successfully creates template inference because deployment already exists": {
callMocks: func(c *utils.MockClient) {
c.On("Create", mock.IsType(context.Background()), mock.IsType(&v1.Deployment{}), mock.Anything).Return(utils.IsAlreadyExistsError())
callMocks: func(c *test.MockClient) {
c.On("Create", mock.IsType(context.Background()), mock.IsType(&v1.Deployment{}), mock.Anything).Return(test.IsAlreadyExistsError())
},
expectedError: nil,
},
"Successfully creates template inference by creating a new deployment": {
callMocks: func(c *utils.MockClient) {
callMocks: func(c *test.MockClient) {
c.On("Create", mock.IsType(context.Background()), mock.IsType(&v1.Deployment{}), mock.Anything).Return(nil)
},
expectedError: nil,
Expand All @@ -40,10 +40,10 @@ func TestCreateTemplateInference(t *testing.T) {

for k, tc := range testcases {
t.Run(k, func(t *testing.T) {
mockClient := utils.NewClient()
mockClient := test.NewClient()
tc.callMocks(mockClient)

obj, err := CreateTemplateInference(context.Background(), utils.MockWorkspaceWithInferenceTemplate, mockClient)
obj, err := CreateTemplateInference(context.Background(), test.MockWorkspaceWithInferenceTemplate, mockClient)
if tc.expectedError == nil {
assert.Check(t, err == nil, "Not expected to return error")
assert.Check(t, obj != nil, "Return object should not be nil")
Expand Down
40 changes: 20 additions & 20 deletions pkg/machine/machine_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,10 @@ package machine
import (
"context"
"errors"
"github.com/azure/kaito/pkg/utils/test"
"testing"

"github.com/aws/karpenter-core/pkg/apis/v1alpha5"
"github.com/azure/kaito/pkg/utils"
"github.com/stretchr/testify/mock"
"gotest.tools/assert"
corev1 "k8s.io/api/core/v1"
Expand All @@ -19,18 +19,18 @@ import (

func TestCreateMachine(t *testing.T) {
testcases := map[string]struct {
callMocks func(c *utils.MockClient)
callMocks func(c *test.MockClient)
machineConditions apis.Conditions
expectedError error
}{
"Machine creation fails": {
callMocks: func(c *utils.MockClient) {
callMocks: func(c *test.MockClient) {
c.On("Create", mock.IsType(context.Background()), mock.IsType(&v1alpha5.Machine{}), mock.Anything).Return(errors.New("Failed to create machine"))
},
expectedError: errors.New("Failed to create machine"),
},
"Machine creation fails because SKU is not available": {
callMocks: func(c *utils.MockClient) {
callMocks: func(c *test.MockClient) {
c.On("Create", mock.IsType(context.Background()), mock.IsType(&v1alpha5.Machine{}), mock.Anything).Return(nil)
c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&v1alpha5.Machine{}), mock.Anything).Return(nil)
},
Expand All @@ -44,7 +44,7 @@ func TestCreateMachine(t *testing.T) {
expectedError: errors.New(ErrorInstanceTypesUnavailable),
},
"A machine is successfully created": {
callMocks: func(c *utils.MockClient) {
callMocks: func(c *test.MockClient) {
c.On("Create", mock.IsType(context.Background()), mock.IsType(&v1alpha5.Machine{}), mock.Anything).Return(nil)
c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&v1alpha5.Machine{}), mock.Anything).Return(nil)
},
Expand All @@ -60,10 +60,10 @@ func TestCreateMachine(t *testing.T) {

for k, tc := range testcases {
t.Run(k, func(t *testing.T) {
mockClient := utils.NewClient()
mockClient := test.NewClient()
tc.callMocks(mockClient)

mockMachine := &utils.MockMachine
mockMachine := &test.MockMachine
mockMachine.Status.Conditions = tc.machineConditions

err := CreateMachine(context.Background(), mockMachine, mockClient)
Expand All @@ -78,24 +78,24 @@ func TestCreateMachine(t *testing.T) {

func TestWaitForPendingMachines(t *testing.T) {
testcases := map[string]struct {
callMocks func(c *utils.MockClient)
callMocks func(c *test.MockClient)
machineConditions apis.Conditions
expectedError error
}{
"Fail to list machines because associated machines cannot be retrieved": {
callMocks: func(c *utils.MockClient) {
callMocks: func(c *test.MockClient) {
c.On("List", mock.IsType(context.Background()), mock.IsType(&v1alpha5.MachineList{}), mock.Anything).Return(errors.New("Failed to retrieve machines"))
},
expectedError: errors.New("Failed to retrieve machines"),
},
"Fail to list machines because machine status cannot be retrieved": {
callMocks: func(c *utils.MockClient) {
machineList := utils.MockMachineList
callMocks: func(c *test.MockClient) {
machineList := test.MockMachineList
relevantMap := c.CreateMapWithType(machineList)
c.CreateOrUpdateObjectInMap(&utils.MockMachine)
c.CreateOrUpdateObjectInMap(&test.MockMachine)

//insert machine objects into the map
for _, obj := range utils.MockMachineList.Items {
for _, obj := range test.MockMachineList.Items {
m := obj
objKey := client.ObjectKeyFromObject(&m)

Expand All @@ -114,13 +114,13 @@ func TestWaitForPendingMachines(t *testing.T) {
expectedError: errors.New("Fail to get machine"),
},
"Successfully waits for all pending machines": {
callMocks: func(c *utils.MockClient) {
machineList := utils.MockMachineList
callMocks: func(c *test.MockClient) {
machineList := test.MockMachineList
relevantMap := c.CreateMapWithType(machineList)
c.CreateOrUpdateObjectInMap(&utils.MockMachine)
c.CreateOrUpdateObjectInMap(&test.MockMachine)

//insert machine objects into the map
for _, obj := range utils.MockMachineList.Items {
for _, obj := range test.MockMachineList.Items {
m := obj
objKey := client.ObjectKeyFromObject(&m)

Expand All @@ -142,7 +142,7 @@ func TestWaitForPendingMachines(t *testing.T) {

for k, tc := range testcases {
t.Run(k, func(t *testing.T) {
mockClient := utils.NewClient()
mockClient := test.NewClient()
tc.callMocks(mockClient)

mockMachine := &v1alpha5.Machine{}
Expand All @@ -153,7 +153,7 @@ func TestWaitForPendingMachines(t *testing.T) {
mockClient.CreateOrUpdateObjectInMap(mockMachine)
}

err := WaitForPendingMachines(context.Background(), utils.MockWorkspaceWithPreset, mockClient)
err := WaitForPendingMachines(context.Background(), test.MockWorkspaceWithPreset, mockClient)
if tc.expectedError == nil {
assert.Check(t, err == nil, "Not expected to return error")
} else {
Expand All @@ -165,7 +165,7 @@ func TestWaitForPendingMachines(t *testing.T) {

func TestGenerateMachineManifiest(t *testing.T) {
t.Run("Should generate a machine object from the given workspace", func(t *testing.T) {
mockWorkspace := utils.MockWorkspaceWithPreset
mockWorkspace := test.MockWorkspaceWithPreset

machine := GenerateMachineManifest(context.Background(), "0", mockWorkspace)

Expand Down
12 changes: 6 additions & 6 deletions pkg/resources/manifests_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,20 +6,20 @@ package resources
import (
"context"
"fmt"
"github.com/azure/kaito/pkg/utils/test"
"reflect"

"testing"

kaitov1alpha1 "github.com/azure/kaito/api/v1alpha1"
"github.com/azure/kaito/pkg/utils"
v1 "k8s.io/api/core/v1"
)

func TestGenerateStatefulSetManifest(t *testing.T) {

t.Run("generate statefulset with headlessSvc", func(t *testing.T) {

workspace := utils.MockWorkspaceWithPreset
workspace := test.MockWorkspaceWithPreset

obj := GenerateStatefulSetManifest(context.TODO(), workspace,
"", //imageName
Expand Down Expand Up @@ -63,7 +63,7 @@ func TestGenerateStatefulSetManifest(t *testing.T) {
func TestGenerateDeploymentManifest(t *testing.T) {
t.Run("generate deployment", func(t *testing.T) {

workspace := utils.MockWorkspaceWithPreset
workspace := test.MockWorkspaceWithPreset

obj := GenerateDeploymentManifest(context.TODO(), workspace,
"", //imageName
Expand Down Expand Up @@ -103,7 +103,7 @@ func TestGenerateDeploymentManifest(t *testing.T) {
func TestGenerateDeploymentManifestWithPodTemplate(t *testing.T) {
t.Run("generate deployment with pod template", func(t *testing.T) {

workspace := utils.MockWorkspaceWithInferenceTemplate
workspace := test.MockWorkspaceWithInferenceTemplate

obj := GenerateDeploymentManifestWithPodTemplate(context.TODO(), workspace, nil)

Expand Down Expand Up @@ -142,7 +142,7 @@ func TestGenerateServiceManifest(t *testing.T) {

for _, isStatefulSet := range options {
t.Run(fmt.Sprintf("generate service, isStatefulSet %v", isStatefulSet), func(t *testing.T) {
workspace := utils.MockWorkspaceWithPreset
workspace := test.MockWorkspaceWithPreset
obj := GenerateServiceManifest(context.TODO(), workspace, v1.ServiceTypeClusterIP, isStatefulSet)

svcSelector := map[string]string{
Expand All @@ -161,7 +161,7 @@ func TestGenerateServiceManifest(t *testing.T) {
func TestGenerateHeadlessServiceManifest(t *testing.T) {

t.Run("generate headless service", func(t *testing.T) {
workspace := utils.MockWorkspaceWithPreset
workspace := test.MockWorkspaceWithPreset
obj := GenerateHeadlessServiceManifest(context.TODO(), workspace)

svcSelector := map[string]string{
Expand Down
Loading

0 comments on commit 9e2ac24

Please sign in to comment.