diff --git a/pkg/prober/prober.go b/pkg/prober/prober.go
index ec9330606..42d5b36eb 100644
--- a/pkg/prober/prober.go
+++ b/pkg/prober/prober.go
@@ -186,7 +186,7 @@ func (m *Manager) doAsync(ctx context.Context, target string, arg interface{}, p
 			result bool
 			inErr  error
 		)
-		err := wait.PollImmediate(period, timeout, func() (bool, error) {
+		err := wait.PollUntilContextTimeout(ctx, period, timeout, true, func(ctx context.Context) (bool, error) {
 			result, inErr = Do(ctx, m.transport, target, ops...)
 			// Do not return error, which is from verifierError, as retry is expected until timeout.
 			return result, nil
diff --git a/pkg/prober/prober_test.go b/pkg/prober/prober_test.go
index dc9c79239..3a49f2934 100644
--- a/pkg/prober/prober_test.go
+++ b/pkg/prober/prober_test.go
@@ -224,7 +224,7 @@ func TestDoAsyncTimeout(t *testing.T) {
 		if done {
 			t.Errorf("done was true")
 		}
-		if !errors.Is(err, wait.ErrWaitTimeout) {
+		if !errors.Is(err, context.DeadlineExceeded) {
 			t.Error("Unexpected error =", err)
 		}
 		wch <- arg
@@ -258,7 +258,7 @@ func TestAsyncMultiple(t *testing.T) {
 	wch <- 2009
 	<-wch
 
-	wait.PollImmediate(probeInterval, probeTimeout, func() (bool, error) {
+	wait.PollUntilContextTimeout(context.Background(), probeInterval, probeTimeout, true, func(ctx context.Context) (bool, error) {
 		return m.len() == 0, nil
 	})
 	if got, want := m.len(), 0; got != want {
diff --git a/test/conformance/certificate/utils.go b/test/conformance/certificate/utils.go
index 362a21037..47bd26644 100644
--- a/test/conformance/certificate/utils.go
+++ b/test/conformance/certificate/utils.go
@@ -82,7 +82,7 @@ func WaitForCertificateSecret(ctx context.Context, t *testing.T, client *test.Cl
 	span := logging.GetEmitableSpan(context.Background(), fmt.Sprintf("WaitForCertificateSecret/%s/%s", cert.Spec.SecretName, desc))
 	defer span.End()
 
-	return wait.PollImmediate(test.PollInterval, test.PollTimeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(ctx, test.PollInterval, test.PollTimeout, true, func(ctx context.Context) (bool, error) {
 		secret, err := client.KubeClient.CoreV1().Secrets(test.ServingNamespace).Get(ctx, cert.Spec.SecretName, metav1.GetOptions{})
 		if apierrs.IsNotFound(err) {
 			return false, nil
@@ -119,7 +119,7 @@ func WaitForCertificateState(ctx context.Context, client *test.NetworkingClients
 	defer span.End()
 
 	var lastState *v1alpha1.Certificate
-	return wait.PollImmediate(test.PollInterval, test.PollTimeout, func() (bool, error) {
+	return wait.PollUntilContextTimeout(ctx, test.PollInterval, test.PollTimeout, true, func(ctx context.Context) (bool, error) {
 		var err error
 		lastState, err = client.Certificates.Get(ctx, name, metav1.GetOptions{})
 		if err != nil {
diff --git a/test/conformance/ingress/util.go b/test/conformance/ingress/util.go
index 3eb41dadc..e1268ce24 100644
--- a/test/conformance/ingress/util.go
+++ b/test/conformance/ingress/util.go
@@ -749,7 +749,7 @@ func createPodAndService(ctx context.Context, t *testing.T, clients *test.Client
 	}
 
 	// Wait for the Pod to show up in the Endpoints resource.
-	waitErr := wait.PollImmediate(test.PollInterval, test.PollTimeout, func() (bool, error) {
+	waitErr := wait.PollUntilContextTimeout(ctx, test.PollInterval, test.PollTimeout, true, func(ctx context.Context) (bool, error) {
 		var ep *corev1.Endpoints
 		err := reconciler.RetryTestErrors(func(attempts int) (err error) {
 			ep, err = clients.KubeClient.CoreV1().Endpoints(svc.Namespace).Get(ctx, svc.Name, metav1.GetOptions{})
@@ -1204,7 +1204,7 @@ func WaitForIngressState(ctx context.Context, client *test.NetworkingClients, na
 	defer span.End()
 
 	var lastState *v1alpha1.Ingress
-	waitErr := wait.PollImmediate(test.PollInterval, test.PollTimeout, func() (bool, error) {
+	waitErr := wait.PollUntilContextTimeout(ctx, test.PollInterval, test.PollTimeout, true, func(ctx context.Context) (bool, error) {
 		err := reconciler.RetryTestErrors(func(attempts int) (err error) {
 			lastState, err = client.Ingresses.Get(ctx, name, metav1.GetOptions{})
 			return err
diff --git a/test/ingress.go b/test/ingress.go
index a8cfd0d14..ef6a3faaa 100644
--- a/test/ingress.go
+++ b/test/ingress.go
@@ -35,7 +35,7 @@ func WaitForIngressState(ctx context.Context, client *NetworkingClients, name st
 	defer span.End()
 
 	var lastState *v1alpha1.Ingress
-	waitErr := wait.PollImmediate(PollInterval, PollTimeout, func() (bool, error) {
+	waitErr := wait.PollUntilContextTimeout(ctx, PollInterval, PollTimeout, true, func(ctx context.Context) (bool, error) {
 		var err error
 		lastState, err = client.Ingresses.Get(ctx, name, metav1.GetOptions{})
 		if err != nil {