From e37715c65962505daa49551f9ec08c70531676a0 Mon Sep 17 00:00:00 2001
From: Adam Ross <adamross@google.com>
Date: Wed, 26 Jul 2023 10:16:22 -0700
Subject: [PATCH] feat: add utils.PollE

---
 infra/blueprint-test/pkg/utils/poll.go      | 20 ++++-
 infra/blueprint-test/pkg/utils/poll_test.go | 92 +++++++++++++++++++++
 2 files changed, 108 insertions(+), 4 deletions(-)
 create mode 100644 infra/blueprint-test/pkg/utils/poll_test.go

diff --git a/infra/blueprint-test/pkg/utils/poll.go b/infra/blueprint-test/pkg/utils/poll.go
index c69a69018b4..3b2703737be 100644
--- a/infra/blueprint-test/pkg/utils/poll.go
+++ b/infra/blueprint-test/pkg/utils/poll.go
@@ -17,6 +17,7 @@
 package utils
 
 import (
+	"fmt"
 	"time"
 
 	"github.com/mitchellh/go-testing-interface"
@@ -25,12 +26,21 @@ import (
 // Polls on a particular condition function while the returns true.
 // It fails the test if the condition is not met within numRetries.
 func Poll(t testing.TB, condition func() (bool, error), numRetries int, interval time.Duration) {
+	err := PollE(t, condition, numRetries, interval)
+	if err != nil {
+		t.Fatalf("failed to pull provided condition after %d retries, last error: %v", numRetries, err)
+	}
+}
+
+// Polls on a particular condition function while the returns true.
+// Returns an error if the condition is not met within numRetries.
+func PollE(t testing.TB, condition func() (bool, error), numRetries int, interval time.Duration) error {
 	if numRetries < 0 {
-		t.Fatal("invalid value for numRetries. Must be >= 0")
+		return fmt.Errorf("invalid value for numRetries. Must be >= 0")
 	}
 
 	if interval <= 0 {
-		t.Fatal("invalid value for numRetries. Must be > 0")
+		return fmt.Errorf("invalid value for numRetries. Must be > 0")
 	}
 
 	retry, err := condition()
@@ -45,10 +55,12 @@ func Poll(t testing.TB, condition func() (bool, error), numRetries int, interval
 	}
 
 	if err != nil {
-		t.Fatalf("failed to pull provided condition after %d retries, last error: %v", numRetries, err)
+		return fmt.Errorf("failed to pull provided condition after %d retries, last error: %v", numRetries, err)
 	}
 
 	if retry {
-		t.Fatalf("polling timed out after %d retries with %d second intervals", numRetries, interval/time.Second)
+		return fmt.Errorf("polling timed out after %d retries with %d second intervals", numRetries, interval/time.Second)
 	}
+
+	return nil
 }
diff --git a/infra/blueprint-test/pkg/utils/poll_test.go b/infra/blueprint-test/pkg/utils/poll_test.go
new file mode 100644
index 00000000000..c76dbc2dd01
--- /dev/null
+++ b/infra/blueprint-test/pkg/utils/poll_test.go
@@ -0,0 +1,92 @@
+/**
+ * Copyright 2023 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 utils_test
+
+import (
+	"errors"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/GoogleCloudPlatform/cloud-foundation-toolkit/infra/blueprint-test/pkg/utils"
+)
+
+func TestPoll(t *testing.T) {
+	testcases := []struct {
+		label     string
+		condition func() (bool, error)
+		want      string
+	}{
+		{
+			label: "error",
+			condition: func() (bool, error) {
+				return true, errors.New("condition failure")
+			},
+			want: "condition failure",
+		},
+		{
+			label: "timeout",
+			condition: func() (bool, error) {
+				return true, nil
+			},
+			want: "polling timed out",
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.label, func(t *testing.T) {
+			it := &inspectableT{t, nil}
+			utils.Poll(it, tc.condition, 3, time.Millisecond)
+			if !strings.Contains(it.err.Error(), tc.want) {
+				t.Errorf("got %v, want %v", it.err, tc.want)
+			}
+		})
+	}
+}
+
+func TestPollE(t *testing.T) {
+	testcases := []struct {
+		label     string
+		condition func() (bool, error)
+		want      string
+	}{
+		{
+			label: "error",
+			condition: func() (bool, error) {
+				return true, errors.New("condition failure")
+			},
+			want: "condition failure",
+		},
+		{
+			label: "timeout",
+			condition: func() (bool, error) {
+				return true, nil
+			},
+			want: "polling timed out",
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.label, func(t *testing.T) {
+			it := &inspectableT{t, nil}
+			err := utils.PollE(it, tc.condition, 3, time.Millisecond)
+			if !strings.Contains(err.Error(), tc.want) {
+				t.Errorf("got %v, want %v", it.err, tc.want)
+			}
+		})
+	}
+}