* 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
Unified · Split
| @@ -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 | func TestFilter(t *testing.T) { | 120 | func TestFilter(t *testing.T) { |
| 99 | t.Parallel() | 121 | t.Parallel() |
| 100 | 122 | ||
| @@ -19,6 +19,7 @@ import ( · 18 unmodified lines | |||
| 19 | "runtime" | 19 | "runtime" |
| 20 | "strings" | 20 | "strings" |
| 21 | "sync" | 21 | "sync" |
| 22 | "sync/atomic" | ||
| 22 | "testing" | 23 | "testing" |
| 23 | "time" | 24 | "time" |
| 24 | ) | 25 | ) |
| @@ -511,6 +512,8 @@ type T struct { · 486 unmodified lines | |||
| 511 | 512 | ||
| 512 | ctx context.Context | 513 | ctx context.Context |
| 513 | cancelCtx context.CancelFunc | 514 | cancelCtx context.CancelFunc |
| 515 | cleanups []func() | ||
| 516 | cleaning atomic.Bool | ||
| 514 | 517 | ||
| 515 | tbLog bool | 518 | tbLog bool |
| 516 | rawLog *log.Logger | 519 | rawLog *log.Logger |
| @@ -550,8 +553,17 @@ func (t *T) shouldLog() bool { · 33 unmodified lines | |||
| 550 | return t.rawLog != nil || t.tbLog | 553 | return t.rawLog != nil || t.tbLog |
| 551 | } | 554 | } |
| 552 | 555 | ||
| 553 | // Context returns a context.Context associated with the test. | 556 | // Context returns a context.Context that is canceled |
| 554 | // It is valid only for the duration of the rapid check. | 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 | func (t *T) Context() context.Context { | 567 | func (t *T) Context() context.Context { |
| 556 | // Fast path: no need to lock if the context is already set. | 568 | // Fast path: no need to lock if the context is already set. |
| 557 | t.mu.RLock() | 569 | t.mu.RLock() |
| @@ -561,6 +573,15 @@ func (t *T) Context() context.Context { · 3 unmodified lines | |||
| 561 | return ctx | 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 | // Slow path: lock and check again, create new context if needed. | 585 | // Slow path: lock and check again, create new context if needed. |
| 565 | t.mu.Lock() | 586 | t.mu.Lock() |
| 566 | defer t.mu.Unlock() | 587 | defer t.mu.Unlock() |
| @@ -587,17 +608,71 @@ func (t *T) Context() context.Context { · 20 unmodified lines | |||
| 587 | return ctx | 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 | // cleanup runs any cleanup tasks associated with the property check. | 633 | // cleanup runs any cleanup tasks associated with the property check. |
| 591 | // It is safe to call multiple times. | 634 | // It is safe to call multiple times. |
| 592 | func (t *T) cleanup() { | 635 | func (t *T) cleanup() { |
| 593 | t.mu.Lock() | 636 | t.cleaning.Store(true) |
| 594 | defer t.mu.Unlock() | 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 | if t.cancelCtx != nil { | 653 | if t.cancelCtx != nil { |
| 597 | t.cancelCtx() | 654 | t.cancelCtx() |
| 598 | t.cancelCtx = nil | 655 | t.cancelCtx = nil |
| 599 | t.ctx = nil | 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 | func (t *T) Logf(format string, args ...any) { | 678 | func (t *T) Logf(format string, args ...any) { |
| @@ -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 | } | ||
| @@ -9,6 +9,7 @@ package rapid · 8 unmodified lines | |||
| 9 | import ( | 9 | import ( |
| 10 | "context" | 10 | "context" |
| 11 | "errors" | 11 | "errors" |
| 12 | "reflect" | ||
| 12 | "strings" | 13 | "strings" |
| 13 | "testing" | 14 | "testing" |
| 14 | ) | 15 | ) |
| @@ -96,6 +97,8 @@ func BenchmarkCheckOverhead(b *testing.B) { · 81 unmodified lines | |||
| 96 | } | 97 | } |
| 97 | 98 | ||
| 98 | func TestCheckContext(t *testing.T) { | 99 | func TestCheckContext(t *testing.T) { |
| 100 | t.Parallel() | ||
| 101 | |||
| 99 | type key struct{} | 102 | type key struct{} |
| 100 | 103 | ||
| 101 | var ctx context.Context | 104 | var ctx context.Context |
| @@ -114,3 +117,144 @@ func TestCheckContext(t *testing.T) { · 12 unmodified lines | |||
| 114 | t.Fatalf("context must have a value") | 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() {} | ||
| @@ -46,7 +46,7 @@ func TestExampleHelper(t *testing.T) { · 45 unmodified lines | |||
| 46 | g.Example(0) | 46 | g.Example(0) |
| 47 | } | 47 | } |
| 48 | 48 | ||
| 49 | func TestExampleContext(t *testing.T) { | 49 | func TestCustomExampleContext(t *testing.T) { |
| 50 | type key struct{} | 50 | type key struct{} |
| 51 | 51 | ||
| 52 | g := Custom(func(t *T) context.Context { | 52 | g := Custom(func(t *T) context.Context { |
| @@ -67,3 +67,17 @@ func TestExampleContext(t *testing.T) { · 14 unmodified lines | |||
| 67 | t.Fatalf("context must have a value") | 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 | } | ||