package ratelimit_test import ( "context" "errors" "fmt" "io" "net/http" "strings" "sync" "testing" "time" "finance-duck/internal/ratelimit" ) func execute(ctx context.Context, gate *ratelimit.Controller, attempt func(context.Context) (*http.Response, error), retry bool) error { if err := gate.Acquire(ctx); err != nil { return err } defer gate.Release() response, err := gate.Do(ctx, attempt, retry) if response != nil { response.Body.Close() } return err } func response(status int) *http.Response { return &http.Response{StatusCode: status, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(""))} } func TestSuccessfulConcurrentCallsArePaced(t *testing.T) { const interval = 30 * time.Millisecond gate := &ratelimit.Controller{MinimumInterval: interval} ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) defer cancel() var mu sync.Mutex var arrivals []time.Time attempt := func(context.Context) (*http.Response, error) { mu.Lock() arrivals = append(arrivals, time.Now()) mu.Unlock() return response(http.StatusOK), nil } done := make(chan error, 4) for range cap(done) { go func() { done <- execute(ctx, gate, attempt, true) }() } for range cap(done) { if err := <-done; err != nil { t.Fatal(err) } } mu.Lock() defer mu.Unlock() if len(arrivals) != 4 { t.Fatalf("got %d outbound calls, want four", len(arrivals)) } for i := 1; i < len(arrivals); i++ { if gap := arrivals[i].Sub(arrivals[i-1]); gap < interval-time.Millisecond { t.Fatalf("successful calls burst: gap %v, minimum %v", gap, interval) } } } func TestRecoveryRetainsLearnedPacing(t *testing.T) { const initial = 20 * time.Millisecond gate := &ratelimit.Controller{InitialBackoff: initial} ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) defer cancel() var arrivals []time.Time attempt := func(context.Context) (*http.Response, error) { arrivals = append(arrivals, time.Now()) if len(arrivals) <= 2 { return response(http.StatusTooManyRequests), nil } return response(http.StatusOK), nil } for range 3 { if err := execute(ctx, gate, attempt, true); err != nil { t.Fatal(err) } } if len(arrivals) != 5 { t.Fatalf("got %d attempts, want two retries and three successes", len(arrivals)) } for i := 2; i < len(arrivals); i++ { if gap := arrivals[i].Sub(arrivals[i-1]); gap < 2*initial-time.Millisecond { t.Fatalf("attempt %d discarded learned cadence after recovery: %v", i+1, gap) } } } func TestExhaustionEscalatesAcrossCalls(t *testing.T) { const initial = 10 * time.Millisecond gate := &ratelimit.Controller{InitialBackoff: initial} ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) defer cancel() var arrivals []time.Time attempt := func(context.Context) (*http.Response, error) { arrivals = append(arrivals, time.Now()) return response(http.StatusTooManyRequests), nil } err := execute(ctx, gate, attempt, true) var exhausted *ratelimit.RateLimitError if !errors.As(err, &exhausted) || len(arrivals) != 4 { t.Fatalf("expected exhaustion after four attempts, got %v and %d", err, len(arrivals)) } for i := 1; i < len(arrivals); i++ { if gap := arrivals[i].Sub(arrivals[i-1]); gap < (initial<<(i-1))-time.Millisecond { t.Fatalf("retry %d did not escalate: %v", i, gap) } } // A new transaction is eligible only after the fourth attempt's cooldown. time.Sleep(time.Until(exhausted.RetryAt())) err = execute(ctx, gate, attempt, false) var next *ratelimit.RateLimitError if !errors.As(err, &next) || len(arrivals) != 5 { t.Fatalf("eligible successor did not make exactly one attempt: %v, %d", err, len(arrivals)) } if delay := next.RetryAt().Sub(arrivals[4]); delay < 16*initial { t.Fatalf("new call reset consecutive-failure backoff: %v, need at least %v", delay, 16*initial) } if err := execute(ctx, gate, attempt, true); !errors.As(err, &next) || len(arrivals) != 5 { t.Fatalf("retained cooldown allowed an early attempt: %v, %d", err, len(arrivals)) } } func TestCanceledPacingWaitMakesNoOutboundCall(t *testing.T) { gate := &ratelimit.Controller{MinimumInterval: 100 * time.Millisecond} ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) defer cancel() calls := 0 attempt := func(context.Context) (*http.Response, error) { calls++ return response(http.StatusOK), nil } if err := execute(ctx, gate, attempt, true); err != nil { t.Fatal(err) } pacedCtx, pacedCancel := context.WithTimeout(ctx, 10*time.Millisecond) err := execute(pacedCtx, gate, attempt, true) pacedCancel() if !errors.Is(err, context.DeadlineExceeded) || calls != 1 { t.Fatalf("paced cancellation sent a request or lost context identity: %v, calls=%d", err, calls) } if err := execute(ctx, gate, attempt, true); err != nil || calls != 2 { t.Fatalf("cancelled waiter leaked an attempt or blocked its successor: %v, calls=%d", err, calls) } } func TestAutomaticWaitBudgetRetainsLongCooldown(t *testing.T) { gate := &ratelimit.Controller{InitialBackoff: 3 * time.Minute} ctx, cancel := context.WithTimeout(context.Background(), time.Second) defer cancel() calls := 0 attempt := func(context.Context) (*http.Response, error) { calls++ return response(http.StatusTooManyRequests), nil } for range 2 { err := execute(ctx, gate, attempt, true) var limit *ratelimit.RateLimitError if !errors.As(err, &limit) || errors.Is(err, context.DeadlineExceeded) || calls != 1 { t.Fatalf("over-budget cooldown waited or replayed: %v, calls=%d", err, calls) } } } func TestProviderSpecificErrorDoesNotPauseUnrelatedCalls(t *testing.T) { for _, deadline := range []time.Time{time.Now().Add(6 * time.Hour), {}} { gate := &ratelimit.Controller{} err := execute(context.Background(), gate, func(context.Context) (*http.Response, error) { return nil, fmt.Errorf("account quota: %w", ratelimit.NewError(deadline)) }, true) var limit *ratelimit.RateLimitError if !errors.As(err, &limit) || !limit.RetryAt().Equal(deadline) { t.Fatalf("provider quota lost its typed retry deadline: %v", err) } calls := 0 err = execute(context.Background(), gate, func(context.Context) (*http.Response, error) { calls++ return response(http.StatusOK), nil }, true) if err != nil || calls != 1 { t.Fatalf("provider-specific quota poisoned common controller: %v, calls=%d", err, calls) } } }