test: added ut for createAndValidateNode

go.mod

@@ -65,6 +65,7 @@ require (
 	github.com/prometheus/procfs v0.9.0 // indirect
 	github.com/prometheus/statsd_exporter v0.21.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
+	github.com/stretchr/objx v0.5.0 // indirect
 	go.opencensus.io v0.24.0 // indirect
 	go.uber.org/atomic v1.9.0 // indirect
 	go.uber.org/automaxprocs v1.4.0 // indirect
@@ -88,6 +89,7 @@ require (
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
+	gotest.tools v2.2.0+incompatible
 	k8s.io/apiextensions-apiserver v0.27.3 // indirect
 	k8s.io/component-base v0.27.3 // indirect
 	k8s.io/kube-openapi v0.0.0-20230501164219-8b0f38b5fd1f // indirect

go.sum

@@ -289,6 +289,7 @@ github.com/stoewer/go-strcase v1.2.0/go.mod h1:IBiWB2sKIp3wVVQ3Y035++gc+knqhUQag
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
+github.com/stretchr/objx v0.5.0 h1:1zr/of2m5FGMsad5YfcqgdqdWrIhu+EBEJRhR1U7z/c=
 github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
@@ -645,6 +646,8 @@ gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C
 gopkg.in/yaml.v3 v3.0.0-20200615113413-eeeca48fe776/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gotest.tools v2.2.0+incompatible h1:VsBPFP1AI068pPrMxtb/S8Zkgf9xEmTLJjfM+P5UIEo=
+gotest.tools v2.2.0+incompatible/go.mod h1:DsYFclhRJ6vuDpmuTbkuFWG+y2sxOXAzmJt81HFBacw=
 honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.0-20190106161140-3f1c8253044a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.0-20190418001031-e561f6794a2a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=

pkg/controllers/workspace_controller_test.go

@@ -3,12 +3,21 @@
 package controllers
 
 import (
+	"context"
+	"errors"
 	"reflect"
 	"sort"
 	"testing"
 
+	"github.com/aws/karpenter-core/pkg/apis/v1alpha5"
+	"github.com/azure/kaito/api/v1alpha1"
+	"github.com/azure/kaito/pkg/machine"
+	"github.com/azure/kaito/pkg/utils"
+	"github.com/stretchr/testify/mock"
+	"gotest.tools/assert"
 	corev1 "k8s.io/api/core/v1"
 	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"knative.dev/pkg/apis"
 )
 
 func TestSelectWorkspaceNodes(t *testing.T) {
@@ -209,3 +218,83 @@ func TestSelectWorkspaceNodes(t *testing.T) {
 		})
 	}
 }
+
+func TestCreateAndValidateNode(t *testing.T) {
+	mockClient := utils.NewClient()
+
+	mockMachine := utils.MockMachine
+
+	testcases := map[string]struct {
+		callMocks     func(c *utils.Client)
+		expectedError error
+	}{
+		"Node is not created because machine creation fails": {
+			callMocks: func(c *utils.Client) {
+				//create machine call should not return any errors
+				c.On("Create", mock.IsType(context.Background()), mock.IsType(&v1alpha5.Machine{}), mock.Anything).Return(nil)
+
+				machineConditions := apis.Conditions{
+					{
+						Type:    v1alpha5.MachineLaunched,
+						Status:  corev1.ConditionFalse,
+						Message: machine.ErrorInstanceTypesUnavailable,
+					},
+				}
+
+				mockMachine.Status.Conditions = machineConditions
+
+				//insert mock machine in map so mock Get call retrieves this object
+				c.InsertObjectInMap(&mockMachine)
+
+				c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&v1alpha5.Machine{}), mock.Anything).Return(nil)
+				c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&v1alpha1.Workspace{}), mock.Anything).Return(nil)
+				c.StatusMock.On("Update", mock.IsType(context.Background()), mock.IsType(&v1alpha1.Workspace{}), mock.Anything).Return(nil)
+			},
+			expectedError: errors.New(machine.ErrorInstanceTypesUnavailable),
+		},
+		"A machine is successfully created": {
+			callMocks: func(c *utils.Client) {
+				//create machine call should not return any errors
+				c.On("Create", mock.IsType(context.Background()), mock.IsType(&v1alpha5.Machine{}), mock.Anything).Return(nil)
+
+				machineConditions := apis.Conditions{
+					{
+						Type:   apis.ConditionReady,
+						Status: corev1.ConditionTrue,
+					},
+				}
+
+				mockMachine.Status.Conditions = machineConditions
+
+				//insert mock machine in map so mock Get call retrieves this object
+				c.InsertObjectInMap(&mockMachine)
+
+				//get machine call should not return any errors
+				c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&v1alpha5.Machine{}), mock.Anything).Return(nil)
+				//get machine call should not return any errors
+				c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&corev1.Node{}), mock.Anything).Return(nil)
+			},
+			expectedError: nil,
+		},
+	}
+
+	for k, tc := range testcases {
+		t.Run(k, func(t *testing.T) {
+			tc.callMocks(mockClient)
+
+			reconciler := &WorkspaceReconciler{
+				Client: mockClient,
+				Scheme: utils.NewTestScheme(),
+			}
+			ctx := context.Background()
+
+			node, err := reconciler.createAndValidateNode(ctx, utils.MockWorkspace)
+			if tc.expectedError == nil {
+				assert.Check(t, err == nil, "Not expected to return error")
+				assert.Check(t, node != nil, "Response node should not be nil")
+			} else {
+				assert.Equal(t, tc.expectedError.Error(), err.Error())
+			}
+		})
+	}
+}

