diff --git a/test/integration/framework/test_server.go b/test/integration/framework/test_server.go
index eff7502d422..3841578c615 100644
--- a/test/integration/framework/test_server.go
+++ b/test/integration/framework/test_server.go
@@ -222,7 +222,7 @@ func StartTestServer(ctx context.Context, t testing.TB, setup TestServerSetup) (
 	kubeAPIServerClientConfig.ServerName = ""
 
 	// wait for health
-	err = wait.PollImmediate(100*time.Millisecond, 10*time.Second, func() (done bool, err error) {
+	err = wait.PollUntilContextTimeout(ctx, 100*time.Millisecond, 10*time.Second, true, func(ctx context.Context) (done bool, err error) {
 		select {
 		case err := <-errCh:
 			return false, err