diff --git a/staging/src/k8s.io/apiserver/pkg/util/flowcontrol/apf_filter_test.go b/staging/src/k8s.io/apiserver/pkg/util/flowcontrol/apf_filter_test.go
index 72e9c73909a..910b1339385 100644
--- a/staging/src/k8s.io/apiserver/pkg/util/flowcontrol/apf_filter_test.go
+++ b/staging/src/k8s.io/apiserver/pkg/util/flowcontrol/apf_filter_test.go
@@ -129,7 +129,7 @@ func TestQueueWaitTimeLatencyTracker(t *testing.T) {
 	}()
 
 	// ensure that the controller has run its first loop.
-	err := wait.PollImmediate(100*time.Millisecond, 5*time.Second, func() (done bool, err error) {
+	err := wait.PollUntilContextTimeout(context.Background(), 100*time.Millisecond, 5*time.Second, true, func(ctx context.Context) (bool, error) {
 		return controller.hasPriorityLevelState(plcObj.Name), nil
 	})
 	if err != nil {
diff --git a/staging/src/k8s.io/apiserver/pkg/util/flowcontrol/borrowing_test.go b/staging/src/k8s.io/apiserver/pkg/util/flowcontrol/borrowing_test.go
index 3a51da0047b..a3db94fa39c 100644
--- a/staging/src/k8s.io/apiserver/pkg/util/flowcontrol/borrowing_test.go
+++ b/staging/src/k8s.io/apiserver/pkg/util/flowcontrol/borrowing_test.go
@@ -165,7 +165,7 @@ func TestBorrowing(t *testing.T) {
 			}()
 
 			// ensure that the controller has run its first loop.
-			err := wait.PollImmediate(100*time.Millisecond, 5*time.Second, func() (done bool, err error) {
+			err := wait.PollUntilContextTimeout(ctx, 100*time.Millisecond, 5*time.Second, true, func(ctx context.Context) (bool, error) {
 				return controller.hasPriorityLevelState(plcObjs[0].Name), nil
 			})
 			if err != nil {
diff --git a/staging/src/k8s.io/apiserver/pkg/util/flowcontrol/controller_test.go b/staging/src/k8s.io/apiserver/pkg/util/flowcontrol/controller_test.go
index c706a7c3c38..2a804628884 100644
--- a/staging/src/k8s.io/apiserver/pkg/util/flowcontrol/controller_test.go
+++ b/staging/src/k8s.io/apiserver/pkg/util/flowcontrol/controller_test.go
@@ -407,7 +407,7 @@ func TestAPFControllerWithGracefulShutdown(t *testing.T) {
 	}()
 
 	// ensure that the controller has run its first loop.
-	err := wait.PollImmediate(100*time.Millisecond, 5*time.Second, func() (done bool, err error) {
+	err := wait.PollUntilContextTimeout(ctx, 100*time.Millisecond, 5*time.Second, true, func(ctx context.Context) (bool, error) {
 		return controller.hasPriorityLevelState(plName), nil
 	})
 	if err != nil {