feat: Add T.Cleanup (#79)

* feat: Add T.Cleanup Adds a new T.Cleanup method to register cleanup functions. These will run after each iteration of a Check. The implementation of how cleanup functions are tracked and run is borrowed heavily from testing.T's own implementation. Namely: - [T.Cleanup puts functions in a slice](https://cs.opensource.google/go/go/+/refs/tags/go1.24.0:src/testing/testing.go;l=1214-1216) - [cleanups are run in reverse order](https://cs.opensource.google/go/go/+/refs/tags/go1.24.0:src/testing/testing.go;l=1433-1446) popping functions from the slice one by one - [panics are not allowed to interrupt cleanup](https://cs.opensource.google/go/go/+/refs/tags/go1.24.0:src/testing/testing.go;l=1420-1427) As with `testing.T`, cleanup runs after the context is canceled. Because Rapid's context is lazily initialized, it needs an additional check to avoid providing a valid context during cleanup. Resolves #62 * doc: try to be clearer

Abhinav Gupta committed 1 year ago

Commit a13e39b1 parent 3f4e8915 files · +277 −5Browse files

Unified · Split

5 files

combinators_external_test.go +22 −0
engine.go +79 −4
engine_fuzz_test.go +17 −0
engine_test.go +144 −0
generator_test.go +15 −1

combinators_external_test.go+22−0

@@ -95,6 +95,28 @@ func TestCustomContext(t *testing.T) { · 94 unmodified lines
95	95	})
96	96	}
97	97
	98	+func TestCustomCleanup(t *testing.T) {
	99	+ t.Parallel()
	100	+
	101	+ var open bool
	102	+ gen := Custom(func(t *T) int {
	103	+ t.Cleanup(func() { open = false })
	104	+ return Int().Draw(t, "")
	105	+ })
	106	+
	107	+ // Cleanup functions registered during a Custom generator
	108	+ // are run after generation, not after the Check.
	109	+ Check(t, func(t *T) {
	110	+ open = true
	111	+ _ = gen.Draw(t, "value")
	112	+
	113	+ // Cleanup must run after each run of the custom generator.
	114	+ if open {
	115	+ t.Fatalf("cleanup must be run")
	116	+ }
	117	+ })
	118	+}
	119	+
