From 0746b63cc526f656d7cb0e21fe9d692f598ac286 Mon Sep 17 00:00:00 2001
From: abhijeet-dhumal <abhijeetdhumal652@gmail.com>
Date: Wed, 31 Jul 2024 12:28:26 +0530
Subject: [PATCH] Add local-queue annotation to use created local queue by
 default in given namespace

---
 support/environment.go |  6 ++++++
 support/kueue.go       | 22 ++++++++++++++++++++--
 support/kueue_test.go  | 12 ++++++++++++
 3 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/support/environment.go b/support/environment.go
index 287f4c2..d377839 100644
--- a/support/environment.go
+++ b/support/environment.go
@@ -54,6 +54,7 @@ const (
 
 	// Storage bucket credentials
 	storageDefaultEndpoint = "AWS_DEFAULT_ENDPOINT"
+	storageDefaultRegion   = "AWS_DEFAULT_REGION"
 	storageAccessKeyId     = "AWS_ACCESS_KEY_ID"
 	storageSecretKey       = "AWS_SECRET_ACCESS_KEY"
 	storageBucketName      = "AWS_STORAGE_BUCKET"
@@ -133,6 +134,11 @@ func GetStorageBucketDefaultEndpoint() (string, bool) {
 	return storage_endpoint, exists
 }
 
+func GetStorageBucketDefaultRegion() (string, bool) {
+	storage_default_region, exists := os.LookupEnv(storageDefaultRegion)
+	return storage_default_region, exists
+}
+
 func GetStorageBucketAccessKeyId() (string, bool) {
 	storage_access_key_id, exists := os.LookupEnv(storageAccessKeyId)
 	return storage_access_key_id, exists
diff --git a/support/kueue.go b/support/kueue.go
index 6283d18..e5bd2a5 100644
--- a/support/kueue.go
+++ b/support/kueue.go
@@ -65,7 +65,20 @@ func CreateKueueClusterQueue(t Test, clusterQueueSpec kueuev1beta1.ClusterQueueS
 	return clusterQueue
 }
 
-func CreateKueueLocalQueue(t Test, namespace, clusterQueueName string) *kueuev1beta1.LocalQueue {
+var AsDefaultQueue = defaultLocalQueueOption{}
+
+type defaultLocalQueueOption struct {
+}
+
+func (d defaultLocalQueueOption) applyTo(to *kueuev1beta1.LocalQueue) error {
+	if to.Annotations == nil {
+		to.Annotations = make(map[string]string)
+	}
+	to.Annotations["kueue.x-k8s.io/default-queue"] = "true"
+	return nil
+}
+
+func CreateKueueLocalQueue(t Test, namespace string, clusterQueueName string, options ...Option[*kueuev1beta1.LocalQueue]) *kueuev1beta1.LocalQueue {
 	t.T().Helper()
 
 	localQueue := &kueuev1beta1.LocalQueue{
@@ -82,7 +95,12 @@ func CreateKueueLocalQueue(t Test, namespace, clusterQueueName string) *kueuev1b
 		},
 	}
 
-	localQueue, err := t.Client().Kueue().KueueV1beta1().LocalQueues(namespace).Create(t.Ctx(), localQueue, metav1.CreateOptions{})
+	//Apply options
+	for _, opt := range options {
+		t.Expect(opt.applyTo(localQueue)).To(gomega.Succeed())
+	}
+
+	localQueue, err := t.Client().Kueue().KueueV1beta1().LocalQueues(localQueue.Namespace).Create(t.Ctx(), localQueue, metav1.CreateOptions{})
 	t.Expect(err).NotTo(gomega.HaveOccurred())
 	t.T().Logf("Created Kueue LocalQueue %s/%s successfully", localQueue.Namespace, localQueue.Name)
 
diff --git a/support/kueue_test.go b/support/kueue_test.go
index 54f5ea1..806b1ce 100644
--- a/support/kueue_test.go
+++ b/support/kueue_test.go
@@ -56,8 +56,20 @@ func TestCreateKueueLocalQueue(t *testing.T) {
 
 	test.Expect(lq).To(gomega.Not(gomega.BeNil()))
 	test.Expect(lq.GenerateName).To(gomega.Equal("lq-"))
+	annotationKey := "kueue.x-k8s.io/default-queue"
+	_, exists := lq.Annotations[annotationKey]
+	test.Expect(exists).To(gomega.BeFalse(), "Annotation key %s should not exist", annotationKey)
 	test.Expect(lq.Namespace).To(gomega.Equal("ns-1"))
 	test.Expect(lq.Spec.ClusterQueue).To(gomega.Equal(kueuev1beta1.ClusterQueueReference("cq-1")))
+
+	default_lq := CreateKueueLocalQueue(test, "ns-2", "cq-2", AsDefaultQueue)
+
+	test.Expect(default_lq).To(gomega.Not(gomega.BeNil()))
+	test.Expect(default_lq.GenerateName).To(gomega.Equal("lq-"))
+	test.Expect(default_lq.Annotations["kueue.x-k8s.io/default-queue"]).To(gomega.Equal("true"))
+	test.Expect(default_lq.Namespace).To(gomega.Equal("ns-2"))
+	test.Expect(default_lq.Spec.ClusterQueue).To(gomega.Equal(kueuev1beta1.ClusterQueueReference("cq-2")))
+
 }
 
 func TestGetKueueWorkloads(t *testing.T) {