proxy-pool/internal/controller/provider/upstream_runtime_test.go

package provider

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

	"proxy-pool/internal/controller/pool"
	"proxy-pool/internal/domain/activitypool"
	proxyDomain "proxy-pool/internal/domain/proxy"
)

func TestUpstreamRuntimeReadsInventoryAndFetchesOnlyInsideLeaderTerm(t *testing.T) {
	var inventoryReads atomic.Int64
	fetched := make(chan struct{}, 1)
	runtime, err := NewUpstreamRuntime(UpstreamRuntimeConfig{
		Provider:          runtimeProviderConfig("provider-a"),
		ReconcilePolicy:   runtimeReconcilePolicy(),
		ReconcileInterval: 10 * time.Millisecond,
	}, UpstreamRuntimeDependencies{
		Coordinator: coordinatorFunc(func(
			ctx context.Context,
			upstreamID string,
			limits CoordinationLimits,
			work func(context.Context, LeaderSession) error,
		) error {
			if upstreamID != "provider-a" || limits.MaxTotal != 10 || limits.MaxInFlight != 1 {
				t.Errorf("coordination = (%q, %+v)", upstreamID, limits)
			}
			return work(ctx, unlimitedLeaderSession{})
		}),
		Inventory: inventoryReaderFunc(func(context.Context, string, time.Duration) (pool.InventorySnapshot, error) {
			inventoryReads.Add(1)
			return pool.InventorySnapshot{}, nil
		}),
		Adapter: adapterFunc(func(context.Context) (FetchResponse, error) {
			fetched <- struct{}{}
			return FetchResponse{Body: []byte("fixture")}, nil
		}),
		Parser: parserFunc(func(context.Context, []byte) ([]proxyDomain.Proxy, error) {
			return []proxyDomain.Proxy{{ID: "proxy-1"}}, nil
		}),
		Activity: activitySinkFunc(func(context.Context, string, activitypool.FetchedBatch) (activitypool.UpsertResult, error) {
			return activitypool.UpsertResult{Accepted: 1, Inserted: 1}, nil
		}),
		Results: resultRecorderFunc(func(Result) {}),
	})
	if err != nil {
		t.Fatalf("NewUpstreamRuntime(): %v", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() { done <- runtime.Run(ctx) }()
	select {
	case <-fetched:
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for leader fetch")
	}
	cancel()
	if err := <-done; !errors.Is(err, context.Canceled) {
		t.Fatalf("Run() error = %v, want context cancellation", err)
	}
	if got := inventoryReads.Load(); got == 0 {
		t.Fatal("inventory was not read inside leader term")
	}
}

func TestUpstreamRuntimeFailsClosedUntilInventoryReadRecovers(t *testing.T) {
	var inventoryReads atomic.Int64
	var providerCalls atomic.Int64
	fetched := make(chan struct{}, 1)
	runtime, err := NewUpstreamRuntime(UpstreamRuntimeConfig{
		Provider:          runtimeProviderConfig("provider-a"),
		ReconcilePolicy:   runtimeReconcilePolicy(),
		ReconcileInterval: 10 * time.Millisecond,
	}, UpstreamRuntimeDependencies{
		Coordinator: coordinatorFunc(func(
			ctx context.Context,
			_ string,
			_ CoordinationLimits,
			work func(context.Context, LeaderSession) error,
		) error {
			return work(ctx, unlimitedLeaderSession{})
		}),
		Inventory: inventoryReaderFunc(func(context.Context, string, time.Duration) (pool.InventorySnapshot, error) {
			if inventoryReads.Add(1) == 1 {
				return pool.InventorySnapshot{}, errors.New("inventory unavailable")
			}
			return pool.InventorySnapshot{}, nil
		}),
		Adapter: adapterFunc(func(context.Context) (FetchResponse, error) {
			providerCalls.Add(1)
			fetched <- struct{}{}
			return FetchResponse{Body: []byte("fixture")}, nil
		}),
		Parser: parserFunc(func(context.Context, []byte) ([]proxyDomain.Proxy, error) {
			return []proxyDomain.Proxy{{ID: "proxy-1"}}, nil
		}),
		Activity: activitySinkFunc(func(context.Context, string, activitypool.FetchedBatch) (activitypool.UpsertResult, error) {
			return activitypool.UpsertResult{Accepted: 1, Inserted: 1}, nil
		}),
		Results: resultRecorderFunc(func(Result) {}),
	})
	if err != nil {
		t.Fatalf("NewUpstreamRuntime(): %v", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() { done <- runtime.Run(ctx) }()
	select {
	case <-fetched:
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for recovered inventory fetch")
	}
	cancel()
	<-done
	if got := inventoryReads.Load(); got < 2 {
		t.Fatalf("inventory reads = %d, want recovery retry", got)
	}
	if got := providerCalls.Load(); got != 1 {
		t.Fatalf("provider calls = %d, want one call after recovery", got)
	}
}

func TestUpstreamRuntimeDelegatesRequestIntervalOnlyToCoordinator(t *testing.T) {
	providerConfig := runtimeProviderConfig("provider-a")
	providerConfig.RequestInterval = 500 * time.Millisecond
	fetched := make(chan struct{}, 2)
	runtime, err := NewUpstreamRuntime(UpstreamRuntimeConfig{
		Provider:          providerConfig,
		ReconcilePolicy:   runtimeReconcilePolicy(),
		ReconcileInterval: 5 * time.Millisecond,
	}, UpstreamRuntimeDependencies{
		Coordinator: coordinatorFunc(func(
			ctx context.Context,
			_ string,
			limits CoordinationLimits,
			work func(context.Context, LeaderSession) error,
		) error {
			if limits.RequestInterval != 500*time.Millisecond {
				t.Errorf("distributed RequestInterval = %s, want 500ms", limits.RequestInterval)
			}
			return work(ctx, unlimitedLeaderSession{})
		}),
		Inventory: inventoryReaderFunc(func(context.Context, string, time.Duration) (pool.InventorySnapshot, error) {
			return pool.InventorySnapshot{}, nil
		}),
		Adapter: adapterFunc(func(context.Context) (FetchResponse, error) {
			fetched <- struct{}{}
			return FetchResponse{Body: []byte("fixture")}, nil
		}),
		Parser: parserFunc(func(context.Context, []byte) ([]proxyDomain.Proxy, error) {
			return []proxyDomain.Proxy{{ID: "proxy-1"}}, nil
		}),
		Activity: activitySinkFunc(func(context.Context, string, activitypool.FetchedBatch) (activitypool.UpsertResult, error) {
			return activitypool.UpsertResult{Accepted: 1, Inserted: 1}, nil
		}),
		Results: resultRecorderFunc(func(Result) {}),
	})
	if err != nil {
		t.Fatalf("NewUpstreamRuntime(): %v", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() { done <- runtime.Run(ctx) }()
	for range 2 {
		select {
		case <-fetched:
		case <-time.After(150 * time.Millisecond):
			cancel()
			<-done
			t.Fatal("local Provider reconciler duplicated the distributed request interval")
		}
	}
	cancel()
	<-done
}

func TestNewUpstreamRuntimeRejectsInvalidDependencies(t *testing.T) {
	validConfig := UpstreamRuntimeConfig{
		Provider:          runtimeProviderConfig("provider-a"),
		ReconcilePolicy:   runtimeReconcilePolicy(),
		ReconcileInterval: time.Second,
	}
	if runtime, err := NewUpstreamRuntime(validConfig, UpstreamRuntimeDependencies{}); err == nil || runtime != nil {
		t.Fatalf("NewUpstreamRuntime() = (%v, %v), want invalid dependencies", runtime, err)
	}
}

func TestInitialReconcileDelayIsStableAndBounded(t *testing.T) {
	const interval = time.Second
	first := initialReconcileDelay("provider-a", interval)
	if first != initialReconcileDelay("provider-a", interval) {
		t.Fatal("initial reconcile delay is not stable")
	}
	if first < 0 || first >= 250*time.Millisecond {
		t.Fatalf("initial reconcile delay = %s, want [0, 250ms)", first)
	}
	if other := initialReconcileDelay("provider-b", interval); other == first {
		t.Fatalf("different upstreams have the same initial delay %s", first)
	}
}

func runtimeProviderConfig(upstreamID string) Config {
	return Config{
		UpstreamID: upstreamID, Timeout: time.Second, MaxAttempts: 1,
		MaxInFlight: 1, MaxTotal: 10, MaxSize: 10,
	}
}

func runtimeReconcilePolicy() pool.ReconcilePolicy {
	return pool.ReconcilePolicy{
		MinimumAvailableSlots: 1, TargetAvailableSlots: 2,
		ExpectedPerFetch: 1, ExpectedSlotsPerFetch: 1,
	}
}

type coordinatorFunc func(
	context.Context,
	string,
	CoordinationLimits,
	func(context.Context, LeaderSession) error,
) error

func (f coordinatorFunc) RunLeader(
	ctx context.Context,
	upstreamID string,
	limits CoordinationLimits,
	work func(context.Context, LeaderSession) error,
) error {
	return f(ctx, upstreamID, limits, work)
}

type inventoryReaderFunc func(context.Context, string, time.Duration) (pool.InventorySnapshot, error)

func (f inventoryReaderFunc) ReadInventory(
	ctx context.Context,
	upstreamID string,
	safetyMargin time.Duration,
) (pool.InventorySnapshot, error) {
	return f(ctx, upstreamID, safetyMargin)
}