pkg/utils/mockClient.go

@@ -0,0 +1,146 @@
+package utils
+
+import (
+	"context"
+	"reflect"
+
+	"github.com/stretchr/testify/mock"
+	"k8s.io/apimachinery/pkg/api/meta"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/runtime/schema"
+	"k8s.io/apimachinery/pkg/types"
+	k8sClient "sigs.k8s.io/controller-runtime/pkg/client"
+)
+
+// Client is a mock for the controller-runtime dynamic client interface.
+type Client struct {
+	mock.Mock
+
+	ObjectMap  map[reflect.Type]map[k8sClient.ObjectKey]k8sClient.Object
+	StatusMock *StatusClient
+}
+
+var _ k8sClient.Client = &Client{}
+
+func NewClient() *Client {
+	return &Client{
+		StatusMock: &StatusClient{},
+		ObjectMap:  map[reflect.Type]map[k8sClient.ObjectKey]k8sClient.Object{},
+	}
+}
+
+// Retrieves or creates a map associated with the type of obj
+func (c *Client) ensureMapFor(obj k8sClient.Object) map[k8sClient.ObjectKey]k8sClient.Object {
+	t := reflect.TypeOf(obj)
+	if _, ok := c.ObjectMap[t]; !ok {
+		//create a new map with the object key if it doesn't exist
+		c.ObjectMap[t] = map[k8sClient.ObjectKey]k8sClient.Object{}
+	}
+	return c.ObjectMap[t]
+}
+
+func (c *Client) InsertObjectInMap(obj k8sClient.Object) {
+	relevantMap := c.ensureMapFor(obj)
+	objKey := k8sClient.ObjectKeyFromObject(obj)
+
+	relevantMap[objKey] = obj
+}
+
+// StatusClient interface
+
+func (c *Client) Status() k8sClient.StatusWriter {
+	return c.StatusMock
+}
+
+// Reader interface
+
+func (c *Client) Get(ctx context.Context, key types.NamespacedName, obj k8sClient.Object, opts ...k8sClient.GetOption) error {
+	relevantMap := c.ensureMapFor(obj)
+
+	for _, val := range relevantMap {
+		v := reflect.ValueOf(obj).Elem()
+		v.Set(reflect.ValueOf(val).Elem())
+		break
+	}
+
+	args := c.Called(ctx, key, obj, opts)
+	return args.Error(0)
+}
+
+func (c *Client) List(ctx context.Context, list k8sClient.ObjectList, opts ...k8sClient.ListOption) error {
+	args := c.Called(ctx, list, opts)
+	return args.Error(0)
+}
+
+// Writer interface
+
+func (c *Client) Create(ctx context.Context, obj k8sClient.Object, opts ...k8sClient.CreateOption) error {
+	args := c.Called(ctx, obj, opts)
+	return args.Error(0)
+}
+
+func (c *Client) Delete(ctx context.Context, obj k8sClient.Object, opts ...k8sClient.DeleteOption) error {
+	args := c.Called(ctx, obj, opts)
+	return args.Error(0)
+}
+
+func (c *Client) Update(ctx context.Context, obj k8sClient.Object, opts ...k8sClient.UpdateOption) error {
+	args := c.Called(ctx, obj, opts)
+	return args.Error(0)
+}
+
+func (c *Client) Patch(ctx context.Context, obj k8sClient.Object, patch k8sClient.Patch, opts ...k8sClient.PatchOption) error {
+	args := c.Called(ctx, obj, patch, opts)
+	return args.Error(0)
+}
+
+func (c *Client) DeleteAllOf(ctx context.Context, obj k8sClient.Object, opts ...k8sClient.DeleteAllOfOption) error {
+	args := c.Called(ctx, obj, opts)
+	return args.Error(0)
+}
+
+// SubResource implements client.Client
+func (*Client) SubResource(subResource string) k8sClient.SubResourceClient {
+	panic("unimplemented")
+}
+
+// GroupVersionKindFor implements client.Client
+func (*Client) GroupVersionKindFor(obj runtime.Object) (schema.GroupVersionKind, error) {
+	panic("unimplemented")
+}
+
+// IsObjectNamespaced implements client.Client
+func (*Client) IsObjectNamespaced(obj runtime.Object) (bool, error) {
+	panic("unimplemented")
+}
+
+func (c *Client) Scheme() *runtime.Scheme {
+	args := c.Called()
+	return args.Get(0).(*runtime.Scheme)
+}
+
+func (c *Client) RESTMapper() meta.RESTMapper {
+	args := c.Called()
+	return args.Get(0).(meta.RESTMapper)
+}
+
+type StatusClient struct {
+	mock.Mock
+}
+
+func (c *StatusClient) Create(ctx context.Context, obj k8sClient.Object, subResource k8sClient.Object, opts ...k8sClient.SubResourceCreateOption) error {
+	args := c.Called(ctx, obj, opts)
+	return args.Error(0)
+}
+
+func (c *StatusClient) Patch(ctx context.Context, obj k8sClient.Object, patch k8sClient.Patch, opts ...k8sClient.SubResourcePatchOption) error {
+	args := c.Called(ctx, obj, opts)
+	return args.Error(0)
+}
+
+func (c *StatusClient) Update(ctx context.Context, obj k8sClient.Object, opts ...k8sClient.SubResourceUpdateOption) error {
+	args := c.Called(ctx, obj, opts)
+	return args.Error(0)
+}
+
+var _ k8sClient.StatusWriter = &StatusClient{}

