RunRetry runs a test function up to `count` times, retrying if it fails. If all attempts fail or the context is canceled, the test will fail. It is safe to use the parent context in the test function, but do note that the context deadline will apply to all attempts. DO NOT USE THIS FUNCTION IN TEST
(t *testing.T, count int, fn func(t testing.TB))
| 36 | // |
| 37 | // Cleanup functions will be executed after each attempt. |
| 38 | func RunRetry(t *testing.T, count int, fn func(t testing.TB)) { |
| 39 | t.Helper() |
| 40 | |
| 41 | for i := 1; i <= count; i++ { |
| 42 | // Canceled in the attempt goroutine before running cleanup functions. |
| 43 | attemptCtx, attemptCancel := context.WithCancel(t.Context()) |
| 44 | attemptT := &fakeT{ |
| 45 | T: t, |
| 46 | ctx: attemptCtx, |
| 47 | name: fmt.Sprintf("%s (attempt %d/%d)", t.Name(), i, count), |
| 48 | } |
| 49 | |
| 50 | // Run the test in a goroutine so we can capture runtime.Goexit() |
| 51 | // and run cleanup functions. |
| 52 | done := make(chan struct{}, 1) |
| 53 | go func() { |
| 54 | defer close(done) |
| 55 | defer func() { |
| 56 | // As per t.Context(), the context is canceled right before |
| 57 | // cleanup functions are executed. |
| 58 | attemptCancel() |
| 59 | attemptT.runCleanupFns() |
| 60 | }() |
| 61 | |
| 62 | t.Logf("testutil.RunRetry: running test: attempt %d/%d", i, count) |
| 63 | fn(attemptT) |
| 64 | }() |
| 65 | |
| 66 | // We don't wait on the context here, because we want to be sure that |
| 67 | // the test function and cleanup functions have finished before |
| 68 | // returning from the test. |
| 69 | <-done |
| 70 | if !attemptT.Failed() { |
| 71 | t.Logf("testutil.RunRetry: test passed on attempt %d/%d", i, count) |
| 72 | return |
| 73 | } |
| 74 | t.Logf("testutil.RunRetry: test failed on attempt %d/%d", i, count) |
| 75 | |
| 76 | // Wait a few seconds in case the test failure was due to system load. |
| 77 | // There's not really a good way to check for this, so we just do it |
| 78 | // every time. |
| 79 | // No point waiting on t.Context() here because it doesn't factor in |
| 80 | // the test deadline, and only gets canceled when the test function |
| 81 | // completes. |
| 82 | time.Sleep(2 * time.Second) |
| 83 | } |
| 84 | t.Fatalf("testutil.RunRetry: all %d attempts failed", count) |
| 85 | } |
| 86 | |
| 87 | // fakeT is a fake implementation of testing.TB that never fails and only logs |
| 88 | // errors. Fatal errors will cause the goroutine to exit without failing the |