Refactor deprecated wait.Poll functions

test/e2e/setup/setup.go

@@ -237,7 +237,7 @@ func DeployOperator(ctx context.Context, config TestConfig, resourceName string,
 		return err
 	}
 
-	if err := wait.PollImmediate(time.Second, 60*time.Second, hasDeploymentRequiredReplicas(ctx, &dep)); err != nil {
+	if err := wait.PollUntilContextTimeout(ctx, time.Second, 60*time.Second, true, hasDeploymentRequiredReplicas(&dep)); err != nil {
 		return errors.New("error building operator deployment: the deployment does not have the required replicas")
 	}
 	fmt.Println("Successfully installed the operator deployment")
@@ -265,8 +265,8 @@ func deployCertManager(config TestConfig) error {
 
 // hasDeploymentRequiredReplicas returns a condition function that indicates whether the given deployment
 // currently has the required amount of replicas in the ready state as specified in spec.replicas
-func hasDeploymentRequiredReplicas(ctx context.Context, dep *appsv1.Deployment) wait.ConditionFunc {
-	return func() (bool, error) {
+func hasDeploymentRequiredReplicas(dep *appsv1.Deployment) wait.ConditionWithContextFunc {
+	return func(ctx context.Context) (bool, error) {
 		err := e2eutil.TestClient.Get(ctx,
 			types.NamespacedName{Name: dep.Name,
 				Namespace: e2eutil.OperatorNamespace},

test/e2e/tlstests/tlstests.go

@@ -92,7 +92,7 @@ func rotateCertManagerSecret(ctx context.Context, secretName types.NamespacedNam
 	assert.NoError(t, err)
 
 	newSecret := corev1.Secret{}
-	err = wait.Poll(5*time.Second, 1*time.Minute, func() (done bool, err error) {
+	err = wait.PollUntilContextTimeout(ctx, 5*time.Second, 1*time.Minute, false, func(ctx context.Context) (done bool, err error) {
 		if err := e2eutil.TestClient.Get(ctx, secretName, &newSecret); err != nil {
 			return false, nil
 		}

test/e2e/util/mongotester/mongotester.go

@@ -255,7 +255,7 @@ func (m *Tester) connectivityCheck(shouldSucceed bool, opts ...OptionApplier) fu
 
 		attempts := 0
 		// There can be a short time before the user can auth as the user
-		err := wait.Poll(connectivityOpts.IntervalTime, connectivityOpts.TimeoutTime, func() (done bool, err error) {
+		err := wait.PollUntilContextTimeout(ctx, connectivityOpts.IntervalTime, connectivityOpts.TimeoutTime, false, func(ctx context.Context) (done bool, err error) {
 			attempts++
 			collection := m.mongoClient.Database(connectivityOpts.Database).Collection(connectivityOpts.Collection)
 			_, err = collection.InsertOne(ctx, bson.M{"name": "pi", "value": 3.14159})
@@ -299,7 +299,7 @@ func (m *Tester) WaitForRotatedCertificate(mdb mdbv1.MongoDBCommunity, initialCe
 		}
 
 		// Ping the cluster until it succeeds. The ping will only succeed with the right certificate.
-		err = wait.Poll(5*time.Second, 5*time.Minute, func() (done bool, err error) {
+		err = wait.PollUntilContextTimeout(m.ctx, 5*time.Second, 5*time.Minute, false, func(ctx context.Context) (done bool, err error) {
 			if err := m.mongoClient.Ping(m.ctx, nil); err != nil {
 				return false, nil
 			}
@@ -316,7 +316,7 @@ func (m *Tester) WaitForRotatedCertificate(mdb mdbv1.MongoDBCommunity, initialCe
 func (m *Tester) EnsureMongodConfig(selector string, expected interface{}) func(*testing.T) {
 	return func(t *testing.T) {
 		connectivityOpts := defaults()
-		err := wait.Poll(connectivityOpts.IntervalTime, connectivityOpts.TimeoutTime, func() (done bool, err error) {
+		err := wait.PollUntilContextTimeout(m.ctx, connectivityOpts.IntervalTime, connectivityOpts.TimeoutTime, false, func(ctx context.Context) (done bool, err error) {
 			opts, err := m.getCommandLineOptions()
 			assert.NoError(t, err)
 
@@ -408,7 +408,7 @@ func (m *Tester) PrometheusEndpointIsReachable(username, password string, useTls
 	client := &http.Client{Transport: customTransport}
 
 	return func(t *testing.T) {
-		_ = wait.Poll(5*time.Second, 60*time.Second, func() (bool, error) {
+		_ = wait.PollUntilContextTimeout(m.ctx, 5*time.Second, 60*time.Second, false, func(ctx context.Context) (bool, error) {
 			var idx int
 
 			// Verify that the Prometheus port is enabled and responding with 200

test/e2e/util/wait/wait.go

@@ -57,7 +57,7 @@ func ForMongoDBMessageStatus(ctx context.Context, t *testing.T, mdb *mdbv1.Mongo
 // waitForMongoDBCondition polls and waits for a given condition to be true
 func waitForMongoDBCondition(ctx context.Context, mdb *mdbv1.MongoDBCommunity, retryInterval, timeout time.Duration, condition func(mdbv1.MongoDBCommunity) bool) error {
 	mdbNew := mdbv1.MongoDBCommunity{}
-	return wait.Poll(retryInterval, timeout, func() (done bool, err error) {
+	return wait.PollUntilContextTimeout(ctx, retryInterval, timeout, false, func(ctx context.Context) (done bool, err error) {
 		err = e2eutil.TestClient.Get(ctx, mdb.NamespacedName(), &mdbNew)
 		if err != nil {
 			return false, err
@@ -144,7 +144,7 @@ func waitForStatefulSetConditionWithSpecificSts(ctx context.Context, t *testing.
 	if statefulSetType == ArbitersStatefulSet {
 		name = mdb.ArbiterNamespacedName()
 	}
-	return wait.Poll(waitOpts.RetryInterval, waitOpts.Timeout, func() (done bool, err error) {
+	return wait.PollUntilContextTimeout(ctx, waitOpts.RetryInterval, waitOpts.Timeout, false, func(ctx context.Context) (done bool, err error) {
 		err = e2eutil.TestClient.Get(ctx, name, &sts)
 		if err != nil {
 			return false, err
@@ -167,7 +167,7 @@ func waitForStatefulSetConditionArbiters(ctx context.Context, t *testing.T, mdb
 }
 
 func ForPodReadiness(ctx context.Context, t *testing.T, isReady bool, containerName string, timeout time.Duration, pod corev1.Pod) error {
-	return wait.Poll(time.Second*3, timeout, func() (done bool, err error) {
+	return wait.PollUntilContextTimeout(ctx, time.Second*3, timeout, false, func(ctx context.Context) (done bool, err error) {
 		err = e2eutil.TestClient.Get(ctx, types.NamespacedName{Name: pod.Name, Namespace: pod.Namespace}, &pod)
 		if err != nil {
 			return false, err
@@ -183,7 +183,7 @@ func ForPodReadiness(ctx context.Context, t *testing.T, isReady bool, containerN
 }
 
 func ForPodPhase(ctx context.Context, t *testing.T, timeout time.Duration, pod corev1.Pod, podPhase corev1.PodPhase) error {
-	return wait.Poll(time.Second*3, timeout, func() (done bool, err error) {
+	return wait.PollUntilContextTimeout(ctx, time.Second*3, timeout, false, func(ctx context.Context) (done bool, err error) {
 		err = e2eutil.TestClient.Get(ctx, types.NamespacedName{Name: pod.Name, Namespace: pod.Namespace}, &pod)
 		if err != nil {
 			return false, err
@@ -196,15 +196,15 @@ func ForPodPhase(ctx context.Context, t *testing.T, timeout time.Duration, pod c
 // waitForRuntimeObjectToExist waits until a runtime.Object of the given name exists
 // using the provided retryInterval and timeout provided.
 func waitForRuntimeObjectToExist(ctx context.Context, name string, retryInterval, timeout time.Duration, obj client.Object, namespace string) error {
-	return wait.Poll(retryInterval, timeout, func() (done bool, err error) {
+	return wait.PollUntilContextTimeout(ctx, retryInterval, timeout, false, func(ctx context.Context) (done bool, err error) {
 		return runtimeObjectExists(ctx, name, obj, namespace)
 	})
 }
 
 // waitForRuntimeObjectToBeDeleted waits until a runtime.Object of the given name is deleted
 // using the provided retryInterval and timeout provided.
 func waitForRuntimeObjectToBeDeleted(ctx context.Context, name string, retryInterval, timeout time.Duration, obj client.Object, namespace string) error {
-	return wait.Poll(retryInterval, timeout, func() (done bool, err error) {
+	return wait.PollUntilContextTimeout(ctx, retryInterval, timeout, false, func(ctx context.Context) (done bool, err error) {
 		exists, err := runtimeObjectExists(ctx, name, obj, namespace)
 		return !exists, err
 	})