proxy-pool/internal/controller/pool/fetch_budget_test.go

package pool

import (
	"errors"
	"sync"
	"testing"
)

func TestFetchBudgetReservesExpectedCapacityAndSeparatesCounters(t *testing.T) {
	budget, err := NewFetchBudget(FetchBudgetConfig{
		UpstreamID:       "provider-a",
		MaxSize:          10,
		MaxTotal:         20,
		ExpectedPerFetch: 4,
		Managed:          2,
		FetchedTotal:     3,
	})
	if err != nil {
		t.Fatalf("NewFetchBudget(): %v", err)
	}

	first, ok, err := budget.ReserveFetch("provider-a")
	if err != nil || !ok {
		t.Fatalf("first ReserveFetch() = (_, %v, %v), want permit", ok, err)
	}
	second, ok, err := budget.ReserveFetch("provider-a")
	if err != nil || !ok {
		t.Fatalf("second ReserveFetch() = (_, %v, %v), want permit", ok, err)
	}
	if _, ok, err := budget.ReserveFetch("provider-a"); err != nil || ok {
		t.Fatalf("third ReserveFetch() = (_, %v, %v), want no capacity", ok, err)
	}

	if err := first.Complete(4, 3); err != nil {
		t.Fatalf("first.Complete(): %v", err)
	}
	if err := second.Cancel(); err != nil {
		t.Fatalf("second.Cancel(): %v", err)
	}
	usage := budget.Snapshot()
	if usage.Managed != 5 || usage.PendingExpected != 0 || usage.FetchedTotal != 7 {
		t.Fatalf("Snapshot() = %+v, want managed=5 pending=0 fetched=7", usage)
	}
	if err := budget.ReleaseManaged(2); err != nil {
		t.Fatalf("ReleaseManaged(): %v", err)
	}
	if usage := budget.Snapshot(); usage.Managed != 3 || usage.FetchedTotal != 7 {
		t.Fatalf("Snapshot() after release = %+v, want managed=3 fetched=7", usage)
	}
}

func TestFetchBudgetRejectsManagedCounterUnderflow(t *testing.T) {
	budget, err := NewFetchBudget(FetchBudgetConfig{
		UpstreamID: "a", MaxSize: 10, ExpectedPerFetch: 1, Managed: 1,
	})
	if err != nil {
		t.Fatalf("NewFetchBudget(): %v", err)
	}
	if err := budget.ReleaseManaged(2); !errors.Is(err, ErrInvalidManagedRelease) {
		t.Fatalf("ReleaseManaged() error = %v, want ErrInvalidManagedRelease", err)
	}
}

func TestFetchBudgetRequiresWholeExpectedBatchToFitLimits(t *testing.T) {
	tests := []struct {
		name   string
		config FetchBudgetConfig
	}{
		{
			name: "pool size",
			config: FetchBudgetConfig{
				UpstreamID: "a", MaxSize: 10, ExpectedPerFetch: 4, Managed: 7,
			},
		},
		{
			name: "cumulative total",
			config: FetchBudgetConfig{
				UpstreamID: "a", MaxSize: 10, MaxTotal: 5, ExpectedPerFetch: 4, FetchedTotal: 2,
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			budget, err := NewFetchBudget(tt.config)
			if err != nil {
				t.Fatalf("NewFetchBudget(): %v", err)
			}
			if _, ok, err := budget.ReserveFetch("a"); err != nil || ok {
				t.Fatalf("ReserveFetch() = (_, %v, %v), want no capacity", ok, err)
			}
		})
	}
}

func TestFetchBudgetConcurrentReservationsNeverExceedMaxSize(t *testing.T) {
	budget, err := NewFetchBudget(FetchBudgetConfig{
		UpstreamID: "a", MaxSize: 10, ExpectedPerFetch: 1,
	})
	if err != nil {
		t.Fatalf("NewFetchBudget(): %v", err)
	}

	var wg sync.WaitGroup
	permits := make(chan interface{ Cancel() error }, 100)
	for range 100 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			permit, ok, reserveErr := budget.ReserveFetch("a")
			if reserveErr != nil {
				t.Errorf("ReserveFetch(): %v", reserveErr)
				return
			}
			if ok {
				permits <- permit
			}
		}()
	}
	wg.Wait()
	close(permits)

	count := 0
	for permit := range permits {
		count++
		if err := permit.Cancel(); err != nil {
			t.Errorf("Cancel(): %v", err)
		}
	}
	if count != 10 {
		t.Fatalf("reserved permits = %d, want 10", count)
	}
}

func TestFetchPermitRejectsDoubleFinishAndInvalidCounts(t *testing.T) {
	budget, err := NewFetchBudget(FetchBudgetConfig{
		UpstreamID: "a", MaxSize: 10, ExpectedPerFetch: 2,
	})
	if err != nil {
		t.Fatalf("NewFetchBudget(): %v", err)
	}
	permit, ok, err := budget.ReserveFetch("a")
	if err != nil || !ok {
		t.Fatalf("ReserveFetch() = (_, %v, %v), want permit", ok, err)
	}
	if err := permit.Complete(1, 2); !errors.Is(err, ErrInvalidFetchCompletion) {
		t.Fatalf("Complete() error = %v, want ErrInvalidFetchCompletion", err)
	}
	if err := permit.Cancel(); err != nil {
		t.Fatalf("Cancel(): %v", err)
	}
	if err := permit.Cancel(); !errors.Is(err, ErrFetchPermitFinished) {
		t.Fatalf("second Cancel() error = %v, want ErrFetchPermitFinished", err)
	}
}