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