e2e framework: support polling with Get

These helper functions can be used in combination with
omega.Eventually/Consistently to implement polling of objects that is aware of
Kubernetes apiserver conventions:

- retry on certain errors instead of giving up,
  with "not found" handling decided by the caller (may or may not
  be fatal, depending on the test)
- sleep if requested by apiserver

test/e2e/framework/get.go

@@ -0,0 +1,131 @@
+/*
+Copyright 2023 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package framework
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/onsi/gomega"
+
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+// GetFunc is a function which retrieves a certain object.
+type GetFunc[T any] func(ctx context.Context) (T, error)
+
+// APIGetFunc is a get functions as used in client-go.
+type APIGetFunc[T any] func(ctx context.Context, name string, getOptions metav1.GetOptions) (T, error)
+
+// GetObject takes a get function like clientset.CoreV1().Pods(ns).Get
+// and the parameters for it and returns a function that executes that get
+// operation in a [gomega.Eventually] or [gomega.Consistently].
+//
+// Delays and retries are handled by [HandleRetry]. A "not found" error is
+// a fatal error that causes polling to stop immediately. If that is not
+// desired, then wrap the result with [IgnoreNotFound].
+func GetObject[T any](get APIGetFunc[T], name string, getOptions metav1.GetOptions) GetFunc[T] {
+	return HandleRetry(func(ctx context.Context) (T, error) {
+		return get(ctx, name, getOptions)
+	})
+}
+
+// HandleRetry wraps an arbitrary get function. When the wrapped function
+// returns an error, HandleGetError will decide whether the call should be
+// retried and if requested, will sleep before doing so.
+//
+// This is meant to be used inside [gomega.Eventually] or [gomega.Consistently].
+func HandleRetry[T any](get GetFunc[T]) GetFunc[T] {
+	return func(ctx context.Context) (T, error) {
+		t, err := get(ctx)
+		if err != nil {
+			if retry, delay := ShouldRetry(err); retry {
+				if delay > 0 {
+					// We could return
+					// gomega.TryAgainAfter(delay) here,
+					// but then we need to funnel that
+					// error through any other
+					// wrappers. Waiting directly is simpler.
+					ctx, cancel := context.WithTimeout(ctx, delay)
+					defer cancel()
+					<-ctx.Done()
+				}
+				return t, err
+			}
+			// Give up polling immediately.
+			var null T
+			return t, gomega.StopTrying(fmt.Sprintf("Unexpected final error while getting %T", null)).Wrap(err)
+		}
+		return t, nil
+	}
+}
+
+// ShouldRetry decides whether to retry an API request. Optionally returns a
+// delay to retry after.
+func ShouldRetry(err error) (retry bool, retryAfter time.Duration) {
+	// if the error sends the Retry-After header, we respect it as an explicit confirmation we should retry.
+	if delay, shouldRetry := apierrors.SuggestsClientDelay(err); shouldRetry {
+		return shouldRetry, time.Duration(delay) * time.Second
+	}
+
+	// these errors indicate a transient error that should be retried.
+	if apierrors.IsTimeout(err) || apierrors.IsTooManyRequests(err) || errors.As(err, &transientError{}) {
+		return true, 0
+	}
+
+	return false, 0
+}
+
+// RetryNotFound wraps an arbitrary get function. When the wrapped function
+// encounters a "not found" error, that error is treated as a transient problem
+// and polling continues.
+//
+// This is meant to be used inside [gomega.Eventually] or [gomega.Consistently].
+func RetryNotFound[T any](get GetFunc[T]) GetFunc[T] {
+	return func(ctx context.Context) (T, error) {
+		t, err := get(ctx)
+		if apierrors.IsNotFound(err) {
+			// If we are wrapping HandleRetry, then the error will
+			// be gomega.StopTrying. We need to get rid of that,
+			// otherwise gomega.Eventually will stop.
+			var stopTryingErr gomega.PollingSignalError
+			if errors.As(err, &stopTryingErr) {
+				if wrappedErr := errors.Unwrap(stopTryingErr); wrappedErr != nil {
+					err = wrappedErr
+				}
+			}
+
+			// Mark the error as transient in case that we get
+			// wrapped by HandleRetry.
+			err = transientError{error: err}
+		}
+		return t, err
+	}
+}
+
+// transientError wraps some other error and indicates that the
+// wrapper error is something that may go away.
+type transientError struct {
+	error
+}
+
+func (err transientError) Unwrap() error {
+	return err.error
+}

test/e2e/framework/pod/wait.go

@@ -817,7 +817,7 @@ func handleWaitingAPIError(err error, retryNotFound bool, taskFormat string, tas
 		framework.Logf("Ignoring NotFound error while " + taskDescription)
 		return false, nil
 	}
-	if retry, delay := shouldRetry(err); retry {
+	if retry, delay := framework.ShouldRetry(err); retry {
 		framework.Logf("Retryable error while %s, retrying after %v: %v", taskDescription, delay, err)
 		if delay > 0 {
 			time.Sleep(delay)
@@ -827,18 +827,3 @@ func handleWaitingAPIError(err error, retryNotFound bool, taskFormat string, tas
 	framework.Logf("Encountered non-retryable error while %s: %v", taskDescription, err)
 	return false, err
 }
-
-// Decide whether to retry an API request. Optionally include a delay to retry after.
-func shouldRetry(err error) (retry bool, retryAfter time.Duration) {
-	// if the error sends the Retry-After header, we respect it as an explicit confirmation we should retry.
-	if delay, shouldRetry := apierrors.SuggestsClientDelay(err); shouldRetry {
-		return shouldRetry, time.Duration(delay) * time.Second
-	}
-
-	// these errors indicate a transient error that should be retried.
-	if apierrors.IsTimeout(err) || apierrors.IsTooManyRequests(err) {
-		return true, 0
-	}
-
-	return false, 0
-}