fix: catch cause variants in context traps

README.md

@@ -328,6 +328,11 @@ if !errors.Is(ctx.Err(), context.DeadlineExceeded) {
 }
 ```
 
+`mClock.Trap().WithTimeout()` traps `WithTimeout` and `WithTimeoutCause` calls, and
+`mClock.Trap().WithDeadline()` traps `WithDeadline` and `WithDeadlineCause` calls, so a test keeps
+working if the code switches variants. The trapped call's `Cause` field holds the cause passed to a
+cause variant.
+
 A canceled parent cancels the child. A parent created with `context.WithDeadline` or
 `context.WithTimeout` can still cancel the child on wall time because Quartz does not convert
 parent deadlines to mock time.

mock.go

@@ -571,22 +571,16 @@ func (t Trapper) Until(tags ...string) *Trap {
 	return t.mock.newTrap(clockFunctionUntil, tags)
 }
 
+// WithDeadline traps calls to WithDeadline and WithDeadlineCause.
 func (t Trapper) WithDeadline(tags ...string) *Trap {
 	return t.mock.newTrap(clockFunctionWithDeadline, tags)
 }
 
-func (t Trapper) WithDeadlineCause(tags ...string) *Trap {
-	return t.mock.newTrap(clockFunctionWithDeadlineCause, tags)
-}
-
+// WithTimeout traps calls to WithTimeout and WithTimeoutCause.
 func (t Trapper) WithTimeout(tags ...string) *Trap {
 	return t.mock.newTrap(clockFunctionWithTimeout, tags)
 }
 
-func (t Trapper) WithTimeoutCause(tags ...string) *Trap {
-	return t.mock.newTrap(clockFunctionWithTimeoutCause, tags)
-}
-
 func (m *Mock) Trap() Trapper {
 	return Trapper{m}
 }
@@ -1085,6 +1079,19 @@ func (c clockFunction) String() string {
 	}
 }
 
+// trapFunction returns the clockFunction of the trap that catches calls to c. Cause variants share
+// the base function's trap so a test does not need to know which variant the code calls.
+func (c clockFunction) trapFunction() clockFunction {
+	switch c {
+	case clockFunctionWithDeadlineCause:
+		return clockFunctionWithDeadline
+	case clockFunctionWithTimeoutCause:
+		return clockFunctionWithTimeout
+	default:
+		return c
+	}
+}
+
 type callArg func(c *apiCall)
 
 // apiCall represents a single call to one of the Clock APIs.
@@ -1142,7 +1149,7 @@ func (a *apiCall) String() string {
 type Call struct {
 	Time     time.Time
 	Duration time.Duration
-	Cause    error // Cause passed to a cause variant, if applicable.
+	Cause    error // Cause passed to WithDeadlineCause or WithTimeoutCause, nil for other calls.
 	Tags     []string
 
 	tb      TestingT
@@ -1234,7 +1241,7 @@ func (t *Trap) catch(c *apiCall) {
 }
 
 func (t *Trap) matches(c *apiCall) bool {
-	if t.fn != c.fn {
+	if t.fn != c.fn.trapFunction() {
 		return false
 	}
 	for _, tag := range t.tags {

mock_test.go

@@ -1128,83 +1128,98 @@ func TestWithTimeout_CancelRace(t *testing.T) {
 	}
 }
 
-func TestWithTimeout_Trap(t *testing.T) {
+func TestTrap_Contexts(t *testing.T) {
 	t.Parallel()
-	testCtx, testCancel := context.WithTimeout(context.Background(), 10*time.Second)
-	defer testCancel()
-
-	mClock := quartz.NewMock(t)
-	trap := mClock.Trap().WithTimeout("ctx")
-	defer trap.Close()
-	ctxs := make(chan context.Context, 1)
-	cancels := make(chan context.CancelFunc, 1)
-	go func() {
-		ctx, cancel := mClock.WithTimeout(context.Background(), time.Hour, "ctx")
-		ctxs <- ctx
-		cancels <- cancel
-	}()
 
-	call := trap.MustWait(testCtx)
-	if call.Duration != time.Hour {
-		t.Fatalf("expected time.Hour, got %v", call.Duration)
-	}
-	call.MustRelease(testCtx)
-	var ctx context.Context
-	select {
-	case ctx = <-ctxs:
-	case <-testCtx.Done():
-		t.Fatal("timeout waiting for trapped WithTimeout context")
-	}
-	var cancel context.CancelFunc
-	select {
-	case cancel = <-cancels:
-	case <-testCtx.Done():
-		t.Fatal("timeout waiting for trapped WithTimeout cancel")
-	}
-	defer cancel()
-	mClock.Advance(time.Hour).MustWait(testCtx)
-	select {
-	case <-ctx.Done():
-		// OK
-	case <-testCtx.Done():
-		t.Fatal("timeout waiting for trapped context")
-	}
-	if !errors.Is(ctx.Err(), context.DeadlineExceeded) {
-		t.Fatalf("expected DeadlineExceeded, got %v", ctx.Err())
-	}
+	cause := errors.New("trap cause")
+	for _, tc := range []struct {
+		name         string
+		trap         func(quartz.Trapper, ...string) *quartz.Trap
+		call         func(*quartz.Mock) (context.Context, context.CancelFunc)
+		wantTime     func(*quartz.Mock) time.Time
+		wantDuration time.Duration
+		wantCause    error
+	}{
+		{
+			name: "WithDeadline",
+			trap: quartz.Trapper.WithDeadline,
+			call: func(m *quartz.Mock) (context.Context, context.CancelFunc) {
+				return m.WithDeadline(context.Background(), m.Now().Add(time.Hour))
+			},
+			wantTime: func(m *quartz.Mock) time.Time { return m.Now().Add(time.Hour) },
+		},
+		{
+			name: "WithDeadlineCause",
+			trap: quartz.Trapper.WithDeadline,
+			call: func(m *quartz.Mock) (context.Context, context.CancelFunc) {
+				return m.WithDeadlineCause(context.Background(), m.Now().Add(time.Hour), cause)
+			},
+			wantTime:  func(m *quartz.Mock) time.Time { return m.Now().Add(time.Hour) },
+			wantCause: cause,
+		},
+		{
+			name: "WithTimeout",
+			trap: quartz.Trapper.WithTimeout,
+			call: func(m *quartz.Mock) (context.Context, context.CancelFunc) {
+				return m.WithTimeout(context.Background(), time.Hour)
+			},
+			wantTime:     func(*quartz.Mock) time.Time { return time.Time{} },
+			wantDuration: time.Hour,
+		},
+		{
+			name: "WithTimeoutCause",
+			trap: quartz.Trapper.WithTimeout,
+			call: func(m *quartz.Mock) (context.Context, context.CancelFunc) {
+				return m.WithTimeoutCause(context.Background(), time.Hour, cause)
+			},
+			wantTime:     func(*quartz.Mock) time.Time { return time.Time{} },
+			wantDuration: time.Hour,
+			wantCause:    cause,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			testCtx, testCancel := context.WithTimeout(context.Background(), 10*time.Second)
+			defer testCancel()
 
-	causeErr := errors.New("trap cause")
-	trapCause := mClock.Trap().WithTimeoutCause("cause")
-	defer trapCause.Close()
-	go func() {
-		ctx, cancel := mClock.WithTimeoutCause(context.Background(), time.Minute, causeErr, "cause")
-		ctxs <- ctx
-		cancels <- cancel
-	}()
+			mClock := quartz.NewMock(t)
+			wantTime := tc.wantTime(mClock)
+			trap := tc.trap(mClock.Trap())
+			defer trap.Close()
+			type result struct {
+				ctx    context.Context
+				cancel context.CancelFunc
+			}
+			results := make(chan result, 1)
+			go func() {
+				ctx, cancel := tc.call(mClock)
+				results <- result{ctx, cancel}
+			}()
+
+			call := trap.MustWait(testCtx)
+			if !call.Time.Equal(wantTime) {
+				t.Fatalf("expected time %v, got %v", wantTime, call.Time)
+			}
+			if call.Duration != tc.wantDuration {
+				t.Fatalf("expected duration %v, got %v", tc.wantDuration, call.Duration)
+			}
+			if call.Cause != tc.wantCause {
+				t.Fatalf("expected cause %v, got %v", tc.wantCause, call.Cause)
+			}
+			call.MustRelease(testCtx)
+			var r result
+			select {
+			case r = <-results:
+			case <-testCtx.Done():
+				t.Fatal("timeout waiting for the released call to return")
+			}
+			defer r.cancel()
 
-	call = trapCause.MustWait(testCtx)
-	if call.Duration != time.Minute {
-		t.Fatalf("expected time.Minute, got %v", call.Duration)
-	}
-	if !errors.Is(call.Cause, causeErr) {
-		t.Fatalf("expected cause %v, got %v", causeErr, call.Cause)
-	}
-	call.MustRelease(testCtx)
-	select {
-	case ctx = <-ctxs:
-	case <-testCtx.Done():
-		t.Fatal("timeout waiting for trapped WithTimeoutCause context")
-	}
-	select {
-	case cancel = <-cancels:
-	case <-testCtx.Done():
-		t.Fatal("timeout waiting for trapped WithTimeoutCause cancel")
-	}
-	cancel()
-	select {
-	case <-ctx.Done():
-	case <-testCtx.Done():
-		t.Fatal("timeout waiting for canceled WithTimeoutCause ctx")
+			mClock.Advance(time.Hour).MustWait(testCtx)
+			if !errors.Is(r.ctx.Err(), context.DeadlineExceeded) {
+				t.Fatalf("expected DeadlineExceeded, got %v", r.ctx.Err())
+			}
+		})
 	}
 }