Updated CreateLocalQueue function parameters to use variadic arguments of type *kueuev1beta1.localQueue

Author: abhijeet-dhumal

support/kueue.go

@@ -18,7 +18,6 @@ package support
 
 import (
 	"github.com/onsi/gomega"
-
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	kueuev1beta1 "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 )
@@ -65,7 +64,21 @@ func CreateKueueClusterQueue(t Test, clusterQueueSpec kueuev1beta1.ClusterQueueS
 	return clusterQueue
 }
 
-func CreateKueueLocalQueue(t Test, namespace, clusterQueueName string, annotations map[string]string) *kueuev1beta1.LocalQueue {
+type DefaultLocalQueueOption struct {
+	annotations map[string]string
+}
+
+func (d DefaultLocalQueueOption) applyTo(to *kueuev1beta1.LocalQueue) error {
+	if to.Annotations == nil {
+		to.Annotations = make(map[string]string)
+	}
+	for key, value := range d.annotations {
+		to.Annotations[key] = value
+	}
+	return nil
+}
+
+func CreateKueueLocalQueue(t Test, namespace string, clusterQueueName string, options ...Option[*kueuev1beta1.LocalQueue]) *kueuev1beta1.LocalQueue {
 	t.T().Helper()
 
 	localQueue := &kueuev1beta1.LocalQueue{
@@ -76,14 +89,18 @@ func CreateKueueLocalQueue(t Test, namespace, clusterQueueName string, annotatio
 		ObjectMeta: metav1.ObjectMeta{
 			GenerateName: "lq-",
 			Namespace:    namespace,
-			Annotations:  annotations,
 		},
 		Spec: kueuev1beta1.LocalQueueSpec{
 			ClusterQueue: kueuev1beta1.ClusterQueueReference(clusterQueueName),
 		},
 	}
 
-	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)
 

support/kueue_test.go

@@ -52,10 +52,12 @@ func TestCreateKueueClusterQueue(t *testing.T) {
 func TestCreateKueueLocalQueue(t *testing.T) {
 	test := NewTest(t)
 
-	annotations := map[string]string{
-		"kueue.x-k8s.io/default-queue": "true",
+	options := DefaultLocalQueueOption{
+		annotations: map[string]string{
+			"kueue.x-k8s.io/default-queue": "true",
+		},
 	}
-	lq := CreateKueueLocalQueue(test, "ns-1", "cq-1", annotations)
+	lq := CreateKueueLocalQueue(test, "ns-1", "cq-1", options)
 
 	test.Expect(lq).To(gomega.Not(gomega.BeNil()))
 	test.Expect(lq.GenerateName).To(gomega.Equal("lq-"))