package main

import (
	"context"
	"errors"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

func TestBudgetConcurrent(t *testing.T) {
	b := NewBudget(5)
	var got atomic.Int64
	var wg sync.WaitGroup
	for range 100 {
		wg.Go(func() {
			if b.Take() {
				got.Add(1)
			}
		})
	}
	wg.Wait()
	if got.Load() != 5 {
		t.Fatal(got.Load())
	}
}
func TestCancellationAndPermanent(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	calls := 0
	f := func(context.Context) error { calls++; return ErrTransient }
	err := Retry(ctx, 3, NewBudget(2), func(int) time.Duration { return 0 }, f)
	if !errors.Is(err, context.Canceled) || calls != 0 {
		t.Fatal(err, calls)
	}
	p := errors.New("invalid argument")
	err = Retry(context.Background(), 3, NewBudget(2), func(int) time.Duration { return 0 }, func(context.Context) error { calls++; return p })
	if err != p || calls != 1 {
		t.Fatal(err, calls)
	}
}
func TestCancelDuringWait(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	calls := 0
	err := Retry(ctx, 3, NewBudget(2), func(int) time.Duration { cancel(); return time.Hour }, func(context.Context) error { calls++; return ErrTransient })
	if !errors.Is(err, context.Canceled) || calls != 1 {
		t.Fatal(err, calls)
	}
}
func TestPanicReleasesSlot(t *testing.T) {
	b := NewBulkhead(1)
	func() {
		defer func() {
			if recover() == nil {
				t.Fatal("missing panic")
			}
		}()
		_ = b.Try(context.Background(), func(context.Context) error { panic("test") })
	}()
	if err := b.Try(context.Background(), func(context.Context) error { return nil }); err != nil {
		t.Fatal(err)
	}
}
func TestJitterAndCap(t *testing.T) {
	for i := 0; i < 1000; i++ {
		d := Backoff(time.Millisecond, time.Second, i)
		j := FullJitter(d)
		if d > time.Second || j < 0 || j >= d {
			t.Fatal(d, j)
		}
	}
}
