diff --git a/e2e/cephfs_helper.go b/e2e/cephfs_helper.go
index e6c734f35..cf0d756ac 100644
--- a/e2e/cephfs_helper.go
+++ b/e2e/cephfs_helper.go
@@ -102,7 +102,7 @@ func createCephfsStorageClass(
 
 	timeout := time.Duration(deployTimeout) * time.Minute
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		_, err = c.StorageV1().StorageClasses().Create(context.TODO(), &sc, metav1.CreateOptions{})
 		if err != nil {
 			framework.Logf("error creating StorageClass %q: %v", sc.Name, err)
diff --git a/e2e/deployment.go b/e2e/deployment.go
index b481221c3..9ecb0e85b 100644
--- a/e2e/deployment.go
+++ b/e2e/deployment.go
@@ -93,7 +93,7 @@ func deleteDeploymentApp(clientSet kubernetes.Interface, name, ns string, deploy
 	start := time.Now()
 	framework.Logf("Waiting for deployment %q to be deleted", name)
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		_, err := clientSet.AppsV1().Deployments(ns).Get(context.TODO(), name, metav1.GetOptions{})
 		if err != nil {
 			if isRetryableAPIError(err) {
@@ -117,7 +117,7 @@ func waitForDeploymentInAvailableState(clientSet kubernetes.Interface, name, ns
 	start := time.Now()
 	framework.Logf("Waiting up to %q to be in Available state", name)
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		d, err := clientSet.AppsV1().Deployments(ns).Get(context.TODO(), name, metav1.GetOptions{})
 		if err != nil {
 			if isRetryableAPIError(err) {
@@ -144,7 +144,7 @@ func waitForDeploymentComplete(clientSet kubernetes.Interface, name, ns string,
 		err        error
 	)
 	timeout := time.Duration(deployTimeout) * time.Minute
-	err = wait.PollImmediate(poll, timeout, func() (bool, error) {
+	err = wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		deployment, err = clientSet.AppsV1().Deployments(ns).Get(context.TODO(), name, metav1.GetOptions{})
 		if err != nil {
 			if isRetryableAPIError(err) {
@@ -310,7 +310,7 @@ func waitForDeploymentUpdateScale(
 ) error {
 	t := time.Duration(timeout) * time.Minute
 	start := time.Now()
-	err := wait.PollImmediate(poll, t, func() (bool, error) {
+	err := wait.PollUntilContextTimeout(context.TODO(), poll, t, true, func(_ context.Context) (bool, error) {
 		scaleResult, upsErr := c.AppsV1().Deployments(ns).UpdateScale(context.TODO(),
 			deploymentName, scale, metav1.UpdateOptions{})
 		if upsErr != nil {
@@ -346,7 +346,7 @@ func waitForDeploymentUpdate(
 ) error {
 	t := time.Duration(timeout) * time.Minute
 	start := time.Now()
-	err := wait.PollImmediate(poll, t, func() (bool, error) {
+	err := wait.PollUntilContextTimeout(context.TODO(), poll, t, true, func(_ context.Context) (bool, error) {
 		_, upErr := c.AppsV1().Deployments(deployment.Namespace).Update(
 			context.TODO(), deployment, metav1.UpdateOptions{})
 		if upErr != nil {
@@ -456,7 +456,7 @@ func waitForContainersArgsUpdate(
 	// wait for scale to become count
 	t := time.Duration(timeout) * time.Minute
 	start := time.Now()
-	err = wait.PollImmediate(poll, t, func() (bool, error) {
+	err = wait.PollUntilContextTimeout(context.TODO(), poll, t, true, func(_ context.Context) (bool, error) {
 		deploy, getErr := c.AppsV1().Deployments(ns).Get(context.TODO(), deploymentName, metav1.GetOptions{})
 		if getErr != nil {
 			if isRetryableAPIError(getErr) {
diff --git a/e2e/namespace.go b/e2e/namespace.go
index 6ba2ff6a6..964418342 100644
--- a/e2e/namespace.go
+++ b/e2e/namespace.go
@@ -43,7 +43,7 @@ func createNamespace(c kubernetes.Interface, name string) error {
 		return fmt.Errorf("failed to create namespace: %w", err)
 	}
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		_, err := c.CoreV1().Namespaces().Get(context.TODO(), name, metav1.GetOptions{})
 		if err != nil {
 			framework.Logf("Error getting namespace: '%s': %v", name, err)
@@ -68,7 +68,7 @@ func deleteNamespace(c kubernetes.Interface, name string) error {
 		return fmt.Errorf("failed to delete namespace: %w", err)
 	}
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		_, err = c.CoreV1().Namespaces().Get(context.TODO(), name, metav1.GetOptions{})
 		if err != nil {
 			if apierrs.IsNotFound(err) {
diff --git a/e2e/nfs.go b/e2e/nfs.go
index 3632c61c2..044a4cf5f 100644
--- a/e2e/nfs.go
+++ b/e2e/nfs.go
@@ -183,7 +183,7 @@ func createNFSStorageClass(
 
 	timeout := time.Duration(deployTimeout) * time.Minute
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		_, err = c.StorageV1().StorageClasses().Create(context.TODO(), &sc, metav1.CreateOptions{})
 		if err != nil {
 			framework.Logf("error creating StorageClass %q: %v", sc.Name, err)
diff --git a/e2e/pod.go b/e2e/pod.go
index d02e05c89..5a8d379ca 100644
--- a/e2e/pod.go
+++ b/e2e/pod.go
@@ -60,7 +60,7 @@ func waitForDaemonSets(name, ns string, c kubernetes.Interface, t int) error {
 	start := time.Now()
 	framework.Logf("Waiting up to %v for all daemonsets in namespace '%s' to start", timeout, ns)
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		ds, err := c.AppsV1().DaemonSets(ns).Get(context.TODO(), name, metav1.GetOptions{})
 		if err != nil {
 			framework.Logf("Error getting daemonsets in namespace: '%s': %v", ns, err)
@@ -97,7 +97,7 @@ func findPodAndContainerName(f *framework.Framework, ns, cn string, opt *metav1.
 		podList *v1.PodList
 		listErr error
 	)
-	err := wait.PollImmediate(poll, timeout, func() (bool, error) {
+	err := wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		podList, listErr = e2epod.PodClientNS(f, ns).List(context.TODO(), *opt)
 		if listErr != nil {
 			if isRetryableAPIError(listErr) {
@@ -215,7 +215,7 @@ func listPods(f *framework.Framework, ns string, opt *metav1.ListOptions) ([]v1.
 func execWithRetry(f *framework.Framework, opts *e2epod.ExecOptions) (string, string, error) {
 	timeout := time.Duration(deployTimeout) * time.Minute
 	var stdOut, stdErr string
-	err := wait.PollImmediate(poll, timeout, func() (bool, error) {
+	err := wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		var execErr error
 		stdOut, stdErr, execErr = e2epod.ExecWithOptions(f, *opts)
 		if execErr != nil {
@@ -353,7 +353,7 @@ func waitForPodInRunningState(name, ns string, c kubernetes.Interface, t int, ex
 	start := time.Now()
 	framework.Logf("Waiting up to %v to be in Running state", name)
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		pod, err := c.CoreV1().Pods(ns).Get(context.TODO(), name, metav1.GetOptions{})
 		if err != nil {
 			if isRetryableAPIError(err) {
@@ -402,7 +402,7 @@ func deletePod(name, ns string, c kubernetes.Interface, t int) error {
 	start := time.Now()
 	framework.Logf("Waiting for pod %v to be deleted", name)
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		_, err := c.CoreV1().Pods(ns).Get(context.TODO(), name, metav1.GetOptions{})
 		if err != nil {
 			if isRetryableAPIError(err) {
diff --git a/e2e/pvc.go b/e2e/pvc.go
index 8ad63c7dc..318fa8b5e 100644
--- a/e2e/pvc.go
+++ b/e2e/pvc.go
@@ -58,7 +58,7 @@ func createPVCAndvalidatePV(c kubernetes.Interface, pvc *v1.PersistentVolumeClai
 	start := time.Now()
 	framework.Logf("Waiting up to %v to be in Bound state", pvc)
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		framework.Logf("waiting for PVC %s (%d seconds elapsed)", name, int(time.Since(start).Seconds()))
 		pvc, err = c.CoreV1().PersistentVolumeClaims(namespace).Get(context.TODO(), name, metav1.GetOptions{})
 		if err != nil {
@@ -130,7 +130,7 @@ func deletePVCAndPV(c kubernetes.Interface, pvc *v1.PersistentVolumeClaim, pv *v
 	start := time.Now()
 
 	pvcToDelete := pvc
-	err = wait.PollImmediate(poll, timeout, func() (bool, error) {
+	err = wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		// Check that the PVC is deleted.
 		framework.Logf(
 			"waiting for PVC %s in state %s to be deleted (%d seconds elapsed)",
@@ -168,7 +168,7 @@ func deletePVCAndPV(c kubernetes.Interface, pvc *v1.PersistentVolumeClaim, pv *v
 	start = time.Now()
 	pvToDelete := pv
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		// Check that the PV is deleted.
 		framework.Logf(
 			"waiting for PV %s in state %s to be deleted (%d seconds elapsed)",
@@ -197,19 +197,24 @@ func getPersistentVolumeClaim(c kubernetes.Interface, namespace, name string) (*
 	var pvc *v1.PersistentVolumeClaim
 	var err error
 	timeout := time.Duration(deployTimeout) * time.Minute
-	err = wait.PollImmediate(1*time.Second, timeout, func() (bool, error) {
-		pvc, err = c.CoreV1().PersistentVolumeClaims(namespace).Get(context.TODO(), name, metav1.GetOptions{})
-		if err != nil {
-			framework.Logf("Error getting pvc %q in namespace %q: %v", name, namespace, err)
-			if isRetryableAPIError(err) {
-				return false, nil
+	err = wait.PollUntilContextTimeout(
+		context.TODO(),
+		1*time.Second,
+		timeout,
+		true,
+		func(_ context.Context) (bool, error) {
+			pvc, err = c.CoreV1().PersistentVolumeClaims(namespace).Get(context.TODO(), name, metav1.GetOptions{})
+			if err != nil {
+				framework.Logf("Error getting pvc %q in namespace %q: %v", name, namespace, err)
+				if isRetryableAPIError(err) {
+					return false, nil
+				}
+
+				return false, fmt.Errorf("failed to get pvc: %w", err)
 			}
 
-			return false, fmt.Errorf("failed to get pvc: %w", err)
-		}
-
-		return true, err
-	})
+			return true, err
+		})
 
 	return pvc, err
 }
@@ -220,19 +225,24 @@ func getPersistentVolume(c kubernetes.Interface, name string) (*v1.PersistentVol
 	var pv *v1.PersistentVolume
 	var err error
 	timeout := time.Duration(deployTimeout) * time.Minute
-	err = wait.PollImmediate(1*time.Second, timeout, func() (bool, error) {
-		pv, err = c.CoreV1().PersistentVolumes().Get(context.TODO(), name, metav1.GetOptions{})
-		if err != nil {
-			framework.Logf("Error getting pv %q: %v", name, err)
-			if isRetryableAPIError(err) {
-				return false, nil
+	err = wait.PollUntilContextTimeout(
+		context.TODO(),
+		1*time.Second,
+		timeout,
+		true,
+		func(_ context.Context) (bool, error) {
+			pv, err = c.CoreV1().PersistentVolumes().Get(context.TODO(), name, metav1.GetOptions{})
+			if err != nil {
+				framework.Logf("Error getting pv %q: %v", name, err)
+				if isRetryableAPIError(err) {
+					return false, nil
+				}
+
+				return false, fmt.Errorf("failed to get pv: %w", err)
 			}
 
-			return false, fmt.Errorf("failed to get pv: %w", err)
-		}
-
-		return true, err
-	})
+			return true, err
+		})
 
 	return pv, err
 }
@@ -275,7 +285,7 @@ func deletePVCAndValidatePV(c kubernetes.Interface, pvc *v1.PersistentVolumeClai
 	}
 	start := time.Now()
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		// Check that the PVC is really deleted.
 		framework.Logf(
 			"waiting for PVC %s in state %s to be deleted (%d seconds elapsed)",
@@ -384,7 +394,7 @@ func getMetricsForPVC(f *framework.Framework, pvc *v1.PersistentVolumeClaim, t i
 	// retry as kubelet does not immediately have the metrics available
 	timeout := time.Duration(t) * time.Minute
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		stdOut, stdErr, err := execCommandInToolBoxPod(f, cmd, rookNamespace)
 		if err != nil {
 			framework.Logf("failed to get metrics for pvc %q (%v): %v", pvc.Name, err, stdErr)
diff --git a/e2e/rbd.go b/e2e/rbd.go
index eb14a545d..37b592942 100644
--- a/e2e/rbd.go
+++ b/e2e/rbd.go
@@ -212,7 +212,7 @@ func checkClusternameInMetadata(f *framework.Framework, ns, pool, image string)
 		stdErr  string
 		execErr error
 	)
-	err := wait.PollImmediate(poll, t, func() (bool, error) {
+	err := wait.PollUntilContextTimeout(context.TODO(), poll, t, true, func(_ context.Context) (bool, error) {
 		coName, stdErr, execErr = execCommandInToolBoxPod(f,
 			fmt.Sprintf("rbd image-meta get %s --image=%s %s", rbdOptions(pool), image, clusterNameKey),
 			ns)
@@ -1887,7 +1887,7 @@ var _ = Describe("RBD", func() {
 
 				timeout := time.Duration(deployTimeout) * time.Minute
 				var reason string
-				err = wait.PollImmediate(poll, timeout, func() (bool, error) {
+				err = wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 					var runningAttachCmd string
 					runningAttachCmd, stdErr, err = execCommandInContainer(
 						f,
diff --git a/e2e/rbd_helper.go b/e2e/rbd_helper.go
index c248734dc..3c58068d1 100644
--- a/e2e/rbd_helper.go
+++ b/e2e/rbd_helper.go
@@ -164,7 +164,7 @@ func createRBDStorageClass(
 
 	timeout := time.Duration(deployTimeout) * time.Minute
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		_, err = c.StorageV1().StorageClasses().Create(context.TODO(), &sc, metav1.CreateOptions{})
 		if err != nil {
 			framework.Logf("error creating StorageClass %q: %v", sc.Name, err)
@@ -1037,7 +1037,7 @@ func listRBDImagesInTrash(f *framework.Framework, poolName string) ([]trashInfo,
 func waitToRemoveImagesFromTrash(f *framework.Framework, poolName string, t int) error {
 	var errReason error
 	timeout := time.Duration(t) * time.Minute
-	err := wait.PollImmediate(poll, timeout, func() (bool, error) {
+	err := wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		imagesInTrash, err := listRBDImagesInTrash(f, poolName)
 		if err != nil {
 			return false, err
diff --git a/e2e/resize.go b/e2e/resize.go
index 8f3b30219..e03148497 100644
--- a/e2e/resize.go
+++ b/e2e/resize.go
@@ -50,7 +50,7 @@ func expandPVCSize(c kubernetes.Interface, pvc *v1.PersistentVolumeClaim, size s
 	start := time.Now()
 	framework.Logf("Waiting up to %v to be in Resized state", pvc)
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		framework.Logf("waiting for PVC %s (%d seconds elapsed)", pvcName, int(time.Since(start).Seconds()))
 		updatedPVC, err = c.CoreV1().
 			PersistentVolumeClaims(pvcNamespace).
@@ -185,7 +185,7 @@ func checkAppMntSize(f *framework.Framework, opt *metav1.ListOptions, size, cmd,
 	timeout := time.Duration(t) * time.Minute
 	start := time.Now()
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		framework.Logf("executing cmd %s (%d seconds elapsed)", cmd, int(time.Since(start).Seconds()))
 		output, stdErr, err := execCommandInPod(f, cmd, ns, opt)
 		if err != nil {
diff --git a/e2e/snapshot.go b/e2e/snapshot.go
index 0839513cc..45ae7964a 100644
--- a/e2e/snapshot.go
+++ b/e2e/snapshot.go
@@ -80,7 +80,7 @@ func createSnapshot(snap *snapapi.VolumeSnapshot, t int) error {
 	start := time.Now()
 	framework.Logf("waiting for %v to be in ready state", snap)
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		framework.Logf("waiting for snapshot %s (%d seconds elapsed)", snap.Name, int(time.Since(start).Seconds()))
 		snaps, err := sclient.
 			VolumeSnapshots(snap.Namespace).
@@ -126,7 +126,7 @@ func deleteSnapshot(snap *snapapi.VolumeSnapshot, t int) error {
 	start := time.Now()
 	framework.Logf("Waiting up to %v to be deleted", snap)
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		framework.Logf("deleting snapshot %s (%d seconds elapsed)", name, int(time.Since(start).Seconds()))
 		_, err := sclient.
 			VolumeSnapshots(snap.Namespace).
@@ -223,7 +223,7 @@ func createNFSSnapshotClass(f *framework.Framework) error {
 
 	timeout := time.Duration(deployTimeout) * time.Minute
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		_, err = sclient.VolumeSnapshotClasses().Create(context.TODO(), &sc, metav1.CreateOptions{})
 		if err != nil {
 			framework.Logf("error creating SnapshotClass %q: %v", sc.Name, err)
@@ -252,7 +252,7 @@ func deleteNFSSnapshotClass() error {
 
 	timeout := time.Duration(deployTimeout) * time.Minute
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		err = sclient.VolumeSnapshotClasses().Delete(context.TODO(), sc.Name, metav1.DeleteOptions{})
 		if err != nil {
 			framework.Logf("error deleting SnapshotClass %q: %v", sc.Name, err)
diff --git a/e2e/utils.go b/e2e/utils.go
index de85a0142..bc9cc4f3a 100644
--- a/e2e/utils.go
+++ b/e2e/utils.go
@@ -248,7 +248,7 @@ func getMons(ns string, c kubernetes.Interface) ([]string, error) {
 
 	var svcList *v1.ServiceList
 	t := time.Duration(deployTimeout) * time.Minute
-	err := wait.PollImmediate(poll, t, func() (bool, error) {
+	err := wait.PollUntilContextTimeout(context.TODO(), poll, t, true, func(_ context.Context) (bool, error) {
 		var svcErr error
 		svcList, svcErr = c.CoreV1().Services(ns).List(context.TODO(), opt)
 		if svcErr != nil {
@@ -1560,7 +1560,7 @@ func waitForJobCompletion(c kubernetes.Interface, ns, job string, timeout int) e
 
 	framework.Logf("waiting for Job %s/%s to be in state %q", ns, job, batch.JobComplete)
 
-	return wait.PollImmediate(poll, t, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, t, true, func(_ context.Context) (bool, error) {
 		j, err := c.BatchV1().Jobs(ns).Get(context.TODO(), job, metav1.GetOptions{})
 		if err != nil {
 			if isRetryableAPIError(err) {
@@ -1608,7 +1608,7 @@ func retryKubectlInput(namespace string, action kubectlAction, data string, t in
 	framework.Logf("waiting for kubectl (%s -f args %s) to finish", action, args)
 	start := time.Now()
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		cmd := []string{}
 		if len(args) != 0 {
 			cmd = append(cmd, strings.Join(args, ""))
@@ -1647,7 +1647,7 @@ func retryKubectlFile(namespace string, action kubectlAction, filename string, t
 	framework.Logf("waiting for kubectl (%s -f %q args %s) to finish", action, filename, args)
 	start := time.Now()
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		cmd := []string{}
 		if len(args) != 0 {
 			cmd = append(cmd, strings.Join(args, ""))
@@ -1690,7 +1690,7 @@ func retryKubectlArgs(namespace string, action kubectlAction, t int, args ...str
 	framework.Logf("waiting for kubectl (%s args) to finish", args)
 	start := time.Now()
 
-	return wait.PollImmediate(poll, timeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(context.TODO(), poll, timeout, true, func(_ context.Context) (bool, error) {
 		_, err := e2ekubectl.RunKubectl(namespace, args...)
 		if err != nil {
 			if isRetryableAPIError(err) {