finance-duck/internal/ratelimit/controller_test.go

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)
		}
	}
}