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

package pool

import (
	"testing"
	"time"

	proxyDomain "proxy-pool/internal/domain/proxy"
	"proxy-pool/internal/domain/upstream"
)

func TestPoolReconcilerSignalsOnlyWhenSlotsAreLowAndWholeFetchFits(t *testing.T) {
	budget, err := NewFetchBudget(FetchBudgetConfig{
		UpstreamID: "provider-a", MaxSize: 10, MaxTotal: 20, ExpectedPerFetch: 2,
	})
	if err != nil {
		t.Fatalf("NewFetchBudget(): %v", err)
	}
	notifier := &recordingFetchNotifier{}
	reconciler, err := NewReconciler(ReconcilePolicy{
		MinimumAvailableSlots: 5,
		TargetAvailableSlots:  8,
		ExpectedPerFetch:      2,
		ExpectedSlotsPerFetch: 2,
		SafetyMargin:          10 * time.Second,
	}, budget, notifier)
	if err != nil {
		t.Fatalf("NewReconciler(): %v", err)
	}
	now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
	inventory := upstream.Inventory{
		Proxies: []upstream.ProxyCapacity{{
			State: proxyDomain.StateAvailable, ExpiresAt: now.Add(time.Minute), Max: 4, Active: 3,
		}},
		MaxSize:  10,
		MaxTotal: 20,
	}

	decision := reconciler.Reconcile(now, inventory)
	if !decision.Triggered || decision.AvailableSlots != 1 || decision.FetchAllowance != 2 {
		t.Fatalf("Reconcile() = %+v, want triggered with one available slot", decision)
	}
	if notifier.calls != 1 {
		t.Fatalf("Notify() calls = %d, want 1", notifier.calls)
	}
}

func TestPoolReconcilerUsesBudgetPendingAndCumulativeCounters(t *testing.T) {
	budget, err := NewFetchBudget(FetchBudgetConfig{
		UpstreamID: "provider-a", MaxSize: 2, MaxTotal: 2, ExpectedPerFetch: 2,
	})
	if err != nil {
		t.Fatalf("NewFetchBudget(): %v", err)
	}
	permit, ok, err := budget.ReserveFetch("provider-a")
	if err != nil || !ok {
		t.Fatalf("ReserveFetch() = (_, %v, %v), want permit", ok, err)
	}
	defer permit.Cancel()
	notifier := &recordingFetchNotifier{}
	reconciler, err := NewReconciler(ReconcilePolicy{
		MinimumAvailableSlots: 1, TargetAvailableSlots: 2,
		ExpectedPerFetch: 2, ExpectedSlotsPerFetch: 2,
	}, budget, notifier)
	if err != nil {
		t.Fatalf("NewReconciler(): %v", err)
	}

	decision := reconciler.Reconcile(time.Now(), upstream.Inventory{MaxSize: 2, MaxTotal: 2})
	if decision.Triggered || decision.PendingExpected != 2 || decision.FetchAllowance != 0 {
		t.Fatalf("Reconcile() = %+v, want pending fetch to suppress signal", decision)
	}
	if notifier.calls != 0 {
		t.Fatalf("Notify() calls = %d, want 0", notifier.calls)
	}
}

func TestPoolReconcilerUsesTargetWatermarkUntilRefillCompletes(t *testing.T) {
	budget, err := NewFetchBudget(FetchBudgetConfig{
		UpstreamID: "provider-a", MaxSize: 10, ExpectedPerFetch: 2,
	})
	if err != nil {
		t.Fatalf("NewFetchBudget(): %v", err)
	}
	notifier := &recordingFetchNotifier{}
	reconciler, err := NewReconciler(ReconcilePolicy{
		MinimumAvailableSlots: 3,
		TargetAvailableSlots:  8,
		ExpectedPerFetch:      2,
		ExpectedSlotsPerFetch: 2,
	}, budget, notifier)
	if err != nil {
		t.Fatalf("NewReconciler(): %v", err)
	}
	now := time.Date(2026, 7, 30, 10, 0, 0, 0, time.UTC)

	belowMinimum := reconciler.Reconcile(now, inventoryWithSlots(now, 2))
	betweenWatermarks := reconciler.Reconcile(now, inventoryWithSlots(now, 5))
	atTarget := reconciler.Reconcile(now, inventoryWithSlots(now, 8))
	aboveMinimumAfterCompletion := reconciler.Reconcile(now, inventoryWithSlots(now, 5))

	if !belowMinimum.Triggered || !betweenWatermarks.Triggered || atTarget.Triggered || aboveMinimumAfterCompletion.Triggered {
		t.Fatalf("triggered states = [%v %v %v %v], want [true true false false]",
			belowMinimum.Triggered, betweenWatermarks.Triggered, atTarget.Triggered, aboveMinimumAfterCompletion.Triggered)
	}
	if notifier.calls != 2 {
		t.Fatalf("Notify() calls = %d, want 2", notifier.calls)
	}
}

func TestPoolReconcilerPendingEstimatePausesWithoutEndingRefillEpisode(t *testing.T) {
	budget, err := NewFetchBudget(FetchBudgetConfig{
		UpstreamID: "provider-a", MaxSize: 10, ExpectedPerFetch: 2,
	})
	if err != nil {
		t.Fatalf("NewFetchBudget(): %v", err)
	}
	notifier := &recordingFetchNotifier{}
	reconciler, err := NewReconciler(ReconcilePolicy{
		MinimumAvailableSlots: 3, TargetAvailableSlots: 8,
		ExpectedPerFetch: 2, ExpectedSlotsPerFetch: 6,
	}, budget, notifier)
	if err != nil {
		t.Fatalf("NewReconciler(): %v", err)
	}
	now := time.Date(2026, 7, 30, 10, 0, 0, 0, time.UTC)
	if decision := reconciler.Reconcile(now, inventoryWithSlots(now, 2)); !decision.Triggered {
		t.Fatalf("initial Reconcile() = %+v, want trigger", decision)
	}
	permit, ok, err := budget.ReserveFetch("provider-a")
	if err != nil || !ok {
		t.Fatalf("ReserveFetch() = (_, %v, %v), want permit", ok, err)
	}
	if decision := reconciler.Reconcile(now, inventoryWithSlots(now, 2)); decision.Triggered || decision.EffectiveSlots != 8 {
		t.Fatalf("pending Reconcile() = %+v, want paused at target estimate", decision)
	}
	if err := permit.Complete(2, 1); err != nil {
		t.Fatalf("Complete(): %v", err)
	}
	if decision := reconciler.Reconcile(now, inventoryWithSlots(now, 5)); !decision.Triggered {
		t.Fatalf("post-fetch Reconcile() = %+v, want refill episode to continue", decision)
	}
}

func inventoryWithSlots(now time.Time, slots int64) upstream.Inventory {
	return upstream.Inventory{Proxies: []upstream.ProxyCapacity{{
		State: proxyDomain.StateAvailable, ExpiresAt: now.Add(time.Minute), Max: slots,
	}}}
}

type recordingFetchNotifier struct{ calls int }

func (n *recordingFetchNotifier) Notify() { n.calls++ }