Add KueueWorkloads function to be used with gomega Eventually

project-codeflare · Apr 30, 2024 · f782f78 · f782f78
1 parent 3801016
commit f782f78
Showing 1 changed file with 14 additions and 9 deletions.
diff --git 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 {