feat: Add karpenter labels to nodeClaim (#373)

**Reason for Change**:
- Add `DoNotDisrupt=true` label to prevent Karpenter from scaling down.
- Make sure to add Fake nodepool name to prevent Karpenter from scaling
up.
- Add Unit tests for generate Nodeclaim and for workspace status.

**Requirements**

- [x] added unit tests and e2e tests (if applicable).

**Issue Fixed**:
<!-- If this PR fixes GitHub issue 4321, add "Fixes #4321" to the next
line. -->

**Notes for Reviewers**:

Signed-off-by: Heba Elayoty <hebaelayoty@gmail.com>

pkg/controllers/workspace_status_test.go

@@ -0,0 +1,172 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+package controllers
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	kaitov1alpha1 "github.com/azure/kaito/api/v1alpha1"
+	"github.com/azure/kaito/pkg/utils/test"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime/schema"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+)
+
+func TestUpdateWorkspaceStatus(t *testing.T) {
+	t.Run("Should update workspace status successfully", func(t *testing.T) {
+		mockClient := test.NewClient()
+		reconciler := &WorkspaceReconciler{
+			Client: mockClient,
+			Scheme: test.NewTestScheme(),
+		}
+		ctx := context.Background()
+		workspace := test.MockWorkspaceDistributedModel
+		condition := metav1.Condition{
+			Type:    "TestCondition",
+			Status:  metav1.ConditionStatus("True"),
+			Reason:  "TestReason",
+			Message: "TestMessage",
+		}
+		workerNodes := []string{"node1", "node2"}
+
+		mockClient.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&kaitov1alpha1.Workspace{}), mock.Anything).Return(nil)
+		mockClient.StatusMock.On("Update", mock.IsType(context.Background()), mock.IsType(&kaitov1alpha1.Workspace{}), mock.Anything).Return(nil)
+
+		err := reconciler.updateWorkspaceStatus(ctx, &client.ObjectKey{Name: workspace.Name, Namespace: workspace.Namespace}, &condition, workerNodes)
+		assert.Nil(t, err)
+	})
+
+	t.Run("Should return error when Get operation fails", func(t *testing.T) {
+		mockClient := test.NewClient()
+		reconciler := &WorkspaceReconciler{
+			Client: mockClient,
+			Scheme: test.NewTestScheme(),
+		}
+		ctx := context.Background()
+		workspace := test.MockWorkspaceDistributedModel
+		condition := metav1.Condition{
+			Type:    "TestCondition",
+			Status:  metav1.ConditionStatus("True"),
+			Reason:  "TestReason",
+			Message: "TestMessage",
+		}
+		workerNodes := []string{"node1", "node2"}
+
+		mockClient.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&kaitov1alpha1.Workspace{}), mock.Anything).Return(errors.New("Get operation failed"))
+
+		err := reconciler.updateWorkspaceStatus(ctx, &client.ObjectKey{Name: workspace.Name, Namespace: workspace.Namespace}, &condition, workerNodes)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("Should return nil when workspace is not found", func(t *testing.T) {
+		mockClient := test.NewClient()
+		reconciler := &WorkspaceReconciler{
+			Client: mockClient,
+			Scheme: test.NewTestScheme(),
+		}
+		ctx := context.Background()
+		workspace := test.MockWorkspaceDistributedModel
+		condition := metav1.Condition{
+			Type:    "TestCondition",
+			Status:  metav1.ConditionStatus("True"),
+			Reason:  "TestReason",
+			Message: "TestMessage",
+		}
+		workerNodes := []string{"node1", "node2"}
+
+		mockClient.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&kaitov1alpha1.Workspace{}), mock.Anything).Return(apierrors.NewNotFound(schema.GroupResource{}, "workspace"))
+
+		err := reconciler.updateWorkspaceStatus(ctx, &client.ObjectKey{Name: workspace.Name, Namespace: workspace.Namespace}, &condition, workerNodes)
+		assert.Nil(t, err)
+	})
+}
+
+func TestUpdateStatusConditionIfNotMatch(t *testing.T) {
+	t.Run("Should not update when condition matches", func(t *testing.T) {
+		mockClient := test.NewClient()
+		reconciler := &WorkspaceReconciler{
+			Client: mockClient,
+			Scheme: test.NewTestScheme(),
+		}
+		ctx := context.Background()
+		workspace := test.MockWorkspaceDistributedModel
+		conditionType := kaitov1alpha1.ConditionType("TestCondition")
+		conditionStatus := metav1.ConditionStatus("True")
+		conditionReason := "TestReason"
+		conditionMessage := "TestMessage"
+
+		workspace.Status.Conditions = []metav1.Condition{
+			{
+				Type:    string(conditionType),
+				Status:  conditionStatus,
+				Reason:  conditionReason,
+				Message: conditionMessage,
+			},
+		}
+
+		err := reconciler.updateStatusConditionIfNotMatch(ctx, workspace, conditionType, conditionStatus, conditionReason, conditionMessage)
+		assert.Nil(t, err)
+	})
+
+	t.Run("Should update when condition does not match", func(t *testing.T) {
+		mockClient := test.NewClient()
+		reconciler := &WorkspaceReconciler{
+			Client: mockClient,
+			Scheme: test.NewTestScheme(),
+		}
+		ctx := context.Background()
+		workspace := test.MockWorkspaceDistributedModel
+		conditionType := kaitov1alpha1.ConditionType("TestCondition")
+		conditionStatus := metav1.ConditionStatus("True")
+		conditionReason := "TestReason"
+		conditionMessage := "TestMessage"
+
+		workspace.Status.Conditions = []metav1.Condition{
+			{
+				Type:    string(conditionType),
+				Status:  conditionStatus,
+				Reason:  conditionReason,
+				Message: "DifferentMessage",
+			},
+		}
+		mockClient.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&kaitov1alpha1.Workspace{}), mock.Anything).Return(nil)
+		mockClient.StatusMock.On("Update", mock.IsType(context.Background()), mock.IsType(&kaitov1alpha1.Workspace{}), mock.Anything).Return(nil)
+
+		err := reconciler.updateStatusConditionIfNotMatch(ctx, workspace, conditionType, conditionStatus, conditionReason, conditionMessage)
+		assert.Nil(t, err)
+	})
+
+	t.Run("Should update when condition is not found", func(t *testing.T) {
+		mockClient := test.NewClient()
+		reconciler := &WorkspaceReconciler{
+			Client: mockClient,
+			Scheme: test.NewTestScheme(),
+		}
+		ctx := context.Background()
+		workspace := test.MockWorkspaceDistributedModel
+		conditionType := kaitov1alpha1.ConditionType("TestCondition")
+		conditionStatus := metav1.ConditionStatus("True")
+		conditionReason := "TestReason"
+		conditionMessage := "TestMessage"
+
+		workspace.Status.Conditions = []metav1.Condition{
+			{
+				Type:    "DifferentCondition",
+				Status:  conditionStatus,
+				Reason:  conditionReason,
+				Message: conditionMessage,
+			},
+		}
+		mockClient.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&kaitov1alpha1.Workspace{}), mock.Anything).Return(nil)
+		mockClient.StatusMock.On("Update", mock.IsType(context.Background()), mock.IsType(&kaitov1alpha1.Workspace{}), mock.Anything).Return(nil)
+
+		err := reconciler.updateStatusConditionIfNotMatch(ctx, workspace, conditionType, conditionStatus, conditionReason, conditionMessage)
+		assert.Nil(t, err)
+	})
+}