98	120	func TestFilter(t *testing.T) {
99	121	t.Parallel()
100	122

engine.go+79−4

@@ -19,6 +19,7 @@ import ( · 18 unmodified lines
19	19	"runtime"
20	20	"strings"
21	21	"sync"
	22	+ "sync/atomic"
22	23	"testing"
23	24	"time"
24	25	)
@@ -511,6 +512,8 @@ type T struct { · 486 unmodified lines
511	512
512	513	ctx context.Context
513	514	cancelCtx context.CancelFunc
	515	+ cleanups []func()
	516	+ cleaning atomic.Bool
514	517
515	518	tbLog bool
516	519	rawLog *log.Logger
@@ -550,8 +553,17 @@ func (t *T) shouldLog() bool { · 33 unmodified lines
550	553	return t.rawLog != nil \|\| t.tbLog
551	554	}
552	555
553		−// Context returns a context.Context associated with the test.
554		−// It is valid only for the duration of the rapid check.
	556	+// Context returns a context.Context that is canceled
	557	+// after the property function exits,
	558	+// before Cleanup-registered functions are run.
	559	+//
	560	+// For [Check], [MakeFuzz], and similar functions,
	561	+// each call to the property function gets a unique context
	562	+// that is canceled after that property function exits.
	563	+//
	564	+// For [Custom], each time a new value is generated,
	565	+// the generator function gets a unique context
	566	+// that is canceled after the generator function exits.
555	567	func (t *T) Context() context.Context {
556	568	// Fast path: no need to lock if the context is already set.
557	569	t.mu.RLock()
@@ -561,6 +573,15 @@ func (t *T) Context() context.Context { · 3 unmodified lines
561	573	return ctx
562	574	}
563	575
	576	+ // If we're in the middle of cleaning up
	577	+ // and the context has already been canceled and cleared,
	578	+ // don't create a new one. Return a canceled context instead.
	579	+ if t.cleaning.Load() {
	580	+ ctx, cancel := context.WithCancel(context.Background())
	581	+ cancel()
	582	+ return ctx
	583	+ }
	584	+
564	585	// Slow path: lock and check again, create new context if needed.
565	586	t.mu.Lock()
566	587	defer t.mu.Unlock()
@@ -587,17 +608,71 @@ func (t *T) Context() context.Context { · 20 unmodified lines
587	608	return ctx
588	609	}
589	610
	611	+// Cleanup registers a function to be called
	612	+// when a property function finishes running.
	613	+//
	614	+// For [Check], [MakeFuzz], and similar functions,
	615	+// each call to the property function registers its cleanup functions,
	616	+// which are called after the property function exits.
	617	+//
	618	+// For [Custom], each time a new value is generated,
	619	+// the generator function registers its cleanup functions,
	620	+// which are called after the generator function exits.
	621	+//
	622	+// Cleanup functions are called in last-in, first-out order.
	623	+//
	624	+// If [T.Context] is used, the context is canceled
	625	+// before the Cleanup functions are executed.
	626	+func (t *T) Cleanup(f func()) {
	627	+ t.mu.Lock()
	628	+ defer t.mu.Unlock()
	629	+
	630	+ t.cleanups = append(t.cleanups, f)
	631	+}
	632	+
590	633	// cleanup runs any cleanup tasks associated with the property check.
591	634	// It is safe to call multiple times.
592	635	func (t *T) cleanup() {
593		− t.mu.Lock()
594		− defer t.mu.Unlock()
	636	+ t.cleaning.Store(true)
	637	+ defer t.cleaning.Store(false)
	638	+
	639	+ // If a cleanup function panics,
	640	+ // we still want to run the remaining cleanup functions.
	641	+ defer func() {
	642	+ t.mu.Lock()
	643	+ recurse := len(t.cleanups) > 0
	644	+ t.mu.Unlock()
	645	+
	646	+ if recurse {
	647	+ t.cleanup()
	648	+ }
	649	+ }()
595	650
	651	+ // Context must be closed before t.Cleanup functions are run.
	652	+ t.mu.Lock()
596	653	if t.cancelCtx != nil {
597	654	t.cancelCtx()
598	655	t.cancelCtx = nil
599	656	t.ctx = nil
600	657	}
	658	+ t.mu.Unlock()
	659	+
	660	+ for {
	661	+ var cleanup func()
	662	+ t.mu.Lock()
	663	+ if len(t.cleanups) > 0 {
	664	+ last := len(t.cleanups) - 1
	665	+ cleanup = t.cleanups[last]
	666	+ t.cleanups = t.cleanups[:last]
	667	+ }
	668	+ t.mu.Unlock()
	669	+
	670	+ if cleanup == nil {
	671	+ break
	672	+ }
	673	+
	674	+ cleanup()
	675	+ }
601	676	}
602	677
603	678	func (t *T) Logf(format string, args ...any) {

engine_fuzz_test.go+17−0

@@ -96,3 +96,20 @@ func FuzzContext(f *testing.F) { · 95 unmodified lines
96	96	}
97	97	}
98	98	}
	99	+
	100	+func FuzzCleanup(f *testing.F) {
	101	+ var state []bool
	102	+ f.Fuzz(MakeFuzz(func(t *T) {
	103	+ idx := len(state)
	104	+ state = append(state, false)
	105	+ t.Cleanup(func() {
	106	+ state[idx] = true
	107	+ })
	108	+ }))
	109	+
	110	+ for _, ok := range state {
	111	+ if !ok {
	112	+ f.Fatalf("cleanup must be called")
	113	+ }
	114	+ }
	115	+}

engine_test.go+144−0

@@ -9,6 +9,7 @@ package rapid · 8 unmodified lines
9	9	import (
10	10	"context"
11	11	"errors"
	12	+ "reflect"
12	13	"strings"
13	14	"testing"
14	15	)
@@ -96,6 +97,8 @@ func BenchmarkCheckOverhead(b *testing.B) { · 81 unmodified lines
96	97	}
97	98
98	99	func TestCheckContext(t *testing.T) {
	100	+ t.Parallel()
	101	+
99	102	type key struct{}
100	103
101	104	var ctx context.Context
@@ -114,3 +117,144 @@ func TestCheckContext(t *testing.T) { · 12 unmodified lines
114	117	t.Fatalf("context must have a value")
115	118	}
116	119	}
	120	+
	121	+func TestCheckCleanup(t *testing.T) {
	122	+ t.Parallel()
	123	+
	124	+ // Each Check iteration will append a true to indicate "open",
	125	+ // and flip it to false on cleanup.
	126	+ //
	127	+ // After the check is done, we expect all values to be false.
	128	+ var state []bool
	129	+
	130	+ Check(t, func(t *T) {
	131	+ idx := len(state)
	132	+ state = append(state, true)
	133	+ t.Cleanup(func() {
	134	+ state[idx] = false
	135	+ })
	136	+ })
	137	+
	138	+ for _, v := range state {
	139	+ if v {
	140	+ t.Fatalf("expected all values to be false")
	141	+ }
	142	+ }
	143	+}
	144	+
	145	+func TestCheckCleanupMultipleOrder(t *testing.T) {
	146	+ t.Parallel()
	147	+
	148	+ // If multiple cleanups are appended during a Check,
	149	+ // they must run in reverse order.
	150	+ var state []int
	151	+ Check(t, func(t *T) {
	152	+ // We just want to capture the result of one iteration,
	153	+ // so we'll keep resetting the state.
	154	+ state = nil
	155	+ t.Cleanup(func() {
	156	+ state = append(state, 1)
	157	+ })
	158	+ t.Cleanup(func() {
	159	+ state = append(state, 2)
	160	+ })
	161	+ t.Cleanup(func() {
	162	+ state = append(state, 3)
	163	+ })
	164	+ })
	165	+
	166	+ if !reflect.DeepEqual(state, []int{3, 2, 1}) {
	167	+ t.Fatalf("expected cleanups to run in reverse order, got: %v", state)
	168	+ }
	169	+}
	170	+
	171	+func TestCheckCleanupPanic(t *testing.T) {
	172	+ t.Parallel()
	173	+
	174	+ // A Cleanup function halfway through will panic.
	175	+ // Deferred assertions will check that all values are false.
	176	+ var state []bool
	177	+ defer func() {
	178	+ for _, v := range state {
	179	+ if v {
	180	+ t.Errorf("expected all values to be false")
	181	+ }
	182	+ }
	183	+ }()
	184	+
	185	+ Check(ignoreErrorsTB{t}, func(t *T) {
	186	+ idx := len(state)
	187	+ state = append(state, true)
	188	+ t.Cleanup(func() {
	189	+ state[idx] = false
	190	+ if idx == len(state)/2 {
	191	+ panic("cleanup panic")
	192	+ }
	193	+ })
	194	+ })
	195	+}
	196	+
	197	+func TestCheckCleanupNewCleanupsDuringCleanup(t *testing.T) {
	198	+ t.Parallel()
	199	+
	200	+ // Cleanups can be added during cleanup.
	201	+ var state []bool
	202	+ Check(t, func(t *T) {
	203	+ idx := len(state)
	204	+ state = append(state, true)
	205	+ t.Cleanup(func() {
	206	+ // Odd numbered events will add a new cleanup.
	207	+ if idx%2 == 0 {
	208	+ state[idx] = false
	209	+ } else {
	210	+ t.Cleanup(func() {
	211	+ state[idx] = false
	212	+ })
	213	+ }
	214	+ })
	215	+ })
	216	+}
	217	+
	218	+func TestCheckCleanupContextIsCanceled(t *testing.T) {
	219	+ t.Parallel()
	220	+
	221	+ // Context created during Check is canceled by the time Cleanup is run.
	222	+ Check(t, func(t *T) {
	223	+ ctx := t.Context()
	224	+ t.Cleanup(func() {
	225	+ if err := ctx.Err(); err == nil \|\| !errors.Is(err, context.Canceled) {
	226	+ t.Fatalf("expected context to be canceled, got: %v", ctx)
	227	+ }
	228	+ })
	229	+ })
	230	+}
	231	+
	232	+func TestCheckCleanupContextCreatedInCleanup(t *testing.T) {
	233	+ t.Parallel()
	234	+
	235	+ // Context created during Cleanup is already canceled.
	236	+ Check(t, func(t *T) {
	237	+ ctx := t.Context()
	238	+ t.Cleanup(func() {
	239	+ // ctx is already cleared on rapid.T by now,
	240	+ // so this will request a new context.
	241	+ newCtx := t.Context()
	242	+ if ctx == newCtx {
	243	+ t.Fatalf("expected new context")
	244	+ }
	245	+
	246	+ if err := newCtx.Err(); err == nil \|\| !errors.Is(err, context.Canceled) {
	247	+ t.Fatalf("expected context to be canceled, got: %v", newCtx)
	248	+ }
	249	+ })
	250	+ })
	251	+}
	252	+
	253	+// ignoreErrorsTB is a TB that ignores all errors posted to it.
	254	+type ignoreErrorsTB struct{ TB }
	255	+
	256	+func (ignoreErrorsTB) Error(...interface{}) {}
	257	+func (ignoreErrorsTB) Errorf(string, ...interface{}) {}
	258	+func (ignoreErrorsTB) Fatal(...interface{}) {}
	259	+func (ignoreErrorsTB) Fatalf(string, ...interface{}) {}
	260	+func (ignoreErrorsTB) Fail() {}

generator_test.go+15−1

@@ -46,7 +46,7 @@ func TestExampleHelper(t *testing.T) { · 45 unmodified lines
46	46	g.Example(0)
47	47	}
48	48
49		−func TestExampleContext(t *testing.T) {
	49	+func TestCustomExampleContext(t *testing.T) {
50	50	type key struct{}
51	51
52	52	g := Custom(func(t *T) context.Context {
@@ -67,3 +67,17 @@ func TestExampleContext(t *testing.T) { · 14 unmodified lines
67	67	t.Fatalf("context must have a value")
68	68	}
69	69	}
	70	+
	71	+func TestCustomExampleCleanup(t *testing.T) {
	72	+ var state bool
	73	+ g := Custom(func(t *T) int {
	74	+ t.Cleanup(func() { state = false })
	75	+ return Int().Draw(t, "x")
	76	+ })
	77	+
	78	+ state = true
	79	+ _ = g.Example(0)
	80	+ if state {
	81	+ t.Fatalf("cleanup must be called")
	82	+ }
	83	+}

feat: Add T.Cleanup (#79)

Keyboard shortcuts

General

Issues and pull requests

Issues, pull requests and Inbox

Repository code