| 1 | package bot |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "errors" |
| 6 | "io" |
| 7 | "log/slog" |
| 8 | "sync/atomic" |
| 9 | "testing" |
| 10 | "time" |
| 11 | ) |
| 12 | |
| 13 | func discardLogger() *slog.Logger { |
| 14 | return slog.New(slog.NewTextHandler(io.Discard, nil)) |
| 15 | } |
| 16 | |
| 17 | func TestSleepCtxCompletes(t *testing.T) { |
| 18 | if !SleepCtx(context.Background(), time.Millisecond) { |
| 19 | t.Fatal("SleepCtx should return true when the full delay elapses") |
| 20 | } |
| 21 | } |
| 22 | |
| 23 | func TestSleepCtxCancelledReturnsPromptly(t *testing.T) { |
| 24 | ctx, cancel := context.WithCancel(context.Background()) |
| 25 | go func() { |
| 26 | time.Sleep(10 * time.Millisecond) |
| 27 | cancel() |
| 28 | }() |
| 29 | start := time.Now() |
| 30 | if SleepCtx(ctx, 10*time.Second) { |
| 31 | t.Fatal("SleepCtx should return false when ctx is cancelled mid-wait") |
| 32 | } |
| 33 | if elapsed := time.Since(start); elapsed > time.Second { |
| 34 | t.Fatalf("SleepCtx ignored cancellation: waited %v", elapsed) |
| 35 | } |
| 36 | } |
| 37 | |
| 38 | func TestSleepCtxAlreadyCancelled(t *testing.T) { |
| 39 | ctx, cancel := context.WithCancel(context.Background()) |
| 40 | cancel() |
| 41 | if SleepCtx(ctx, time.Second) { |
| 42 | t.Fatal("SleepCtx should return false immediately when ctx is already cancelled") |
| 43 | } |
| 44 | } |
| 45 | |
| 46 | func TestNextDelay(t *testing.T) { |
| 47 | maxD := 30 * time.Second |
| 48 | cases := []struct{ cur, want time.Duration }{ |
| 49 | {1 * time.Second, 2 * time.Second}, |
| 50 | {8 * time.Second, 16 * time.Second}, |
| 51 | {16 * time.Second, 30 * time.Second}, // doubling past max → capped |
| 52 | {30 * time.Second, 30 * time.Second}, // stays at max |
| 53 | } |
| 54 | for _, c := range cases { |
| 55 | if got := nextDelay(c.cur, maxD); got != c.want { |
| 56 | t.Errorf("nextDelay(%v, %v) = %v, want %v", c.cur, maxD, got, c.want) |
| 57 | } |
| 58 | } |
| 59 | } |
| 60 | |
| 61 | func TestRetryConfigDefaults(t *testing.T) { |
| 62 | got := RetryConfig{}.withDefaults() |
| 63 | if got.InitialDelay != defaultInitialDelay || got.MaxDelay != defaultMaxDelay || got.ResetAfter != defaultResetAfter { |
| 64 | t.Fatalf("zero RetryConfig defaults = %+v", got) |
| 65 | } |
| 66 | // MaxDelay below InitialDelay is clamped up to InitialDelay. |
| 67 | got = RetryConfig{InitialDelay: 5 * time.Second, MaxDelay: time.Second}.withDefaults() |
| 68 | if got.MaxDelay != 5*time.Second { |
| 69 | t.Fatalf("MaxDelay should clamp up to InitialDelay, got %v", got.MaxDelay) |
| 70 | } |
| 71 | } |
| 72 | |
| 73 | func TestRunWithRetryNotCalledWhenContextAlreadyCancelled(t *testing.T) { |
| 74 | ctx, cancel := context.WithCancel(context.Background()) |
| 75 | cancel() |
| 76 | var calls atomic.Int32 |
| 77 | RunWithRetry(ctx, discardLogger(), "test", RetryConfig{InitialDelay: time.Millisecond}, func(context.Context) error { |
| 78 | calls.Add(1) |
| 79 | return nil |
| 80 | }) |
| 81 | if n := calls.Load(); n != 0 { |
| 82 | t.Fatalf("attempt ran %d times for an already-cancelled ctx, want 0", n) |
| 83 | } |
| 84 | } |
| 85 | |
| 86 | func TestRunWithRetryRetriesUntilCancel(t *testing.T) { |
| 87 | ctx, cancel := context.WithCancel(context.Background()) |
| 88 | var calls atomic.Int32 |
| 89 | done := make(chan struct{}) |
| 90 | go func() { |
| 91 | RunWithRetry(ctx, discardLogger(), "test", RetryConfig{InitialDelay: time.Millisecond, MaxDelay: time.Millisecond}, func(context.Context) error { |
| 92 | // Stop the loop from inside the attempt once we've reconnected 3 times. |
| 93 | if calls.Add(1) >= 3 { |
| 94 | cancel() |
| 95 | } |
| 96 | return errors.New("dropped") |
| 97 | }) |
| 98 | close(done) |
| 99 | }() |
| 100 | select { |
| 101 | case <-done: |
| 102 | case <-time.After(5 * time.Second): |
| 103 | t.Fatal("RunWithRetry did not return after ctx cancellation") |
| 104 | } |
| 105 | if n := calls.Load(); n != 3 { |
| 106 | t.Fatalf("attempt ran %d times, want exactly 3", n) |
| 107 | } |
| 108 | } |
| 109 | |
| 110 | func TestRunWithRetryCancelDuringBackoffReturnsPromptly(t *testing.T) { |
| 111 | ctx, cancel := context.WithCancel(context.Background()) |
| 112 | done := make(chan struct{}) |
| 113 | start := time.Now() |
| 114 | go func() { |
| 115 | // Large backoff: the only way this returns quickly is if the backoff wait |
| 116 | // honors ctx cancellation. |
| 117 | RunWithRetry(ctx, discardLogger(), "test", RetryConfig{InitialDelay: 10 * time.Second, MaxDelay: 10 * time.Second}, func(context.Context) error { |
| 118 | return errors.New("dropped") |
| 119 | }) |
| 120 | close(done) |
| 121 | }() |
| 122 | time.Sleep(20 * time.Millisecond) |
| 123 | cancel() |
| 124 | select { |
| 125 | case <-done: |
| 126 | case <-time.After(2 * time.Second): |
| 127 | t.Fatal("RunWithRetry ignored cancellation during backoff") |
| 128 | } |
| 129 | if elapsed := time.Since(start); elapsed > time.Second { |
| 130 | t.Fatalf("RunWithRetry took %v to honor cancellation during backoff", elapsed) |
| 131 | } |
| 132 | } |
| 133 |