Allow StopTrying() to be wrapped

onsi · Oct 11, 2022 · bf3cba9 · bf3cba9
1 parent c5d5c39
commit bf3cba9
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 12 deletions.
diff --git a/internal/async_assertion.go b/internal/async_assertion.go
@@ -7,6 +7,7 @@ import (
 	"runtime"
 	"sync"
 	"time"
+	"errors"
 
 	"github.com/onsi/gomega/types"
 )
@@ -17,6 +18,23 @@ type StopTryingError interface {
 	wasViaPanic() bool
 }
 
+
+func asStopTryingError(actual interface{}) (StopTryingError, bool) {
+	if actual == nil {
+		return nil, false
+	}
+	if actualErr, ok := actual.(error); ok {
+		var target *stopTryingError
+		if errors.As(actualErr, &target) {
+			return target, true
+		} else {
+			return nil, false
+		}
+	}
+
+	return nil, false
+}
+
 type stopTryingError struct {
 	message  string
 	viaPanic bool
@@ -35,8 +53,6 @@ func (s *stopTryingError) wasViaPanic() bool {
 	return s.viaPanic
 }
 
-var stopTryingErrorType = reflect.TypeOf(&stopTryingError{})
-
 var StopTrying = func(message string) StopTryingError {
 	return &stopTryingError{message: message}
 }
@@ -157,20 +173,19 @@ func (assertion *AsyncAssertion) processReturnValues(values []reflect.Value) (in
 		return nil, fmt.Errorf("No values were returned by the function passed to Gomega"), stopTrying
 	}
 	actual := values[0].Interface()
-	if actual != nil && reflect.TypeOf(actual) == stopTryingErrorType {
-		stopTrying = actual.(StopTryingError)
+	if stopTryingErr, ok := asStopTryingError(actual); ok{
+		stopTrying = stopTryingErr
 	}
 	for i, extraValue := range values[1:] {
 		extra := extraValue.Interface()
 		if extra == nil {
 			continue
 		}
-		extraType := reflect.TypeOf(extra)
-		if extraType == stopTryingErrorType {
-			stopTrying = extra.(StopTryingError)
+		if stopTryingErr, ok := asStopTryingError(extra); ok{
+			stopTrying = stopTryingErr
 			continue
 		}
-		zero := reflect.Zero(extraType).Interface()
+		zero := reflect.Zero(reflect.TypeOf(extra)).Interface()
 		if reflect.DeepEqual(extra, zero) {
 			continue
 		}
@@ -226,8 +241,8 @@ func (assertion *AsyncAssertion) buildActualPoller() (func() (interface{}, error
 		return func() (actual interface{}, err error, stopTrying StopTryingError) {
 			defer func() {
 				if e := recover(); e != nil {
-					if reflect.TypeOf(e) == stopTryingErrorType {
-						stopTrying = e.(StopTryingError)
+					if stopTryingErr, ok := asStopTryingError(e); ok {
+						stopTrying = stopTryingErr
 					} else {
 						panic(e)
 					}
@@ -291,8 +306,8 @@ func (assertion *AsyncAssertion) buildActualPoller() (func() (interface{}, error
 				}
 			}
 			if e := recover(); e != nil {
-				if reflect.TypeOf(e) == stopTryingErrorType {
-					stopTrying = e.(StopTryingError)
+				if stopTryingErr, ok := asStopTryingError(e); ok {
+					stopTrying = stopTryingErr
 				} else if assertionFailure == nil {
 					panic(e)
 				}

diff --git a/internal/async_assertion_test.go b/internal/async_assertion_test.go
@@ -987,6 +987,23 @@ var _ = Describe("Asynchronous Assertions", func() {
 					Ω(ig.FailureMessage).Should(ContainSubstring("Reached the end - after"))
 					Ω(ig.FailureMessage).Should(ContainSubstring("Expected\n    <string>: C\nto equal\n    <string>: D"))
 				})
+
+				It("works even if the error is wrapped", func() {
+					possibilities := []string{"A", "B", "C"}
+					i := 0
+					ig.G.Eventually(func() (string, error) {
+						possibility := possibilities[i]
+						i += 1
+						if i == len(possibilities) {
+							return possibility, fmt.Errorf("Wrapped error: %w", StopTrying("Reached the end"))
+						} else {
+							return possibility, nil
+						}
+					}).Should(Equal("D"))
+					Ω(i).Should(Equal(3))
+					Ω(ig.FailureMessage).Should(ContainSubstring("Reached the end - after"))
+					Ω(ig.FailureMessage).Should(ContainSubstring("Expected\n    <string>: C\nto equal\n    <string>: D"))
+				})
 			})
 
 			Context("when StopTrying().Now() is called", func() {