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

Author: abhinav

combinators_external_test.go

@@ -95,6 +95,28 @@ func TestCustomContext(t *testing.T) {
 	})
 }
 
+func TestCustomCleanup(t *testing.T) {
+	t.Parallel()
+
+	var open bool
+	gen := Custom(func(t *T) int {
+		t.Cleanup(func() { open = false })
+		return Int().Draw(t, "")
+	})
+
+	// Cleanup functions registered during a Custom generator
+	// are run after generation, not after the Check.
+	Check(t, func(t *T) {
+		open = true
+		_ = gen.Draw(t, "value")
+
+		// Cleanup must run after each run of the custom generator.
+		if open {
+			t.Fatalf("cleanup must be run")
+		}
+	})
+}
+
 func TestFilter(t *testing.T) {
 	t.Parallel()
 

engine.go

@@ -19,6 +19,7 @@ import (
 	"runtime"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 )
@@ -506,6 +507,8 @@ type T struct {
 
 	ctx       context.Context
 	cancelCtx context.CancelFunc
+	cleanups  []func()
+	cleaning  atomic.Bool
 
 	tbLog    bool
 	rawLog   *log.Logger
@@ -545,8 +548,17 @@ func (t *T) shouldLog() bool {
 	return t.rawLog != nil || t.tbLog
 }
 
-// Context returns a context.Context associated with the test.
-// It is valid only for the duration of the rapid check.
+// Context returns a context.Context that is canceled
+// after the property function exits,
+// before Cleanup-registered functions are run.
+//
+// For [Check], [MakeFuzz], and similar functions,
+// each call to the property function gets a unique context
+// that is canceled after that property function exits.
+//
+// For [Custom], each time a new value is generated,
+// the generator function gets a unique context
+// that is canceled after the generator function exits.
 func (t *T) Context() context.Context {
 	// Fast path: no need to lock if the context is already set.
 	t.mu.RLock()
@@ -556,6 +568,15 @@ func (t *T) Context() context.Context {
 		return ctx
 	}
 
+	// If we're in the middle of cleaning up
+	// and the context has already been canceled and cleared,
+	// don't create a new one. Return a canceled context instead.
+	if t.cleaning.Load() {
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		return ctx
+	}
+
 	// Slow path: lock and check again, create new context if needed.
 	t.mu.Lock()
 	defer t.mu.Unlock()
@@ -582,17 +603,71 @@ func (t *T) Context() context.Context {
 	return ctx
 }
 
+// Cleanup registers a function to be called
+// when a property function finishes running.
+//
+// For [Check], [MakeFuzz], and similar functions,
+// each call to the property function registers its cleanup functions,
+// which are called after the property function exits.
+//
+// For [Custom], each time a new value is generated,
+// the generator function registers its cleanup functions,
+// which are called after the generator function exits.
+//
+// Cleanup functions are called in last-in, first-out order.
+//
+// If [T.Context] is used, the context is canceled
+// before the Cleanup functions are executed.
+func (t *T) Cleanup(f func()) {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
+	t.cleanups = append(t.cleanups, f)
+}
+
 // cleanup runs any cleanup tasks associated with the property check.
 // It is safe to call multiple times.
 func (t *T) cleanup() {
-	t.mu.Lock()
-	defer t.mu.Unlock()
+	t.cleaning.Store(true)
+	defer t.cleaning.Store(false)
+
+	// If a cleanup function panics,
+	// we still want to run the remaining cleanup functions.
+	defer func() {
+		t.mu.Lock()
+		recurse := len(t.cleanups) > 0
+		t.mu.Unlock()
+
+		if recurse {
+			t.cleanup()
+		}
+	}()
 
+	// Context must be closed before t.Cleanup functions are run.
+	t.mu.Lock()
 	if t.cancelCtx != nil {
 		t.cancelCtx()
 		t.cancelCtx = nil
 		t.ctx = nil
 	}
+	t.mu.Unlock()
+
+	for {
+		var cleanup func()
+		t.mu.Lock()
+		if len(t.cleanups) > 0 {
+			last := len(t.cleanups) - 1
+			cleanup = t.cleanups[last]
+			t.cleanups = t.cleanups[:last]
+		}
+		t.mu.Unlock()
+
+		if cleanup == nil {
+			break
+		}
+
+		cleanup()
+	}
 }
 
 func (t *T) Logf(format string, args ...any) {

engine_fuzz_test.go

@@ -92,3 +92,20 @@ func FuzzContext(f *testing.F) {
 		}
 	}
 }
+
+func FuzzCleanup(f *testing.F) {
+	var state []bool
+	f.Fuzz(MakeFuzz(func(t *T) {
+		idx := len(state)
+		state = append(state, false)
+		t.Cleanup(func() {
+			state[idx] = true
+		})
+	}))
+
+	for _, ok := range state {
+		if !ok {
+			f.Fatalf("cleanup must be called")
+		}
+	}
+}

engine_test.go

@@ -9,6 +9,7 @@ package rapid
 import (
 	"context"
 	"errors"
+	"reflect"
 	"strings"
 	"testing"
 )
@@ -97,6 +98,8 @@ func BenchmarkCheckOverhead(b *testing.B) {
 }
 
 func TestCheckContext(t *testing.T) {
+	t.Parallel()
+
 	type key struct{}
 
 	var ctx context.Context
@@ -115,3 +118,144 @@ func TestCheckContext(t *testing.T) {
 		t.Fatalf("context must have a value")
 	}
 }
+
+func TestCheckCleanup(t *testing.T) {
+	t.Parallel()
+
+	// Each Check iteration will append a true to indicate "open",
+	// and flip it to false on cleanup.
+	//
+	// After the check is done, we expect all values to be false.
+	var state []bool
+
+	Check(t, func(t *T) {
+		idx := len(state)
+		state = append(state, true)
+		t.Cleanup(func() {
+			state[idx] = false
+		})
+	})
+
+	for _, v := range state {
+		if v {
+			t.Fatalf("expected all values to be false")
+		}
+	}
+}
+
+func TestCheckCleanupMultipleOrder(t *testing.T) {
+	t.Parallel()
+
+	// If multiple cleanups are appended during a Check,
+	// they must run in reverse order.
+	var state []int
+	Check(t, func(t *T) {
+		// We just want to capture the result of one iteration,
+		// so we'll keep resetting the state.
+		state = nil
+		t.Cleanup(func() {
+			state = append(state, 1)
+		})
+		t.Cleanup(func() {
+			state = append(state, 2)
+		})
+		t.Cleanup(func() {
+			state = append(state, 3)
+		})
+	})
+
+	if !reflect.DeepEqual(state, []int{3, 2, 1}) {
+		t.Fatalf("expected cleanups to run in reverse order, got: %v", state)
+	}
+}
+
+func TestCheckCleanupPanic(t *testing.T) {
+	t.Parallel()
+
+	// A Cleanup function halfway through will panic.
+	// Deferred assertions will check that all values are false.
+	var state []bool
+	defer func() {
+		for _, v := range state {
+			if v {
+				t.Errorf("expected all values to be false")
+			}
+		}
+	}()
+
+	Check(ignoreErrorsTB{t}, func(t *T) {
+		idx := len(state)
+		state = append(state, true)
+		t.Cleanup(func() {
+			state[idx] = false
+			if idx == len(state)/2 {
+				panic("cleanup panic")
+			}
+		})
+	})
+}
+
+func TestCheckCleanupNewCleanupsDuringCleanup(t *testing.T) {
+	t.Parallel()
+
+	// Cleanups can be added during cleanup.
+	var state []bool
+	Check(t, func(t *T) {
+		idx := len(state)
+		state = append(state, true)
+		t.Cleanup(func() {
+			// Odd numbered events will add a new cleanup.
+			if idx%2 == 0 {
+				state[idx] = false
+			} else {
+				t.Cleanup(func() {
+					state[idx] = false
+				})
+			}
+		})
+	})
+}
+
+func TestCheckCleanupContextIsCanceled(t *testing.T) {
+	t.Parallel()
+
+	// Context created during Check is canceled by the time Cleanup is run.
+	Check(t, func(t *T) {
+		ctx := t.Context()
+		t.Cleanup(func() {
+			if err := ctx.Err(); err == nil || !errors.Is(err, context.Canceled) {
+				t.Fatalf("expected context to be canceled, got: %v", ctx)
+			}
+		})
+	})
+}
+
+func TestCheckCleanupContextCreatedInCleanup(t *testing.T) {
+	t.Parallel()
+
+	// Context created during Cleanup is already canceled.
+	Check(t, func(t *T) {
+		ctx := t.Context()
+		t.Cleanup(func() {
+			// ctx is already cleared on rapid.T by now,
+			// so this will request a new context.
+			newCtx := t.Context()
+			if ctx == newCtx {
+				t.Fatalf("expected new context")
+			}
+
+			if err := newCtx.Err(); err == nil || !errors.Is(err, context.Canceled) {
+				t.Fatalf("expected context to be canceled, got: %v", newCtx)
+			}
+		})
+	})
+}
+
+// ignoreErrorsTB is a TB that ignores all errors posted to it.
+type ignoreErrorsTB struct{ TB }
+
+func (ignoreErrorsTB) Error(...interface{})          {}
+func (ignoreErrorsTB) Errorf(string, ...interface{}) {}
+func (ignoreErrorsTB) Fatal(...interface{})          {}
+func (ignoreErrorsTB) Fatalf(string, ...interface{}) {}
+func (ignoreErrorsTB) Fail()                         {}

generator_test.go

@@ -46,7 +46,7 @@ func TestExampleHelper(t *testing.T) {
 	g.Example(0)
 }
 
-func TestExampleContext(t *testing.T) {
+func TestCustomExampleContext(t *testing.T) {
 	type key struct{}
 
 	g := Custom(func(t *T) context.Context {
@@ -67,3 +67,17 @@ func TestExampleContext(t *testing.T) {
 		t.Fatalf("context must have a value")
 	}
 }
+
+func TestCustomExampleCleanup(t *testing.T) {
+	var state bool
+	g := Custom(func(t *T) int {
+		t.Cleanup(func() { state = false })
+		return Int().Draw(t, "x")
+	})
+
+	state = true
+	_ = g.Example(0)
+	if state {
+		t.Fatalf("cleanup must be called")
+	}
+}