pkg/nodeclaim/nodeclaim.go

@@ -37,17 +37,19 @@ var (
 
 // GenerateNodeClaimManifest generates a nodeClaim object from the given workspace.
 func GenerateNodeClaimManifest(ctx context.Context, storageRequirement string, workspaceObj *kaitov1alpha1.Workspace) *v1beta1.NodeClaim {
-	digest := sha256.Sum256([]byte(workspaceObj.Namespace + workspaceObj.Name + time.Now().Format("2006-01-02 15:04:05.000000000"))) // We make sure the nodeClaim name is not fixed to the a workspace
+	digest := sha256.Sum256([]byte(workspaceObj.Namespace + workspaceObj.Name + time.Now().
+		Format("2006-01-02 15:04:05.000000000"))) // We make sure the nodeClaim name is not fixed to the workspace
 	nodeClaimName := "ws" + hex.EncodeToString(digest[0:])[0:9]
+
 	nodeClaimLabels := map[string]string{
-		LabelNodePool:                         KaitoNodePoolName,
+		LabelNodePool:                         KaitoNodePoolName, // Fake nodepool name to prevent Karpenter from scaling up.
 		kaitov1alpha1.LabelWorkspaceName:      workspaceObj.Name,
 		kaitov1alpha1.LabelWorkspaceNamespace: workspaceObj.Namespace,
+		v1beta1.DoNotDisruptAnnotationKey:     "true", // To prevent Karpenter from scaling down.
 	}
 	if workspaceObj.Resource.LabelSelector != nil &&
 		len(workspaceObj.Resource.LabelSelector.MatchLabels) != 0 {
 		nodeClaimLabels = lo.Assign(nodeClaimLabels, workspaceObj.Resource.LabelSelector.MatchLabels)
-
 	}
 
 	return &v1beta1.NodeClaim{

pkg/nodeclaim/nodeclaim_test.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 	"testing"
 
+	kaitov1alpha1 "github.com/azure/kaito/api/v1alpha1"
 	"github.com/azure/kaito/pkg/utils/test"
 	"github.com/stretchr/testify/mock"
 	"gotest.tools/assert"
@@ -171,5 +172,10 @@ func TestGenerateNodeClaimManifest(t *testing.T) {
 
 		assert.Check(t, nodeClaim != nil, "NodeClaim must not be nil")
 		assert.Equal(t, nodeClaim.Namespace, mockWorkspace.Namespace, "NodeClaim must have same namespace as workspace")
+		assert.Equal(t, nodeClaim.Labels[kaitov1alpha1.LabelWorkspaceName], mockWorkspace.Name, "label must have same workspace name as workspace")
+		assert.Equal(t, nodeClaim.Labels[kaitov1alpha1.LabelWorkspaceNamespace], mockWorkspace.Namespace, "label must have same workspace namespace as workspace")
+		assert.Equal(t, nodeClaim.Labels[LabelNodePool], KaitoNodePoolName, "label must have same labels as workspace label selector")
+		assert.Equal(t, nodeClaim.Labels[v1beta1.DoNotDisruptAnnotationKey], "true", "label must have do not disrupt annotation")
+		assert.Equal(t, nodeClaim.Spec.Requirements[0].Values[0], mockWorkspace.Resource.InstanceType, "NodeClaim must have same instance type as workspace")
 	})
 }