Files
2026-07-13 13:00:08 +08:00

133 lines
3.9 KiB
Go

package bot
import (
"context"
"errors"
"io"
"log/slog"
"sync/atomic"
"testing"
"time"
)
func discardLogger() *slog.Logger {
return slog.New(slog.NewTextHandler(io.Discard, nil))
}
func TestSleepCtxCompletes(t *testing.T) {
if !SleepCtx(context.Background(), time.Millisecond) {
t.Fatal("SleepCtx should return true when the full delay elapses")
}
}
func TestSleepCtxCancelledReturnsPromptly(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
go func() {
time.Sleep(10 * time.Millisecond)
cancel()
}()
start := time.Now()
if SleepCtx(ctx, 10*time.Second) {
t.Fatal("SleepCtx should return false when ctx is cancelled mid-wait")
}
if elapsed := time.Since(start); elapsed > time.Second {
t.Fatalf("SleepCtx ignored cancellation: waited %v", elapsed)
}
}
func TestSleepCtxAlreadyCancelled(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
cancel()
if SleepCtx(ctx, time.Second) {
t.Fatal("SleepCtx should return false immediately when ctx is already cancelled")
}
}
func TestNextDelay(t *testing.T) {
maxD := 30 * time.Second
cases := []struct{ cur, want time.Duration }{
{1 * time.Second, 2 * time.Second},
{8 * time.Second, 16 * time.Second},
{16 * time.Second, 30 * time.Second}, // doubling past max → capped
{30 * time.Second, 30 * time.Second}, // stays at max
}
for _, c := range cases {
if got := nextDelay(c.cur, maxD); got != c.want {
t.Errorf("nextDelay(%v, %v) = %v, want %v", c.cur, maxD, got, c.want)
}
}
}
func TestRetryConfigDefaults(t *testing.T) {
got := RetryConfig{}.withDefaults()
if got.InitialDelay != defaultInitialDelay || got.MaxDelay != defaultMaxDelay || got.ResetAfter != defaultResetAfter {
t.Fatalf("zero RetryConfig defaults = %+v", got)
}
// MaxDelay below InitialDelay is clamped up to InitialDelay.
got = RetryConfig{InitialDelay: 5 * time.Second, MaxDelay: time.Second}.withDefaults()
if got.MaxDelay != 5*time.Second {
t.Fatalf("MaxDelay should clamp up to InitialDelay, got %v", got.MaxDelay)
}
}
func TestRunWithRetryNotCalledWhenContextAlreadyCancelled(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
cancel()
var calls atomic.Int32
RunWithRetry(ctx, discardLogger(), "test", RetryConfig{InitialDelay: time.Millisecond}, func(context.Context) error {
calls.Add(1)
return nil
})
if n := calls.Load(); n != 0 {
t.Fatalf("attempt ran %d times for an already-cancelled ctx, want 0", n)
}
}
func TestRunWithRetryRetriesUntilCancel(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
var calls atomic.Int32
done := make(chan struct{})
go func() {
RunWithRetry(ctx, discardLogger(), "test", RetryConfig{InitialDelay: time.Millisecond, MaxDelay: time.Millisecond}, func(context.Context) error {
// Stop the loop from inside the attempt once we've reconnected 3 times.
if calls.Add(1) >= 3 {
cancel()
}
return errors.New("dropped")
})
close(done)
}()
select {
case <-done:
case <-time.After(5 * time.Second):
t.Fatal("RunWithRetry did not return after ctx cancellation")
}
if n := calls.Load(); n != 3 {
t.Fatalf("attempt ran %d times, want exactly 3", n)
}
}
func TestRunWithRetryCancelDuringBackoffReturnsPromptly(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
done := make(chan struct{})
start := time.Now()
go func() {
// Large backoff: the only way this returns quickly is if the backoff wait
// honors ctx cancellation.
RunWithRetry(ctx, discardLogger(), "test", RetryConfig{InitialDelay: 10 * time.Second, MaxDelay: 10 * time.Second}, func(context.Context) error {
return errors.New("dropped")
})
close(done)
}()
time.Sleep(20 * time.Millisecond)
cancel()
select {
case <-done:
case <-time.After(2 * time.Second):
t.Fatal("RunWithRetry ignored cancellation during backoff")
}
if elapsed := time.Since(start); elapsed > time.Second {
t.Fatalf("RunWithRetry took %v to honor cancellation during backoff", elapsed)
}
}