From 27d5cda8114bb5c6b439b1f8814a80aedcbb767f Mon Sep 17 00:00:00 2001
From: Maciej Szulik <maszulik@redhat.com>
Date: Mon, 6 Mar 2023 12:43:53 +0100
Subject: [PATCH] Simplify waitForStatusCurrentReplicas helper

---
 test/e2e/apps/statefulset.go           |  2 +-
 test/e2e/apps/wait.go                  | 13 +++++++++++++
 test/e2e/framework/statefulset/wait.go | 25 -------------------------
 3 files changed, 14 insertions(+), 26 deletions(-)

diff --git a/test/e2e/apps/statefulset.go b/test/e2e/apps/statefulset.go
index c494e45f1c4..de5c4649094 100644
--- a/test/e2e/apps/statefulset.go
+++ b/test/e2e/apps/statefulset.go
@@ -1433,7 +1433,7 @@ var _ = SIGDescribe("StatefulSet", func() {
 			framework.ExpectNoError(err)
 
 			// wait for the pod to be recreated
-			e2estatefulset.WaitForStatusCurrentReplicas(c, ss, 1)
+			waitForStatusCurrentReplicas(ctx, c, ss, 1)
 			_, err = c.CoreV1().Pods(ns).Get(ctx, podName, metav1.GetOptions{})
 			framework.ExpectNoError(err)
 
diff --git a/test/e2e/apps/wait.go b/test/e2e/apps/wait.go
index ad0c84d74d3..92c25f067e5 100644
--- a/test/e2e/apps/wait.go
+++ b/test/e2e/apps/wait.go
@@ -97,6 +97,19 @@ func waitForStatus(ctx context.Context, c clientset.Interface, set *appsv1.State
 	return set
 }
 
+// waitForStatus waits for the StatefulSetStatus's CurrentReplicas to be equal to expectedReplicas
+// The returned StatefulSet contains such a StatefulSetStatus
+func waitForStatusCurrentReplicas(ctx context.Context, c clientset.Interface, set *appsv1.StatefulSet, expectedReplicas int32) *appsv1.StatefulSet {
+	e2estatefulset.WaitForState(ctx, c, set, func(set2 *appsv1.StatefulSet, pods *v1.PodList) (bool, error) {
+		if set2.Status.ObservedGeneration >= set.Generation && set2.Status.CurrentReplicas == expectedReplicas {
+			set = set2
+			return true, nil
+		}
+		return false, nil
+	})
+	return set
+}
+
 // waitForPodNotReady waits for the Pod named podName in set to exist and to not have a Ready condition.
 func waitForPodNotReady(ctx context.Context, c clientset.Interface, set *appsv1.StatefulSet, podName string) (*appsv1.StatefulSet, *v1.PodList) {
 	var pods *v1.PodList
diff --git a/test/e2e/framework/statefulset/wait.go b/test/e2e/framework/statefulset/wait.go
index 7e4d7053788..f05c50dc353 100644
--- a/test/e2e/framework/statefulset/wait.go
+++ b/test/e2e/framework/statefulset/wait.go
@@ -171,31 +171,6 @@ func WaitForStatusReplicas(ctx context.Context, c clientset.Interface, ss *appsv
 	}
 }
 
-// WaitForStatusCurrentReplicas waits for the ss.Status.CurrentReplicas to be equal to expectedReplicas
-func WaitForStatusCurrentReplicas(c clientset.Interface, ss *appsv1.StatefulSet, expectedReplicas int32) {
-	framework.Logf("Waiting for statefulset status.currentReplicas updated to %d", expectedReplicas)
-
-	ns, name := ss.Namespace, ss.Name
-	pollErr := wait.PollImmediate(StatefulSetPoll, StatefulSetTimeout,
-		func() (bool, error) {
-			ssGet, err := c.AppsV1().StatefulSets(ns).Get(context.TODO(), name, metav1.GetOptions{})
-			if err != nil {
-				return false, err
-			}
-			if ssGet.Status.ObservedGeneration < ss.Generation {
-				return false, nil
-			}
-			if ssGet.Status.CurrentReplicas != expectedReplicas {
-				framework.Logf("Waiting for stateful set status.currentReplicas to become %d, currently %d", expectedReplicas, ssGet.Status.CurrentReplicas)
-				return false, nil
-			}
-			return true, nil
-		})
-	if pollErr != nil {
-		framework.Failf("Failed waiting for stateful set status.currentReplicas updated to %d: %v", expectedReplicas, pollErr)
-	}
-}
-
 // Saturate waits for all Pods in ss to become Running and Ready.
 func Saturate(ctx context.Context, c clientset.Interface, ss *appsv1.StatefulSet) {
 	var i int32