pkg/utils/testUtils.go

@@ -0,0 +1,68 @@
+package utils
+
+import (
+	"github.com/aws/karpenter-core/pkg/apis/v1alpha5"
+	"github.com/azure/kaito/api/v1alpha1"
+	appsv1 "k8s.io/api/apps/v1"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+)
+
+var (
+	MockWorkspace = &v1alpha1.Workspace{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "testWorkspace",
+			Namespace: "kaito",
+		},
+		Resource: v1alpha1.ResourceSpec{
+			InstanceType: "Standard_NC12s_v3",
+		},
+		Inference: v1alpha1.InferenceSpec{
+			Preset: &v1alpha1.PresetSpec{
+				PresetMeta: v1alpha1.PresetMeta{
+					Name: "llama-2-7b-chat",
+				},
+			},
+		},
+	}
+)
+
+var (
+	machineLabels = map[string]string{
+		"karpenter.sh/provisioner-name": "default",
+		"kaito.sh/workspace":            "none",
+	}
+)
+
+var (
+	MockMachine = v1alpha5.Machine{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:   "testmachine",
+			Labels: machineLabels,
+		},
+		Spec: v1alpha5.MachineSpec{
+			MachineTemplateRef: &v1alpha5.MachineTemplateRef{
+				Name: "test-machine",
+			},
+			Requirements: []corev1.NodeSelectorRequirement{
+				{
+					Key:      corev1.LabelInstanceTypeStable,
+					Operator: corev1.NodeSelectorOpIn,
+					Values:   []string{"Standard_NC12s_v3"},
+				},
+				{
+					Key:      "karpenter.sh/provisioner-name",
+					Operator: corev1.NodeSelectorOpIn,
+					Values:   []string{"default"},
+				},
+			},
+		},
+	}
+)
+
+func NewTestScheme() *runtime.Scheme {
+	testScheme := runtime.NewScheme()
+	_ = appsv1.AddToScheme(testScheme)
+	return testScheme
+}