Align Await timeout with test context

common/testing/await/config.go

@@ -5,6 +5,7 @@ import (
 	"time"
 
 	"go.temporal.io/server/common/debug"
+	"go.temporal.io/server/common/testing/testcontext"
 )
 
 const attemptTimeoutEnvVar = "TEMPORAL_AWAIT_ATTEMPT_TIMEOUT"
@@ -17,13 +18,13 @@ type config struct {
 
 func newConfig() config {
 	return config{
+		totalTimeout:   testcontext.DefaultTimeout(),
 		attemptTimeout: envDuration(attemptTimeoutEnvVar, 10*time.Second) * debug.TimeoutMultiplier,
 	}
 }
 
-func legacyConfig(timeout time.Duration, timeoutMsg string) config {
+func legacyConfig(timeoutMsg string) config {
 	cfg := newConfig()
-	cfg.totalTimeout = timeout
 	cfg.timeoutMsg = timeoutMsg
 	return cfg
 }

common/testing/await/config_test.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/stretchr/testify/require"
 	"go.temporal.io/server/common/debug"
+	"go.temporal.io/server/common/testing/testcontext"
 )
 
 func TestConfig_OverrideAttemptTimeout(t *testing.T) {
@@ -15,6 +16,13 @@ func TestConfig_OverrideAttemptTimeout(t *testing.T) {
 	require.Equal(t, 250*time.Millisecond*debug.TimeoutMultiplier, cfg.attemptTimeout)
 }
 
+func TestConfig_UsesTestContextTimeout(t *testing.T) {
+	t.Setenv("TEMPORAL_TEST_TIMEOUT", "250ms")
+
+	cfg := newConfig()
+	require.Equal(t, testcontext.DefaultTimeout(), cfg.totalTimeout)
+}
+
 func TestNextPollIntervalCapsAtMaximum(t *testing.T) {
 	require.Equal(t, 500*time.Millisecond, nextPollInterval(1))
 	require.Equal(t, time.Second, nextPollInterval(2))

common/testing/await/doc.go

@@ -5,8 +5,9 @@
 // their formatted variants. By default, they enforce a 10s timeout for each
 // await attempt.
 //
-// Polling backs off from 500ms to 2s. The poll interval arguments remain for
-// source compatibility and are ignored.
+// The total wait is bounded by the test context timeout, and polling backs off
+// from 500ms to 2s. The timeout and poll interval arguments remain for source
+// compatibility and are ignored.
 //
 // Improvements over testify's eventually functions:
 //

common/testing/await/require_ctx.go

@@ -43,21 +43,23 @@ func hardDeadlockTimeout() time.Duration {
 const postAwaitTimeoutReserve = 10 * time.Second
 
 // Require polls condition until it returns without assertion failures, or
-// until ctx is canceled or timeout expires (whichever is earliest).
+// until ctx is canceled or the test context timeout expires (whichever is
+// earliest).
 //
 // Pass the *await.T to require.*/assert.* — failures cause a retry, not a
 // test failure. Use t.Context() inside the callback to honor the timeout.
-// The poll interval argument is retained for source compatibility and ignored.
-func Require(ctx context.Context, tb testing.TB, condition func(*T), timeout, _ time.Duration) {
+// The timeout and poll interval arguments are retained for source compatibility
+// and ignored.
+func Require(ctx context.Context, tb testing.TB, condition func(*T), _, _ time.Duration) {
 	tb.Helper()
-	run(ctx, tb, condition, legacyConfig(timeout, ""), "Require", requireMisuseHint, true)
+	run(ctx, tb, condition, legacyConfig(""), "Require", requireMisuseHint, true)
 }
 
 // Requiref is like [Require] but adds a formatted message to the timeout
-// failure. Its poll interval argument is also ignored.
-func Requiref(ctx context.Context, tb testing.TB, condition func(*T), timeout, _ time.Duration, msg string, args ...any) {
+// failure. Its timeout and poll interval arguments are also ignored.
+func Requiref(ctx context.Context, tb testing.TB, condition func(*T), _, _ time.Duration, msg string, args ...any) {
 	tb.Helper()
-	run(ctx, tb, condition, legacyConfig(timeout, fmt.Sprintf(msg, args...)), "Requiref", requireMisuseHint, true)
+	run(ctx, tb, condition, legacyConfig(fmt.Sprintf(msg, args...)), "Requiref", requireMisuseHint, true)
 }
 
 func run(

common/testing/await/require_ctx_test.go

@@ -99,28 +99,28 @@ func TestRequire_PropagatesParentContextValues(t *testing.T) {
 	require.Equal(t, "value", got)
 }
 
-func TestRequire_SetsTimeoutContextDeadline(t *testing.T) {
+func TestRequire_IgnoresLegacyTimeoutArgument(t *testing.T) {
 	t.Parallel()
 
 	longCtx := testcontext.For(t)
 	longDeadline, ok := longCtx.Deadline()
 	require.True(t, ok)
 
-	shortTimeout := 1 * time.Second
+	attemptTimeout := 10 * time.Second * debug.TimeoutMultiplier
 
-	var shortCtx context.Context
+	var attemptCtx context.Context
 	await.Require(longCtx, t, func(t *await.T) {
-		shortCtx = t.Context()
-	}, shortTimeout, 100*time.Millisecond)
+		attemptCtx = t.Context()
+	}, time.Nanosecond, time.Hour)
 
-	require.NotNil(t, shortCtx)
-	require.NotSame(t, longCtx, shortCtx)
+	require.NotNil(t, attemptCtx)
+	require.NotSame(t, longCtx, attemptCtx)
 
-	shortDeadline, ok := shortCtx.Deadline()
+	attemptDeadline, ok := attemptCtx.Deadline()
 	require.True(t, ok)
-	require.True(t, shortDeadline.Before(longDeadline))
-	require.LessOrEqual(t, time.Until(shortDeadline), shortTimeout)
-	require.Greater(t, time.Until(shortDeadline), shortTimeout-200*time.Millisecond)
+	require.True(t, attemptDeadline.Before(longDeadline))
+	require.LessOrEqual(t, time.Until(attemptDeadline), attemptTimeout)
+	require.Greater(t, time.Until(attemptDeadline), attemptTimeout-200*time.Millisecond)
 }
 
 func TestRequire_ExtendsCachedTestContextPastActiveExpiration(t *testing.T) {
@@ -128,13 +128,9 @@ func TestRequire_ExtendsCachedTestContextPastActiveExpiration(t *testing.T) {
 
 	synctest.Test(t, func(t *testing.T) {
 		ctx := testcontext.For(t)
-		completeAt := time.Now().Add(testcontext.DefaultTimeout() + time.Second)
 
-		await.Require(ctx, t, func(t *await.T) {
-			if time.Now().Before(completeAt) {
-				t.Error("not ready")
-			}
-		}, testcontext.DefaultTimeout()+5*time.Second, testcontext.DefaultTimeout()+time.Second)
+		await.Require(ctx, t, func(*await.T) {}, testcontext.DefaultTimeout()+5*time.Second, testcontext.DefaultTimeout()+time.Second)
+		time.Sleep(testcontext.DefaultTimeout() + time.Second) //nolint:forbidigo // advance past the original active expiration
 
 		require.Same(t, ctx, testcontext.For(t))
 		require.NoError(t, ctx.Err())
@@ -173,7 +169,8 @@ func TestRequire_FailureScenarios(t *testing.T) {
 	t.Run("reports timeout", func(t *testing.T) {
 		t.Parallel()
 
-		ctx := testcontext.For(t)
+		ctx, cancel := context.WithTimeout(testcontext.For(t), time.Second)
+		defer cancel()
 		tb := newRecordingTB()
 		tb.run(func() {
 			await.Require(ctx, tb, func(t *await.T) {
@@ -187,7 +184,8 @@ func TestRequire_FailureScenarios(t *testing.T) {
 	t.Run("cancels attempt context on timeout", func(t *testing.T) {
 		t.Parallel()
 
-		ctx := testcontext.For(t)
+		ctx, cancel := context.WithTimeout(testcontext.For(t), 2*time.Second)
+		defer cancel()
 		tb := newRecordingTB()
 		tb.run(func() {
 			await.Require(ctx, tb, func(t *await.T) {
@@ -207,7 +205,9 @@ func TestRequire_FailureScenarios(t *testing.T) {
 		pollInterval := 100 * time.Millisecond
 		t.Setenv("TEMPORAL_AWAIT_ATTEMPT_TIMEOUT", attemptTimeoutEnv.String())
 
-		ctx := testcontext.For(t)
+		awaitTimeout := 2*attemptTimeout + time.Second
+		ctx, cancel := context.WithTimeout(testcontext.For(t), awaitTimeout)
+		defer cancel()
 		var attempts atomic.Int32
 		var firstAttemptRemaining time.Duration
 
@@ -219,7 +219,7 @@ func TestRequire_FailureScenarios(t *testing.T) {
 					firstAttemptRemaining = time.Until(deadline)
 				}
 				<-t.Context().Done()
-			}, 2*attemptTimeout+time.Second, pollInterval)
+			}, awaitTimeout, pollInterval)
 		})
 
 		require.True(t, tb.Failed())
@@ -232,7 +232,8 @@ func TestRequire_FailureScenarios(t *testing.T) {
 	t.Run("does not poll again after attempt consumes timeout", func(t *testing.T) {
 		t.Parallel()
 
-		ctx := testcontext.For(t)
+		ctx, cancel := context.WithTimeout(testcontext.For(t), time.Second)
+		defer cancel()
 		var attempts atomic.Int32
 
 		tb := newRecordingTB()
@@ -297,7 +298,8 @@ func TestRequire_FailureScenarios(t *testing.T) {
 	t.Run("reports all attempt errors on timeout", func(t *testing.T) {
 		t.Parallel()
 
-		ctx := testcontext.For(t)
+		ctx, cancel := context.WithTimeout(testcontext.For(t), time.Second)
+		defer cancel()
 		var attempts atomic.Int32
 		tb := newRecordingTB()
 		tb.run(func() {
@@ -318,7 +320,8 @@ func TestRequire_FailureScenarios(t *testing.T) {
 
 	t.Run("truncates middle attempts when many fail", func(t *testing.T) {
 		synctest.Test(t, func(t *testing.T) {
-			ctx := testcontext.For(t)
+			ctx, cancel := context.WithTimeout(testcontext.For(t), 6*time.Second)
+			defer cancel()
 			var attempts atomic.Int32
 			tb := newRecordingTB()
 			tb.run(func() {
@@ -346,7 +349,8 @@ func TestRequire_FailureScenarios(t *testing.T) {
 	t.Run("Requiref includes message on timeout", func(t *testing.T) {
 		t.Parallel()
 
-		ctx := testcontext.For(t)
+		ctx, cancel := context.WithTimeout(testcontext.For(t), time.Second)
+		defer cancel()
 		tb := newRecordingTB()
 		tb.run(func() {
 			await.Requiref(ctx, tb, func(t *await.T) {
@@ -465,7 +469,8 @@ func TestRequire_WaitsForInFlightAttemptOnTimeout(t *testing.T) {
 	t.Parallel()
 
 	var finished atomic.Bool
-	ctx := testcontext.For(t)
+	ctx, cancel := context.WithTimeout(testcontext.For(t), time.Second)
+	defer cancel()
 	tb := newRecordingTB()
 	tb.run(func() {
 		await.Require(ctx, tb, func(t *await.T) {
@@ -481,6 +486,7 @@ func TestRequire_WaitsForInFlightAttemptOnTimeout(t *testing.T) {
 // recordingTB is a minimal testing.TB implementation for testing failure scenarios.
 type recordingTB struct {
 	testing.TB    // embed for interface satisfaction
+	ctx           context.Context
 	mu            sync.Mutex
 	failed        atomic.Bool
 	errorMessages []string
@@ -490,7 +496,14 @@ type recordingTB struct {
 }
 
 func newRecordingTB() *recordingTB {
-	return &recordingTB{}
+	return &recordingTB{ctx: context.Background()}
+}
+
+func newRecordingTBWithTimeout(timeout time.Duration) *recordingTB {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	tb := &recordingTB{ctx: ctx}
+	tb.Cleanup(cancel)
+	return tb
 }
 
 func (r *recordingTB) Helper()      {}
@@ -502,7 +515,7 @@ func (r *recordingTB) Logf(format string, args ...any) {
 	r.logMessages = append(r.logMessages, fmt.Sprintf(format, args...))
 }
 func (r *recordingTB) Context() context.Context {
-	return context.Background()
+	return r.ctx
 }
 
 func (r *recordingTB) Cleanup(fn func()) {

common/testing/await/require_true.go

@@ -11,28 +11,29 @@ import (
 const requireTrueMisuseHint = "do not use test assertions inside the predicate - return false to retry or use await.Require for assertions"
 
 // RequireTrue runs `condition` repeatedly until it returns true, or until the
-// timeout expires. The timeout is capped at the test's deadline, if one is set.
-// The poll interval argument is retained for source compatibility and ignored.
+// test context timeout expires. The timeout is capped at the test's deadline,
+// if one is set. The timeout and poll interval arguments are retained for source
+// compatibility and ignored.
 //
 // Use [RequireTrue] for simple local predicates only. Do not use assertions or
 // side effects in the predicate - use [Require] for these.
-func RequireTrue(tb testing.TB, condition func() bool, timeout, _ time.Duration) {
+func RequireTrue(tb testing.TB, condition func() bool, _, _ time.Duration) {
 	tb.Helper()
 	run(testcontext.For(tb), tb, func(t *T) {
 		if !condition() {
 			t.Fail()
 		}
-	}, legacyConfig(timeout, ""), "RequireTrue", requireTrueMisuseHint, false)
+	}, legacyConfig(""), "RequireTrue", requireTrueMisuseHint, false)
 }
 
 // RequireTruef is like [RequireTrue] but accepts a format string that is included
 // in the failure message when the condition is not satisfied before the timeout.
-// Its poll interval argument is also ignored.
-func RequireTruef(tb testing.TB, condition func() bool, timeout, _ time.Duration, msg string, args ...any) {
+// Its timeout and poll interval arguments are also ignored.
+func RequireTruef(tb testing.TB, condition func() bool, _, _ time.Duration, msg string, args ...any) {
 	tb.Helper()
 	run(testcontext.For(tb), tb, func(t *T) {
 		if !condition() {
 			t.Fail()
 		}
-	}, legacyConfig(timeout, fmt.Sprintf(msg, args...)), "RequireTruef", requireTrueMisuseHint, false)
+	}, legacyConfig(fmt.Sprintf(msg, args...)), "RequireTruef", requireTrueMisuseHint, false)
 }

common/testing/await/require_true_test.go

@@ -58,7 +58,7 @@ func TestRequireTrue_FailureScenarios(t *testing.T) {
 	t.Run("reports timeout", func(t *testing.T) {
 		t.Parallel()
 
-		tb := newRecordingTB()
+		tb := newRecordingTBWithTimeout(time.Second)
 		tb.run(func() {
 			await.RequireTrue(tb, func() bool {
 				return false
@@ -71,7 +71,7 @@ func TestRequireTrue_FailureScenarios(t *testing.T) {
 	t.Run("RequireTruef includes message on timeout", func(t *testing.T) {
 		t.Parallel()
 
-		tb := newRecordingTB()
+		tb := newRecordingTBWithTimeout(time.Second)
 		tb.run(func() {
 			await.RequireTruef(tb, func() bool {
 				return false

common/testing/parallelsuite/suite_test.go

@@ -94,9 +94,8 @@ func (s *contextSuite) TestAwaitUsesSuiteContext() {
 
 	s.Await(func(s *contextSuite) {
 		s.Equal("decorated", s.Context().Value(key{}))
-		deadline, ok := s.Context().Deadline()
+		_, ok := s.Context().Deadline()
 		s.True(ok)
-		s.Less(time.Until(deadline), 200*time.Millisecond)
 	}, 100*time.Millisecond, time.Millisecond)
 }