diff --git a/test/e2e/storage/csi_mock/csi_storage_capacity.go b/test/e2e/storage/csi_mock/csi_storage_capacity.go
index 48fda168fcd..b2cb03668b1 100644
--- a/test/e2e/storage/csi_mock/csi_storage_capacity.go
+++ b/test/e2e/storage/csi_mock/csi_storage_capacity.go
@@ -177,7 +177,7 @@ var _ = utils.SIGDescribe("CSI Mock volume storage capacity", func() {
 				}
 
 				var calls []drivers.MockCSICall
-				err = wait.PollImmediateUntilWithContext(ctx, time.Second, func(ctx context.Context) (done bool, err error) {
+				err = wait.PollUntilContextCancel(ctx, time.Second, true, func(ctx context.Context) (done bool, err error) {
 					c, index, err := compareCSICalls(ctx, deterministicCalls, expected, m.driver.GetCalls)
 					if err != nil {
 						return true, fmt.Errorf("error waiting for expected CSI calls: %w", err)
diff --git a/test/integration/apiserver/podlogs/podlogs_test.go b/test/integration/apiserver/podlogs/podlogs_test.go
index b1d0edc4c98..949f26eed7f 100644
--- a/test/integration/apiserver/podlogs/podlogs_test.go
+++ b/test/integration/apiserver/podlogs/podlogs_test.go
@@ -283,7 +283,7 @@ func TestPodLogsKubeletClientCertReload(t *testing.T) {
 	}
 
 	// wait until the kubelet observes the new CA
-	if err := wait.PollImmediateUntilWithContext(ctx, time.Second, func(ctx context.Context) (bool, error) {
+	if err := wait.PollUntilContextCancel(ctx, time.Second, true, func(ctx context.Context) (bool, error) {
 		_, errLog := clientSet.CoreV1().Pods("ns").GetLogs(pod.Name, &corev1.PodLogOptions{}).DoRaw(ctx)
 		if errors.IsUnauthorized(errLog) {
 			return true, nil
@@ -302,7 +302,7 @@ func TestPodLogsKubeletClientCertReload(t *testing.T) {
 	}
 
 	// confirm that the API server observes the new client cert and closes existing connections to use it
-	if err := wait.PollImmediateUntilWithContext(ctx, time.Second, func(ctx context.Context) (bool, error) {
+	if err := wait.PollUntilContextCancel(ctx, time.Second, true, func(ctx context.Context) (bool, error) {
 		fixedPodLogs, errLog := clientSet.CoreV1().Pods("ns").GetLogs(pod.Name, &corev1.PodLogOptions{}).DoRaw(ctx)
 		if errors.IsUnauthorized(errLog) {
 			t.Log("api server has not observed new client cert")