Export EnforceDefaultTimeoutsWhenUsingContexts and DisableDefaultTime…

…outsWhenUsingContext fix: #790 Signed-off-by: Toshikuni Fukaya <[email protected]>
onsi · Oct 31, 2024 · ca36da1 · ca36da1
1 parent d6331f9
commit ca36da1
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 13 deletions.
diff --git a/docs/index.md b/docs/index.md
@@ -604,7 +604,7 @@ SetDefaultConsistentlyPollingInterval(t time.Duration)
 
 You can also adjust these global timeouts by setting the `GOMEGA_DEFAULT_EVENTUALLY_TIMEOUT`, `GOMEGA_DEFAULT_EVENTUALLY_POLLING_INTERVAL`, `GOMEGA_DEFAULT_CONSISTENTLY_DURATION`, and `GOMEGA_DEFAULT_CONSISTENTLY_POLLING_INTERVAL` environment variables to a parseable duration string. The environment variables have a lower precedence than `SetDefault...()`.
 
-As discussed [above](#category-2-making-eventually-assertions-on-functions) `Eventually`s that are passed a `context` object without an explicit timeout will only stop polling when the context is cancelled.  If you would like to enforce the default timeout when a context is provided you can call `EnforceDefaultTimeoutsWhenUsingContexts()` (to go back to the default behavior call `DoNotEnforceDefaultTimeoutsWhenUsingContexts()`).   You can also set the `GOMEGA_ENFORCE_DEFAULT_TIMEOUTS_WHEN_USING_CONTEXTS` environment variable to enforce the default timeout when a context is provided.
+As discussed [above](#category-2-making-eventually-assertions-on-functions) `Eventually`s that are passed a `context` object without an explicit timeout will only stop polling when the context is cancelled.  If you would like to enforce the default timeout when a context is provided you can call `EnforceDefaultTimeoutsWhenUsingContexts()` (to go back to the default behavior call `DisableDefaultTimeoutsWhenUsingContexts()`).   You can also set the `GOMEGA_ENFORCE_DEFAULT_TIMEOUTS_WHEN_USING_CONTEXTS` environment variable to enforce the default timeout when a context is provided.
 
 ## Making Assertions in Helper Functions
 

diff --git a/gomega_dsl.go b/gomega_dsl.go
@@ -503,6 +503,16 @@ func SetDefaultConsistentlyPollingInterval(t time.Duration) {
 	Default.SetDefaultConsistentlyPollingInterval(t)
 }
 
+// EnforceDefaultTimeoutsWhenUsingContexts forces `Eventually` to apply a default timeout even when a context is provided.
+func EnforceDefaultTimeoutsWhenUsingContexts() {
+	Default.EnforceDefaultTimeoutsWhenUsingContexts()
+}
+
+// DisableDefaultTimeoutsWhenUsingContext disables the default timeout when a context is provided to `Eventually`.
+func DisableDefaultTimeoutsWhenUsingContext() {
+	Default.DisableDefaultTimeoutsWhenUsingContext()
+}
+
 // AsyncAssertion is returned by Eventually and Consistently and polls the actual value passed into Eventually against
 // the matcher passed to the Should and ShouldNot methods.
 //

diff --git a/internal/dsl_test.go b/internal/dsl_test.go
@@ -19,6 +19,11 @@ func setGlobalDurationBundle(bundle internal.DurationBundle) {
 	SetDefaultEventuallyPollingInterval(bundle.EventuallyPollingInterval)
 	SetDefaultConsistentlyDuration(bundle.ConsistentlyDuration)
 	SetDefaultConsistentlyPollingInterval(bundle.ConsistentlyPollingInterval)
+	if bundle.EnforceDefaultTimeoutsWhenUsingContexts {
+		EnforceDefaultTimeoutsWhenUsingContexts()
+	} else {
+		DisableDefaultTimeoutsWhenUsingContext()
+	}
 }
 
 var _ = Describe("Gomega DSL", func() {
@@ -51,10 +56,11 @@ var _ = Describe("Gomega DSL", func() {
 	Describe("NewGomega", func() {
 		It("creates and configures a new Gomega, using the global duration bundle", func() {
 			bundle := internal.DurationBundle{
-				EventuallyTimeout:           time.Minute,
-				EventuallyPollingInterval:   2 * time.Minute,
-				ConsistentlyDuration:        3 * time.Minute,
-				ConsistentlyPollingInterval: 4 * time.Minute,
+				EventuallyTimeout:                       time.Minute,
+				EventuallyPollingInterval:               2 * time.Minute,
+				ConsistentlyDuration:                    3 * time.Minute,
+				ConsistentlyPollingInterval:             4 * time.Minute,
+				EnforceDefaultTimeoutsWhenUsingContexts: true,
 			}
 			setGlobalDurationBundle(bundle)
 
@@ -74,10 +80,11 @@ var _ = Describe("Gomega DSL", func() {
 	Describe("NewWithT", func() {
 		It("creates and configure a new Gomega with the passed-in T, using the global duration bundle", func() {
 			bundle := internal.DurationBundle{
-				EventuallyTimeout:           time.Minute,
-				EventuallyPollingInterval:   2 * time.Minute,
-				ConsistentlyDuration:        3 * time.Minute,
-				ConsistentlyPollingInterval: 4 * time.Minute,
+				EventuallyTimeout:                       time.Minute,
+				EventuallyPollingInterval:               2 * time.Minute,
+				ConsistentlyDuration:                    3 * time.Minute,
+				ConsistentlyPollingInterval:             4 * time.Minute,
+				EnforceDefaultTimeoutsWhenUsingContexts: true,
 			}
 			setGlobalDurationBundle(bundle)
 
@@ -191,16 +198,18 @@ var _ = Describe("Gomega DSL", func() {
 	Describe("specifying default durations globally", func() {
 		It("should update the durations on the Default gomega", func() {
 			bundle := internal.DurationBundle{
-				EventuallyTimeout:           time.Minute,
-				EventuallyPollingInterval:   2 * time.Minute,
-				ConsistentlyDuration:        3 * time.Minute,
-				ConsistentlyPollingInterval: 4 * time.Minute,
+				EventuallyTimeout:                       time.Minute,
+				EventuallyPollingInterval:               2 * time.Minute,
+				ConsistentlyDuration:                    3 * time.Minute,
+				ConsistentlyPollingInterval:             4 * time.Minute,
+				EnforceDefaultTimeoutsWhenUsingContexts: true,
 			}
 
 			SetDefaultEventuallyTimeout(bundle.EventuallyTimeout)
 			SetDefaultEventuallyPollingInterval(bundle.EventuallyPollingInterval)
 			SetDefaultConsistentlyDuration(bundle.ConsistentlyDuration)
 			SetDefaultConsistentlyPollingInterval(bundle.ConsistentlyPollingInterval)
+			EnforceDefaultTimeoutsWhenUsingContexts()
 
 			Ω(Default.(*internal.Gomega).DurationBundle).Should(Equal(bundle))
 		})

diff --git a/types/types.go b/types/types.go
@@ -29,6 +29,8 @@ type Gomega interface {
 	SetDefaultEventuallyPollingInterval(time.Duration)
 	SetDefaultConsistentlyDuration(time.Duration)
 	SetDefaultConsistentlyPollingInterval(time.Duration)
+	EnforceDefaultTimeoutsWhenUsingContexts()
+	DisableDefaultTimeoutsWhenUsingContext()
 }
 
 // All Gomega matchers must implement the GomegaMatcher interface