From f782f78e5bb83df4723ba679108ac121484253bb Mon Sep 17 00:00:00 2001
From: Karel Suta <ksuta@redhat.com>
Date: Tue, 30 Apr 2024 09:03:00 +0200
Subject: [PATCH] Add KueueWorkloads function to be used with gomega Eventually

---
 support/kueue.go | 23 ++++++++++++++---------
 1 file changed, 14 insertions(+), 9 deletions(-)

diff --git a/support/kueue.go b/support/kueue.go
index 3f582be..6283d18 100644
--- a/support/kueue.go
+++ b/support/kueue.go
@@ -89,18 +89,23 @@ func CreateKueueLocalQueue(t Test, namespace, clusterQueueName string) *kueuev1b
 	return localQueue
 }
 
-func GetKueueWorkloads(t Test, namespace string) []*kueuev1beta1.Workload {
-	t.T().Helper()
-
-	workloads, err := t.Client().Kueue().KueueV1beta1().Workloads(namespace).List(t.Ctx(), metav1.ListOptions{})
-	t.Expect(err).NotTo(gomega.HaveOccurred())
+func KueueWorkloads(t Test, namespace string) func(g gomega.Gomega) []*kueuev1beta1.Workload {
+	return func(g gomega.Gomega) []*kueuev1beta1.Workload {
+		workloads, err := t.Client().Kueue().KueueV1beta1().Workloads(namespace).List(t.Ctx(), metav1.ListOptions{})
+		g.Expect(err).NotTo(gomega.HaveOccurred())
+
+		workloadsp := []*kueuev1beta1.Workload{}
+		for _, v := range workloads.Items {
+			workloadsp = append(workloadsp, &v)
+		}
 
-	workloadsp := []*kueuev1beta1.Workload{}
-	for _, v := range workloads.Items {
-		workloadsp = append(workloadsp, &v)
+		return workloadsp
 	}
+}
 
-	return workloadsp
+func GetKueueWorkloads(t Test, namespace string) []*kueuev1beta1.Workload {
+	t.T().Helper()
+	return KueueWorkloads(t, namespace)(t)
 }
 
 func KueueWorkloadAdmitted(workload *kueuev1beta1.Workload) bool {