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

package pool

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

	extractionDomain "github.com/proxy-pool/proxy-pool/internal/domain/extraction"
)

func TestOwnershipManagerPreventsDualAssignment(t *testing.T) {
	manager := newTestOwnershipManager(t, "proxy-1")
	now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
	var succeeded atomic.Int64
	var wg sync.WaitGroup

	for index := range 100 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			_, err := manager.Assign(now, "proxy-1", fmt.Sprintf("worker-%d", index), time.Minute)
			if err == nil {
				succeeded.Add(1)
				return
			}
			if !errors.Is(err, ErrAlreadyOwned) {
				t.Errorf("Assign(): %v", err)
			}
		}()
	}
	wg.Wait()

	if got := succeeded.Load(); got != 1 {
		t.Fatalf("successful assignments = %d, want 1", got)
	}
}

func TestOwnershipManagerRenewsOnlyCurrentAssignment(t *testing.T) {
	manager := newTestOwnershipManager(t, "proxy-1")
	now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
	assigned, err := manager.Assign(now, "proxy-1", "worker-1", time.Minute)
	if err != nil {
		t.Fatalf("Assign(): %v", err)
	}
	renewed, err := manager.Renew(now.Add(30*time.Second), "proxy-1", "worker-1", assigned.Epoch, time.Minute)
	if err != nil {
		t.Fatalf("Renew(): %v", err)
	}
	if renewed.Version != assigned.Version+1 || !renewed.ExpiresAt.Equal(now.Add(90*time.Second)) {
		t.Fatalf("renewed assignment = %+v", renewed)
	}
	if _, err := manager.Renew(now, "proxy-1", "worker-2", assigned.Epoch, time.Minute); !errors.Is(err, ErrStaleAssignment) {
		t.Fatalf("Renew(stale) error = %v, want ErrStaleAssignment", err)
	}
	if expired := manager.Expire(now.Add(time.Minute)); len(expired) != 0 {
		t.Fatalf("renewed assignment expired at old deadline: %+v", expired)
	}
}

func TestSharedRepositoryMakesOwnershipAndExtractionMutuallyExclusive(t *testing.T) {
	for iteration := range 100 {
		now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
		store := extractionDomain.NewMemoryStore([]extractionDomain.Candidate{{
			ID: "proxy-1", State: extractionDomain.Available,
		}})
		manager, err := NewOwnershipManager(store)
		if err != nil {
			t.Fatalf("iteration %d NewOwnershipManager(): %v", iteration, err)
		}

		start := make(chan struct{})
		assigned := make(chan bool, 1)
		extracted := make(chan bool, 1)
		go func() {
			<-start
			_, assignErr := manager.Assign(now, "proxy-1", "worker-1", time.Minute)
			if assignErr != nil && !errors.Is(assignErr, ErrOwnershipUnavailable) {
				t.Errorf("iteration %d Assign(): %v", iteration, assignErr)
			}
			assigned <- assignErr == nil
		}()
		go func() {
			<-start
			result, extractErr := store.Extract(context.Background(), extractionDomain.Command{
				ClientID: "client-1", Requested: 1, Fulfillment: extractionDomain.Partial, Now: now,
			})
			if extractErr != nil {
				t.Errorf("iteration %d Extract(): %v", iteration, extractErr)
			}
			extracted <- result.Returned == 1
		}()
		close(start)

		wins := 0
		if <-assigned {
			wins++
		}
		if <-extracted {
			wins++
		}
		if wins != 1 {
			t.Fatalf("iteration %d successful ownership/extraction operations = %d, want 1", iteration, wins)
		}
	}
}

func TestOwnershipManagerRequiresDrainAckAtZeroRuntime(t *testing.T) {
	manager := newTestOwnershipManager(t, "proxy-1")
	now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
	assignment, err := manager.Assign(now, "proxy-1", "worker-1", time.Minute)
	if err != nil {
		t.Fatalf("Assign(): %v", err)
	}
	draining, err := manager.BeginDrain("proxy-1", "worker-1", assignment.Epoch)
	if err != nil {
		t.Fatalf("BeginDrain(): %v", err)
	}
	if !draining.Draining || draining.Version != assignment.Version+1 {
		t.Fatalf("draining assignment = %+v", draining)
	}

	if err := manager.AcknowledgeDrain("proxy-1", "worker-1", assignment.Epoch, 1, 0); !errors.Is(err, ErrDrainNotReady) {
		t.Fatalf("AcknowledgeDrain(active) error = %v, want ErrDrainNotReady", err)
	}
	if err := manager.AcknowledgeDrain("proxy-1", "worker-1", assignment.Epoch, 0, 0); err != nil {
		t.Fatalf("AcknowledgeDrain(zero): %v", err)
	}
	if _, ok := manager.Get("proxy-1"); ok {
		t.Fatal("assignment still exists after drain acknowledgement")
	}
}

func TestOwnershipManagerExpiresCrashedWorkerAssignment(t *testing.T) {
	manager := newTestOwnershipManager(t, "proxy-1")
	now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
	first, err := manager.Assign(now, "proxy-1", "worker-1", time.Minute)
	if err != nil {
		t.Fatalf("Assign(first): %v", err)
	}
	if expired := manager.Expire(now.Add(59 * time.Second)); len(expired) != 0 {
		t.Fatalf("expired early: %+v", expired)
	}
	if expired := manager.Expire(now.Add(time.Minute)); len(expired) != 1 || expired[0].ProxyID != "proxy-1" {
		t.Fatalf("Expire() = %+v, want proxy-1", expired)
	}

	second, err := manager.Assign(now.Add(time.Minute), "proxy-1", "worker-2", time.Minute)
	if err != nil {
		t.Fatalf("Assign(second): %v", err)
	}
	if second.Epoch <= first.Epoch {
		t.Fatalf("second epoch = %d, want greater than %d", second.Epoch, first.Epoch)
	}
}

func newTestOwnershipManager(t *testing.T, proxyIDs ...string) *OwnershipManager {
	t.Helper()
	candidates := make([]extractionDomain.Candidate, 0, len(proxyIDs))
	for _, proxyID := range proxyIDs {
		candidates = append(candidates, extractionDomain.Candidate{
			ID: proxyID, State: extractionDomain.Available,
		})
	}
	manager, err := NewOwnershipManager(extractionDomain.NewMemoryStore(candidates))
	if err != nil {
		t.Fatalf("NewOwnershipManager(): %v", err)
	}
	return manager
}