From 386ef912061c205bea16c701ef5fc9f230c8b79f Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Tue, 3 Nov 2026 12:22:19 +0000 Subject: [PATCH 1/6] fix: catch cause variants in context traps Trap().WithTimeout() did not catch WithTimeoutCause() calls, and Trap().WithDeadline() did not catch WithDeadlineCause(), so a test had to know which variant the code under test calls to trap it. The cause variants now match their base function's trap, and Call.Cause still exposes the cause. This removes the separate WithDeadlineCause and WithTimeoutCause trappers, which are not on main or in any release. --- mock.go | 27 +++++++---- mock_test.go | 127 ++++++++++++++++++++++++++++++++++++++------------- 2 files changed, 113 insertions(+), 41 deletions(-) diff --git a/mock.go b/mock.go index 5e9a4bb..a38476e 100644 --- a/mock.go +++ b/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 function of the Trapper method that catches calls to c. Cause variants +// have no Trapper method of their own, so a trap for the base function catches both. +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 { diff --git a/mock_test.go b/mock_test.go index 37a3fe1..7712d0e 100644 --- a/mock_test.go +++ b/mock_test.go @@ -7,6 +7,7 @@ import ( "fmt" "os" "runtime/pprof" + "slices" "strings" "sync" "testing" @@ -1172,39 +1173,103 @@ func TestWithTimeout_Trap(t *testing.T) { if !errors.Is(ctx.Err(), context.DeadlineExceeded) { t.Fatalf("expected DeadlineExceeded, got %v", ctx.Err()) } +} - 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 - }() +func TestTrap_CatchesCauseVariants(t *testing.T) { + t.Parallel() - 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") + cause := errors.New("trap cause") + deadline := time.Date(2030, 1, 1, 0, 0, 0, 0, time.UTC) + for _, tc := range []struct { + name string + trap func(quartz.Trapper, ...string) *quartz.Trap + call func(m *quartz.Mock, tags ...string) (context.Context, context.CancelFunc) + wantTime time.Time + wantDuration time.Duration + wantCause error + }{ + { + name: "WithDeadline", + trap: quartz.Trapper.WithDeadline, + call: func(m *quartz.Mock, tags ...string) (context.Context, context.CancelFunc) { + return m.WithDeadline(context.Background(), deadline, tags...) + }, + wantTime: deadline, + }, + { + name: "WithDeadlineCause", + trap: quartz.Trapper.WithDeadline, + call: func(m *quartz.Mock, tags ...string) (context.Context, context.CancelFunc) { + return m.WithDeadlineCause(context.Background(), deadline, cause, tags...) + }, + wantTime: deadline, + wantCause: cause, + }, + { + name: "WithTimeout", + trap: quartz.Trapper.WithTimeout, + call: func(m *quartz.Mock, tags ...string) (context.Context, context.CancelFunc) { + return m.WithTimeout(context.Background(), time.Hour, tags...) + }, + wantDuration: time.Hour, + }, + { + name: "WithTimeoutCause", + trap: quartz.Trapper.WithTimeout, + call: func(m *quartz.Mock, tags ...string) (context.Context, context.CancelFunc) { + return m.WithTimeoutCause(context.Background(), time.Hour, cause, tags...) + }, + 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() + + mClock := quartz.NewMock(t) + trap := tc.trap(mClock.Trap(), "match") + defer trap.Close() + cancels := make(chan context.CancelFunc, 1) + + // The trap does not catch a call without its tags, so the call returns unreleased. + go func() { + _, cancel := tc.call(mClock, "other") + cancels <- cancel + }() + select { + case cancel := <-cancels: + cancel() + case <-testCtx.Done(): + t.Fatal("timeout waiting for call without the trap's tags") + } + + go func() { + _, cancel := tc.call(mClock, "match") + cancels <- cancel + }() + call := trap.MustWait(testCtx) + if !slices.Equal(call.Tags, []string{"match"}) { + t.Fatalf("expected tags [match], got %v", call.Tags) + } + if !call.Time.Equal(tc.wantTime) { + t.Fatalf("expected time %v, got %v", tc.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) + select { + case cancel := <-cancels: + cancel() + case <-testCtx.Done(): + t.Fatal("timeout waiting for released call to return") + } + }) } } From 1a3ee3a4a74b3027724f48a3b9f4e3db9221ce77 Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Tue, 3 Nov 2026 15:06:26 +0000 Subject: [PATCH 2/6] docs(mock.go): say why cause variants share the base trap --- mock.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mock.go b/mock.go index a38476e..edd781e 100644 --- a/mock.go +++ b/mock.go @@ -1079,8 +1079,8 @@ func (c clockFunction) String() string { } } -// trapFunction returns the function of the Trapper method that catches calls to c. Cause variants -// have no Trapper method of their own, so a trap for the base function catches both. +// 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: From 29e91800322311f92441b25846a6a1bc3547aa55 Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Tue, 3 Nov 2026 15:09:08 +0000 Subject: [PATCH 3/6] test(mock_test.go): cover context traps in one table TestWithTimeout_Trap duplicated the WithTimeout row of the cause variant test. TestTrap_Contexts now has one row per context function, and each row also checks that the released context expires on Advance, which only the WithTimeout case checked before. The negative control for a call without the trap's tags is gone. It tested generic tag filtering, which Test_MultipleTraps covers, and did not control for which function a trap catches. --- mock_test.go | 120 +++++++++++++++------------------------------------ 1 file changed, 35 insertions(+), 85 deletions(-) diff --git a/mock_test.go b/mock_test.go index 7712d0e..5a22c6a 100644 --- a/mock_test.go +++ b/mock_test.go @@ -7,7 +7,6 @@ import ( "fmt" "os" "runtime/pprof" - "slices" "strings" "sync" "testing" @@ -1129,96 +1128,51 @@ func TestWithTimeout_CancelRace(t *testing.T) { } } -func TestWithTimeout_Trap(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()) - } -} - -func TestTrap_CatchesCauseVariants(t *testing.T) { +func TestTrap_Contexts(t *testing.T) { t.Parallel() cause := errors.New("trap cause") - deadline := time.Date(2030, 1, 1, 0, 0, 0, 0, time.UTC) for _, tc := range []struct { name string trap func(quartz.Trapper, ...string) *quartz.Trap - call func(m *quartz.Mock, tags ...string) (context.Context, context.CancelFunc) - wantTime time.Time + 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, tags ...string) (context.Context, context.CancelFunc) { - return m.WithDeadline(context.Background(), deadline, tags...) + call: func(m *quartz.Mock) (context.Context, context.CancelFunc) { + return m.WithDeadline(context.Background(), m.Now().Add(time.Hour)) }, - wantTime: deadline, + wantTime: func(m *quartz.Mock) time.Time { return m.Now().Add(time.Hour) }, }, { name: "WithDeadlineCause", trap: quartz.Trapper.WithDeadline, - call: func(m *quartz.Mock, tags ...string) (context.Context, context.CancelFunc) { - return m.WithDeadlineCause(context.Background(), deadline, cause, tags...) + call: func(m *quartz.Mock) (context.Context, context.CancelFunc) { + return m.WithDeadlineCause(context.Background(), m.Now().Add(time.Hour), cause) }, - wantTime: deadline, + 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, tags ...string) (context.Context, context.CancelFunc) { - return m.WithTimeout(context.Background(), time.Hour, tags...) + 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, tags ...string) (context.Context, context.CancelFunc) { - return m.WithTimeoutCause(context.Background(), time.Hour, cause, tags...) + 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, }, @@ -1229,32 +1183,22 @@ func TestTrap_CatchesCauseVariants(t *testing.T) { defer testCancel() mClock := quartz.NewMock(t) - trap := tc.trap(mClock.Trap(), "match") + wantTime := tc.wantTime(mClock) + trap := tc.trap(mClock.Trap()) defer trap.Close() - cancels := make(chan context.CancelFunc, 1) - - // The trap does not catch a call without its tags, so the call returns unreleased. - go func() { - _, cancel := tc.call(mClock, "other") - cancels <- cancel - }() - select { - case cancel := <-cancels: - cancel() - case <-testCtx.Done(): - t.Fatal("timeout waiting for call without the trap's tags") + type result struct { + ctx context.Context + cancel context.CancelFunc } - + results := make(chan result, 1) go func() { - _, cancel := tc.call(mClock, "match") - cancels <- cancel + ctx, cancel := tc.call(mClock) + results <- result{ctx, cancel} }() + call := trap.MustWait(testCtx) - if !slices.Equal(call.Tags, []string{"match"}) { - t.Fatalf("expected tags [match], got %v", call.Tags) - } - if !call.Time.Equal(tc.wantTime) { - t.Fatalf("expected time %v, got %v", tc.wantTime, call.Time) + 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) @@ -1263,11 +1207,17 @@ func TestTrap_CatchesCauseVariants(t *testing.T) { t.Fatalf("expected cause %v, got %v", tc.wantCause, call.Cause) } call.MustRelease(testCtx) + var r result select { - case cancel := <-cancels: - cancel() + case r = <-results: case <-testCtx.Done(): - t.Fatal("timeout waiting for released call to return") + t.Fatal("timeout waiting for the released call to return") + } + defer r.cancel() + + mClock.Advance(time.Hour).MustWait(testCtx) + if !errors.Is(r.ctx.Err(), context.DeadlineExceeded) { + t.Fatalf("expected DeadlineExceeded, got %v", r.ctx.Err()) } }) } From fdd6ba1605b51f1243c55c6bf5d882710d0e937d Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Wed, 4 Nov 2026 18:20:04 +0000 Subject: [PATCH 4/6] docs(README.md): say the context traps catch the cause variants A user reading the Contexts section had no way to learn that Trap().WithTimeout() and Trap().WithDeadline() also catch the cause variants, now that the separate cause trappers are gone. --- README.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/README.md b/README.md index 1f9dc0f..049eb5b 100644 --- a/README.md +++ b/README.md @@ -367,6 +367,11 @@ contexts created after the test ended are returned already canceled with that ca Call the returned cancel function as you would for `context.WithTimeout`. It is idempotent and removes the pending mock event. Cause variants preserve the cause for `context.Cause`. +To trap the calls that create a context, use `mClock.Trap().WithTimeout()` or +`mClock.Trap().WithDeadline()`. They also catch calls to `WithTimeoutCause` and `WithDeadlineCause`, +so the test keeps working if the code switches variants. The trapped call's `Cause` field holds the +cause passed to a cause variant. + #### `Deadline()` does not report mock time With a `*Mock`, `ctx.Done()` and `ctx.Err()` follow mock time, but `ctx.Deadline()` reports the From 3b702b1d5d3963e84357c511b4b91c199edd9434 Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Wed, 4 Nov 2026 18:48:27 +0000 Subject: [PATCH 5/6] docs(README.md): pair each context trap with its cause variant --- README.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 049eb5b..0c5db2e 100644 --- a/README.md +++ b/README.md @@ -367,10 +367,10 @@ contexts created after the test ended are returned already canceled with that ca Call the returned cancel function as you would for `context.WithTimeout`. It is idempotent and removes the pending mock event. Cause variants preserve the cause for `context.Cause`. -To trap the calls that create a context, use `mClock.Trap().WithTimeout()` or -`mClock.Trap().WithDeadline()`. They also catch calls to `WithTimeoutCause` and `WithDeadlineCause`, -so the test keeps working if the code switches variants. The trapped call's `Cause` field holds the -cause passed to a cause variant. +To trap the calls that create a context, use `mClock.Trap().WithTimeout()`, which also catches +`WithTimeoutCause` calls, or `mClock.Trap().WithDeadline()`, which also catches `WithDeadlineCause` +calls. The test then keeps working if the code switches variants. The trapped call's `Cause` field +holds the cause passed to a cause variant. #### `Deadline()` does not report mock time From e722958265f1f0409ae85c96ff366d03381182ee Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Wed, 4 Nov 2026 19:10:15 +0000 Subject: [PATCH 6/6] docs(README.md): move the context trap note next to the example The note now follows the first Contexts example instead of the cancel paragraph, where #32 and #33 also add text, so the three branches do not insert at the same place. It also pairs each trap with the calls it catches in one sentence. --- README.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 0c5db2e..0dadacb 100644 --- a/README.md +++ b/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. @@ -367,11 +372,6 @@ contexts created after the test ended are returned already canceled with that ca Call the returned cancel function as you would for `context.WithTimeout`. It is idempotent and removes the pending mock event. Cause variants preserve the cause for `context.Cause`. -To trap the calls that create a context, use `mClock.Trap().WithTimeout()`, which also catches -`WithTimeoutCause` calls, or `mClock.Trap().WithDeadline()`, which also catches `WithDeadlineCause` -calls. The test then keeps working if the code switches variants. The trapped call's `Cause` field -holds the cause passed to a cause variant. - #### `Deadline()` does not report mock time With a `*Mock`, `ctx.Done()` and `ctx.Err()` follow mock time, but `ctx.Deadline()` reports the