tests(logging): Minimized risk of flaky metrics/sinks tests by adding retries

logging/go.mod

@@ -3,7 +3,7 @@ module cloud.google.com/go/logging
 go 1.21
 
 require (
-	cloud.google.com/go v0.116.0
+	cloud.google.com/go v0.117.0
 	cloud.google.com/go/compute/metadata v0.5.2
 	cloud.google.com/go/iam v1.2.2
 	cloud.google.com/go/longrunning v0.6.2

logging/go.sum

@@ -1,6 +1,6 @@
 cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
-cloud.google.com/go v0.116.0 h1:B3fRrSDkLRt5qSHWe40ERJvhvnQwdZiHu0bJOpldweE=
-cloud.google.com/go v0.116.0/go.mod h1:cEPSRWPzZEswwdr9BxE6ChEn01dWlTaF05LiC2Xs70U=
+cloud.google.com/go v0.117.0 h1:Z5TNFfQxj7WG2FgOGX1ekC5RiXrYgms6QscOm32M/4s=
+cloud.google.com/go v0.117.0/go.mod h1:ZbwhVTb1DBGt2Iwb3tNO6SEK4q+cplHZmLWH+DelYYc=
 cloud.google.com/go/auth v0.12.1 h1:n2Bj25BUMM0nvE9D2XLTiImanwZhO3DkfWSYS/SAJP4=
 cloud.google.com/go/auth v0.12.1/go.mod h1:BFMu+TNpF3DmvfBO9ClqTR/SiqVIm7LukKF9mbendF4=
 cloud.google.com/go/auth/oauth2adapt v0.2.6 h1:V6a6XDu2lTwPZWOawrAa9HUK+DB2zfJyTuciBG5hFkU=

logging/internal/testing/retry.go

@@ -0,0 +1,110 @@
+/*
+Copyright 2024 Google LLC
+
+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 testing
+
+import (
+	"testing"
+	"time"
+
+	"cloud.google.com/go/internal/testutil"
+	"google.golang.org/api/iterator"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+var defaultMaxAttempts = 10
+var defaultSleep = 10 * time.Second
+var defaultRetryableCodes = map[codes.Code]bool{
+	codes.Unavailable: true,
+}
+
+// Iterator is a wrapper interface type for iterators in the logadmin
+// library that have a Next function that gets the next item/error, or returns
+// nil/iterator.Done if the object has no next item.
+type Iterator[T any] interface {
+	Next() (*T, error)
+}
+
+// handleError handles the given error for the retry attempt.
+func handleError(r *testutil.R, err error) {
+	if err != nil {
+		s, ok := status.FromError(err)
+
+		// Throw a fatal error if the error is not retryable or if it cannot be converted into
+		// a status object.
+		if ok && !defaultRetryableCodes[s.Code()] {
+			r.Fatalf("%+v\n", err)
+		} else if ok {
+			r.Errorf("%+v\n", err)
+		} else {
+			r.Fatalf("%+v\n", err)
+		}
+	}
+}
+
+// Retry is a wrapper around testutil.Retry that retries the test function on Unavailable errors, otherwise, Fatalfs.
+func Retry(t *testing.T, f func(r *testutil.R) error) bool {
+	retryFunc := func(r *testutil.R) {
+		err := f(r)
+		handleError(r, err)
+	}
+	return testutil.Retry(t, defaultMaxAttempts, defaultSleep, retryFunc)
+}
+
+// RetryAndExpectError retries the test function on Unavailable errors, otherwise passes
+// if a different error was thrown. If no non-retryable error is returned, fails.
+func RetryAndExpectError(t *testing.T, f func(r *testutil.R) error) bool {
+	retryFunc := func(r *testutil.R) {
+		err := f(r)
+
+		if err != nil {
+			s, ok := status.FromError(err)
+
+			// Only retry on retryable errors, otherwise pass.
+			if ok && defaultRetryableCodes[s.Code()] {
+				r.Errorf("%+v\n", err)
+			}
+		} else {
+			r.Fatalf("got no error, expected one")
+		}
+	}
+
+	return testutil.Retry(t, defaultMaxAttempts, defaultSleep, retryFunc)
+}
+
+// RetryIteratorNext is a wrapper around testutil.Retry that retries the given iterator's Next function
+// and returns the next object, retrying if a retryable error is found. If a non-retryable error is found, fail
+// the test.
+func RetryIteratorNext[T any](t *testing.T, it Iterator[T]) (*T, bool) {
+	var next *T
+	var err error
+	retryFunc := func(r *testutil.R) {
+		next, err = it.Next()
+		if err != nil {
+			if err == iterator.Done {
+				return
+			}
+
+			handleError(r, err)
+		}
+	}
+	testutil.Retry(t, defaultMaxAttempts, defaultSleep, retryFunc)
+	if err == iterator.Done {
+		return nil, true
+	}
+	return next, false
+}

logging/logadmin/metrics_test.go

@@ -22,6 +22,7 @@ import (
 
 	"cloud.google.com/go/internal/testutil"
 	"cloud.google.com/go/internal/uid"
+	ltest "cloud.google.com/go/logging/internal/testing"
 	"google.golang.org/api/iterator"
 )
 
@@ -55,26 +56,32 @@ func TestCreateDeleteMetric(t *testing.T) {
 		Description: "DESC",
 		Filter:      "FILTER",
 	}
+
 	if err := client.CreateMetric(ctx, metric); err != nil {
 		t.Fatal(err)
 	}
 	defer client.DeleteMetric(ctx, metric.ID)
 
-	got, err := client.Metric(ctx, metric.ID)
-	if err != nil {
-		t.Fatal(err)
-	}
+	var got *Metric
+	ltest.Retry(t, func(r *testutil.R) error {
+		var err error
+		got, err = client.Metric(ctx, metric.ID)
+		return err
+	})
 	if want := metric; !testutil.Equal(got, want) {
 		t.Errorf("got %+v, want %+v", got, want)
 	}
 
-	if err := client.DeleteMetric(ctx, metric.ID); err != nil {
-		t.Fatal(err)
-	}
+	ltest.Retry(t, func(r *testutil.R) error {
+		return client.DeleteMetric(ctx, metric.ID)
+	})
 
-	if _, err := client.Metric(ctx, metric.ID); err == nil {
-		t.Fatal("got no error, expected one")
-	}
+	// client.Metric should give an error. Test if this is the case, but retry on
+	// retryable errors.
+	ltest.RetryAndExpectError(t, func(r *testutil.R) error {
+		_, err := client.Metric(ctx, metric.ID)
+		return err
+	})
 }
 
 func TestUpdateMetric(t *testing.T) {
@@ -86,27 +93,31 @@ func TestUpdateMetric(t *testing.T) {
 	}
 
 	// Updating a non-existent metric creates a new one.
-	if err := client.UpdateMetric(ctx, metric); err != nil {
-		t.Fatal(err)
-	}
+	ltest.Retry(t, func(r *testutil.R) error {
+		return client.UpdateMetric(ctx, metric)
+	})
 	defer client.DeleteMetric(ctx, metric.ID)
-	got, err := client.Metric(ctx, metric.ID)
-	if err != nil {
-		t.Fatal(err)
-	}
+
+	var got *Metric
+	ltest.Retry(t, func(r *testutil.R) error {
+		var err error
+		got, err = client.Metric(ctx, metric.ID)
+		return err
+	})
 	if want := metric; !testutil.Equal(got, want) {
 		t.Errorf("got %+v, want %+v", got, want)
 	}
 
 	// Updating an existing metric changes it.
 	metric.Description = "CHANGED"
-	if err := client.UpdateMetric(ctx, metric); err != nil {
-		t.Fatal(err)
-	}
-	got, err = client.Metric(ctx, metric.ID)
-	if err != nil {
-		t.Fatal(err)
-	}
+	ltest.Retry(t, func(r *testutil.R) error {
+		return client.UpdateMetric(ctx, metric)
+	})
+	ltest.Retry(t, func(r *testutil.R) error {
+		var err error
+		got, err = client.Metric(ctx, metric.ID)
+		return err
+	})
 	if want := metric; !testutil.Equal(got, want) {
 		t.Errorf("got %+v, want %+v", got, want)
 	}
@@ -127,22 +138,15 @@ func TestListMetrics(t *testing.T) {
 		want[m.ID] = m
 	}
 	for _, m := range metrics {
-		if err := client.CreateMetric(ctx, m); err != nil {
-			t.Fatalf("Create(%q): %v", m.ID, err)
-		}
+		ltest.Retry(t, func(r *testutil.R) error {
+			return client.CreateMetric(ctx, m)
+		})
 		defer client.DeleteMetric(ctx, m.ID)
 	}
 
 	got := map[string]*Metric{}
 	it := client.Metrics(ctx)
-	for {
-		m, err := it.Next()
-		if err == iterator.Done {
-			break
-		}
-		if err != nil {
-			t.Fatal(err)
-		}
+	for m, done := ltest.RetryIteratorNext(t, it); !done; m, done = ltest.RetryIteratorNext(t, it) {
 		// If tests run simultaneously, we may have more metrics than we
 		// created. So only check for our own.
 		if _, ok := want[m.ID]; ok {