proxy-pool/internal/domain/activitypool/contracttest/contract.go

package contracttest

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

	"proxy-pool/internal/domain/activitypool"
	extractionDomain "proxy-pool/internal/domain/extraction"
	ownershipDomain "proxy-pool/internal/domain/ownership"
	proxyDomain "proxy-pool/internal/domain/proxy"
)

type Store interface {
	activitypool.Upserter
	activitypool.HealthStore
	activitypool.InventoryReader
	activitypool.StateInventoryReader
	activitypool.Maintainer
	extractionDomain.Store
	ownershipDomain.Repository
	ownershipDomain.DrainTicketStore
}

type Factory func(*testing.T) (Store, func())

func Run(t *testing.T, factory Factory) {
	t.Helper()
	t.Run("upsert capacity and incumbent lifecycle", func(t *testing.T) {
		runUpsertContract(t, newStore(t, factory))
	})
	t.Run("health and filtered extraction", func(t *testing.T) {
		runHealthAndExtractionContract(t, newStore(t, factory))
	})
	t.Run("fulfillment and gateway reserve", func(t *testing.T) {
		runFulfillmentContract(t, factory)
	})
	t.Run("business idempotency", func(t *testing.T) {
		runIdempotencyContract(t, factory)
	})
	t.Run("idempotency is bounded by proxy expiry", func(t *testing.T) {
		runIdempotencyExpiryContract(t, newStore(t, factory))
	})
	t.Run("ownership lifecycle", func(t *testing.T) {
		runOwnershipContract(t, factory)
	})
	t.Run("inventory and bounded maintenance", func(t *testing.T) {
		runMaintenanceContract(t, newStore(t, factory))
	})
	t.Run("state inventory lifecycle", func(t *testing.T) {
		runStateInventoryContract(t, newStore(t, factory))
	})
	t.Run("concurrent exclusivity", func(t *testing.T) {
		runConcurrencyContract(t, factory)
	})
	t.Run("canceled contexts", func(t *testing.T) {
		runCancellationContract(t, newStore(t, factory))
	})
}

func runUpsertContract(t *testing.T, store Store) {
	t.Helper()
	now := contractNow()
	firstProxy := contractProxy("proxy-a", "192.0.2.10", proxyDomain.StateFetched)
	secondProxy := contractProxy("proxy-b", "192.0.2.11", proxyDomain.StateFetched)
	result, err := store.UpsertFetched(context.Background(), "provider-a", activitypool.FetchedBatch{
		ObservedAt: now, ConfiguredTTL: 30 * time.Second, AllocationSafetyMargin: 3 * time.Second,
		MaxSize: 1, Proxies: []proxyDomain.Proxy{firstProxy, secondProxy},
	})
	if err != nil || result.Accepted != 2 || result.Inserted != 1 || result.Dropped != 1 {
		t.Fatalf("UpsertFetched(capacity) = %+v, %v", result, err)
	}
	assertInventory(t, store, "provider-a", now, 1)

	refreshed, err := store.UpsertFetched(context.Background(), "provider-b", activitypool.FetchedBatch{
		ObservedAt: now.Add(time.Second), ConfiguredTTL: 5 * time.Minute,
		AllocationSafetyMargin: 10 * time.Second, MaxSize: 10, Proxies: []proxyDomain.Proxy{firstProxy},
	})
	if err != nil || refreshed.Inserted != 0 || refreshed.Refreshed != 1 {
		t.Fatalf("UpsertFetched(cross-provider refresh) = %+v, %v", refreshed, err)
	}
	assertInventory(t, store, "provider-a", now.Add(time.Second), 1)
	assertInventory(t, store, "provider-b", now.Add(time.Second), 0)

	replaced, err := store.UpsertFetched(context.Background(), "provider-b", activitypool.FetchedBatch{
		ObservedAt: now.Add(31 * time.Second), ConfiguredTTL: time.Minute,
		AllocationSafetyMargin: 5 * time.Second, MaxSize: 10, Proxies: []proxyDomain.Proxy{firstProxy},
	})
	if err != nil || replaced.Inserted != 1 || replaced.Refreshed != 0 {
		t.Fatalf("UpsertFetched(after incumbent expiry) = %+v, %v", replaced, err)
	}
	assertInventory(t, store, "provider-a", now.Add(31*time.Second), 0)
	assertInventory(t, store, "provider-b", now.Add(31*time.Second), 1)
}

func runHealthAndExtractionContract(t *testing.T, store Store) {
	t.Helper()
	now := contractNow()
	candidate := contractProxy("proxy-a", "192.0.2.10", proxyDomain.StateFetched)
	candidate.Tags = map[string]string{"region": "cn", "carrier": "ct"}
	upsertOne(t, store, "provider-a", now, time.Minute, candidate)
	checking, err := store.ApplyHealth(context.Background(), activitypool.HealthUpdate{
		ProxyID: candidate.ID, CheckedAt: now.Add(time.Second), NextState: proxyDomain.StateChecking,
	})
	if err != nil || checking.State != proxyDomain.StateChecking {
		t.Fatalf("ApplyHealth(checking) = %+v, %v", checking, err)
	}
	available, err := store.ApplyHealth(context.Background(), activitypool.HealthUpdate{
		ProxyID: candidate.ID, CheckedAt: now.Add(2 * time.Second),
		NextState: proxyDomain.StateAvailable, Latency: 25 * time.Millisecond,
	})
	if err != nil || available.State != proxyDomain.StateAvailable ||
		available.Proxy.LastCheckedAt == nil || !available.Proxy.LastCheckedAt.Equal(now.Add(2*time.Second)) ||
		available.Proxy.LastSuccessAt == nil || !available.Proxy.LastSuccessAt.Equal(now.Add(2*time.Second)) {
		t.Fatalf("ApplyHealth(available) = %+v, %v", available, err)
	}
	if _, err := store.ApplyHealth(context.Background(), activitypool.HealthUpdate{
		ProxyID: candidate.ID, CheckedAt: now.Add(time.Second), NextState: proxyDomain.StateSuspect,
	}); !errors.Is(err, activitypool.ErrStaleHealthUpdate) {
		t.Fatalf("ApplyHealth(stale) error = %v", err)
	}

	result, err := store.Extract(context.Background(), extractionDomain.Command{
		RequestID: "req-filter", ClientID: "client-a", Requested: 1,
		Fulfillment: extractionDomain.Partial, Now: now.Add(3 * time.Second),
		MinRemainingTTL: 10 * time.Second, MaxHealthCheckAge: 5 * time.Second,
		Protocols: []string{"http"}, Regions: []string{"cn"}, Carriers: []string{"ct"},
		Upstreams: []string{"provider-a"},
	})
	if err != nil || result.Returned != 1 || result.Items[0].ID != candidate.ID ||
		result.Items[0].State != extractionDomain.Extracted {
		t.Fatalf("Extract(filtered) = %+v, %v", result, err)
	}
	assertInventory(t, store, "provider-a", now.Add(3*time.Second), 0)
}

func runFulfillmentContract(t *testing.T, factory Factory) {
	t.Helper()
	now := contractNow()
	store := newStore(t, factory)
	seedAvailable(t, store, "provider-a", now, time.Minute,
		contractProxy("proxy-a", "192.0.2.10", proxyDomain.StateAvailable))
	if result, err := store.Extract(context.Background(), extractionDomain.Command{
		RequestID: "req-all", ClientID: "client-a", Requested: 2,
		Fulfillment: extractionDomain.AllOrNothing, Now: now.Add(time.Second),
	}); !errors.Is(err, extractionDomain.ErrInsufficientProxies) || result.Returned != 0 {
		t.Fatalf("Extract(allOrNothing) = %+v, %v", result, err)
	}
	if result, err := store.Extract(context.Background(), extractionDomain.Command{
		RequestID: "req-partial", ClientID: "client-a", Requested: 1,
		Fulfillment: extractionDomain.Partial, Now: now.Add(time.Second),
	}); err != nil || result.Returned != 1 {
		t.Fatalf("Extract(partial after insufficient) = %+v, %v", result, err)
	}

	reservedStore := newStore(t, factory)
	seedAvailable(t, reservedStore, "provider-a", now, time.Minute,
		contractProxy("proxy-reserved", "192.0.2.11", proxyDomain.StateAvailable))
	reserved, err := reservedStore.Extract(context.Background(), extractionDomain.Command{
		RequestID: "req-reserved", ClientID: "client-a", Requested: 1,
		Fulfillment: extractionDomain.Partial, Now: now.Add(time.Second), ReserveForGateway: 1,
	})
	if err != nil || reserved.Returned != 0 || len(reserved.Items) != 0 {
		t.Fatalf("Extract(reserved) = %+v, %v", reserved, err)
	}

	marginStore := newStore(t, factory)
	marginProxy := contractProxy("proxy-margin", "192.0.2.12", proxyDomain.StateAvailable)
	if result, err := marginStore.UpsertFetched(context.Background(), "provider-a", activitypool.FetchedBatch{
		ObservedAt: now, ConfiguredTTL: 30 * time.Second, AllocationSafetyMargin: 5 * time.Second,
		MaxSize: 10, Proxies: []proxyDomain.Proxy{marginProxy},
	}); err != nil || result.Inserted != 1 {
		t.Fatalf("UpsertFetched(safety margin) = %+v, %v", result, err)
	}
	if result, err := marginStore.Extract(context.Background(), extractionDomain.Command{
		RequestID: "req-margin", ClientID: "client-a", Requested: 1,
		Fulfillment: extractionDomain.Partial, Now: now.Add(26 * time.Second),
	}); err != nil || result.Returned != 0 {
		t.Fatalf("Extract(after usableUntil) = %+v, %v", result, err)
	}
}

func runIdempotencyContract(t *testing.T, factory Factory) {
	t.Helper()
	now := contractNow()
	store := newStore(t, factory)
	candidate := contractProxy("proxy-a", "192.0.2.10", proxyDomain.StateAvailable)
	candidate.Tags = map[string]string{"region": "cn", "carrier": "ct"}
	checkedAt := now
	candidate.LastCheckedAt = &checkedAt
	seedAvailable(t, store, "provider-a", now, time.Minute, candidate)
	command := extractionDomain.Command{
		RequestID: "req-idem-first", ClientID: "client-a", SourceIP: "192.0.2.100",
		IdempotencyKey: "idem-contract", IdempotencyTTL: time.Minute,
		Requested: 1, Fulfillment: extractionDomain.Partial, Now: now.Add(time.Second),
		MinRemainingTTL: time.Second, MaxHealthCheckAge: time.Minute,
		Protocols: []string{"http", "http"}, Regions: []string{"cn", "cn"},
		Carriers: []string{"ct", "ct"}, Upstreams: []string{"provider-a", "provider-a"},
	}
	first, err := store.Extract(context.Background(), command)
	if err != nil || first.Returned != 1 {
		t.Fatalf("Extract(idempotent first) = %+v, %v", first, err)
	}
	command.RequestID = "req-idem-replay"
	command.SourceIP = "198.51.100.200"
	command.Now = now.Add(2 * time.Second)
	command.IdempotencyTTL = 2 * time.Minute
	command.MinRemainingTTL = 2 * time.Second
	command.MaxHealthCheckAge = 2 * time.Minute
	command.Protocols = []string{"http"}
	command.Regions = []string{"cn"}
	command.Carriers = []string{"ct"}
	command.Upstreams = []string{"provider-a"}
	replayed, err := store.Extract(context.Background(), command)
	if err != nil || replayed.Returned != 1 || replayed.Items[0].ID != first.Items[0].ID ||
		!replayed.ExtractedAt.Equal(first.ExtractedAt) {
		t.Fatalf("Extract(idempotent replay) = %+v, %v", replayed, err)
	}
	command.RequestID = "req-idem-conflict"
	command.Requested = 2
	if _, err := store.Extract(context.Background(), command); !errors.Is(err, extractionDomain.ErrIdempotencyConflict) {
		t.Fatalf("Extract(idempotent conflict) error = %v", err)
	}

	zeroStore := newStore(t, factory)
	zero := extractionDomain.Command{
		RequestID: "req-zero", ClientID: "client-a", IdempotencyKey: "idem-zero",
		Requested: 0, Fulfillment: extractionDomain.Partial, Now: now,
	}
	if result, err := zeroStore.Extract(context.Background(), zero); err != nil || result.Returned != 0 {
		t.Fatalf("Extract(zero first) = %+v, %v", result, err)
	}
	zero.RequestID = "req-zero-replay"
	zero.Now = now.Add(time.Second)
	if result, err := zeroStore.Extract(context.Background(), zero); err != nil || result.Returned != 0 {
		t.Fatalf("Extract(zero replay) = %+v, %v", result, err)
	}
	zero.RequestID = "req-zero-conflict"
	zero.Requested = 1
	if _, err := zeroStore.Extract(context.Background(), zero); !errors.Is(err, extractionDomain.ErrIdempotencyConflict) {
		t.Fatalf("Extract(zero conflict) error = %v", err)
	}
}

func runIdempotencyExpiryContract(t *testing.T, store Store) {
	t.Helper()
	now := time.Now().UTC().Truncate(time.Millisecond)
	candidate := contractProxy("short-lived", "192.0.2.20", proxyDomain.StateAvailable)
	upsertOne(t, store, "provider-a", now, 600*time.Millisecond, candidate)
	command := extractionDomain.Command{
		RequestID: "req-expiry-first", ClientID: "client-a", IdempotencyKey: "idem-expiry-contract",
		IdempotencyTTL: time.Minute, Requested: 1,
		Fulfillment: extractionDomain.Partial, Now: now.Add(time.Millisecond),
	}
	first, err := store.Extract(context.Background(), command)
	if err != nil || first.Returned != 1 {
		t.Fatalf("Extract(short-lived first) = %+v, %v", first, err)
	}
	time.Sleep(750 * time.Millisecond)
	replacementAt := time.Now().UTC().Truncate(time.Millisecond)
	upsertOne(t, store, "provider-a", replacementAt, time.Minute, candidate)
	command.RequestID = "req-expiry-second"
	command.Now = replacementAt.Add(time.Millisecond)
	again, err := store.Extract(context.Background(), command)
	if err != nil || again.Returned != 1 || again.ExtractedAt.Equal(first.ExtractedAt) {
		t.Fatalf("Extract(after idempotency expiry) = %+v, %v", again, err)
	}
}

func runOwnershipContract(t *testing.T, factory Factory) {
	t.Helper()
	now := contractNow()
	store := newStore(t, factory)
	seedAvailable(t, store, "provider-a", now, 2*time.Minute,
		contractProxy("proxy-a", "192.0.2.10", proxyDomain.StateAvailable))
	assigned, err := store.Assign(context.Background(), now.Add(time.Second), "proxy-a", "worker-a", 20*time.Second)
	if err != nil || assigned.Epoch == 0 || assigned.Version != 1 {
		t.Fatalf("Assign() = %+v, %v", assigned, err)
	}
	if _, err := store.Assign(context.Background(), now.Add(2*time.Second), "proxy-a", "worker-b", time.Minute); !errors.Is(err, ownershipDomain.ErrAlreadyOwned) {
		t.Fatalf("Assign(already owned) error = %v", err)
	}
	renewed, err := store.Renew(context.Background(), now.Add(10*time.Second), "proxy-a", "worker-a", assigned.Epoch, 5*time.Minute)
	if err != nil || renewed.Version != 2 || !renewed.ExpiresAt.Equal(now.Add(2*time.Minute)) {
		t.Fatalf("Renew() = %+v, %v", renewed, err)
	}
	draining, err := store.BeginDrain(context.Background(), "proxy-a", "worker-a", assigned.Epoch)
	if err != nil || !draining.Draining || draining.Version != 3 {
		t.Fatalf("BeginDrain() = %+v, %v", draining, err)
	}
	if replayed, err := store.BeginDrain(context.Background(), "proxy-a", "worker-a", assigned.Epoch); err != nil || replayed != draining {
		t.Fatalf("BeginDrain(replay) = %+v, %v", replayed, err)
	}
	pending, err := store.PendingDrains(context.Background(), "worker-a", 1)
	if err != nil || len(pending) != 1 || pending[0].ProxyID != "proxy-a" ||
		pending[0].AssignmentEpoch != assigned.Epoch || pending[0].RequiredSnapshotEpoch <= assigned.Epoch {
		t.Fatalf("PendingDrains() = %+v, %v", pending, err)
	}
	if err := store.AcknowledgeDrain(context.Background(), "proxy-a", "worker-a", assigned.Epoch, 1, 0); !errors.Is(err, ownershipDomain.ErrDrainNotReady) {
		t.Fatalf("AcknowledgeDrain(active) error = %v", err)
	}
	if err := store.AcknowledgeDrain(context.Background(), "proxy-a", "worker-a", assigned.Epoch, 0, 0); err != nil {
		t.Fatalf("AcknowledgeDrain(): %v", err)
	}
	if current, ok, err := store.Get(context.Background(), "proxy-a"); err != nil || ok {
		t.Fatalf("Get(after ACK) = %+v, %t, %v", current, ok, err)
	}
	if pending, err := store.PendingDrains(context.Background(), "worker-a", 1); err != nil || len(pending) != 0 {
		t.Fatalf("PendingDrains(after ACK) = %+v, %v", pending, err)
	}
	if result, err := store.Extract(context.Background(), extractionDomain.Command{
		RequestID: "req-after-ack", ClientID: "client-a", Requested: 1,
		Fulfillment: extractionDomain.Partial, Now: now.Add(20 * time.Second),
	}); err != nil || result.Returned != 1 {
		t.Fatalf("Extract(after ACK) = %+v, %v", result, err)
	}

	takeoverStore := newStore(t, factory)
	seedAvailable(t, takeoverStore, "provider-a", now, 2*time.Minute,
		contractProxy("takeover", "192.0.2.11", proxyDomain.StateAvailable))
	old, err := takeoverStore.Assign(context.Background(), now.Add(time.Second), "takeover", "worker-a", time.Second)
	if err != nil {
		t.Fatalf("Assign(takeover old): %v", err)
	}
	newAssignment, err := takeoverStore.Assign(context.Background(), now.Add(3*time.Second), "takeover", "worker-b", time.Minute)
	if err != nil || newAssignment.Epoch <= old.Epoch || newAssignment.WorkerID != "worker-b" {
		t.Fatalf("Assign(takeover new) = %+v, %v; old=%+v", newAssignment, err, old)
	}
}

func runMaintenanceContract(t *testing.T, store Store) {
	t.Helper()
	now := contractNow()
	for index := range 2 {
		upsertOne(t, store, "provider-a", now, 5*time.Second,
			contractProxy(fmt.Sprintf("proxy-%d", index), fmt.Sprintf("192.0.2.%d", index+10), proxyDomain.StateFetched))
	}
	assertInventory(t, store, "provider-a", now.Add(time.Second), 2)
	if removed, err := store.SweepExpired(context.Background(), now.Add(6*time.Second), 1); err != nil || removed != 1 {
		t.Fatalf("SweepExpired(first) = %d, %v", removed, err)
	}
	if removed, err := store.SweepExpired(context.Background(), now.Add(6*time.Second), 1); err != nil || removed != 1 {
		t.Fatalf("SweepExpired(second) = %d, %v", removed, err)
	}
	if removed, err := store.SweepExpired(context.Background(), now.Add(6*time.Second), 1); err != nil || removed != 0 {
		t.Fatalf("SweepExpired(empty) = %d, %v", removed, err)
	}
	assertInventory(t, store, "provider-a", now.Add(6*time.Second), 0)
}

func runStateInventoryContract(t *testing.T, store Store) {
	t.Helper()
	now := contractNow()
	upstreamA := "provider:a:FETCHED"
	upstreamB := "provider:a"
	states := []proxyDomain.State{
		proxyDomain.StateFetched,
		proxyDomain.StateChecking,
		proxyDomain.StateAvailable,
		proxyDomain.StateSuspect,
		proxyDomain.StateDraining,
		proxyDomain.StateUnhealthy,
		proxyDomain.StateExtracted,
	}
	for index, state := range states {
		upsertOne(t, store, upstreamA, now, time.Minute,
			contractProxy(fmt.Sprintf("state-%d", index), fmt.Sprintf("192.0.2.%d", index+30), state))
	}
	upsertOne(t, store, upstreamB, now, time.Minute,
		contractProxy("collision-control", "198.51.100.30", proxyDomain.StateFetched))

	inventories, err := store.ReadStateInventory(context.Background(), []string{upstreamB, upstreamA}, now)
	if err != nil || len(inventories) != 2 {
		t.Fatalf("ReadStateInventory() = %+v, %v", inventories, err)
	}
	if inventories[0] != (activitypool.StateInventory{UpstreamID: upstreamB, Fetched: 1}) {
		t.Fatalf("ReadStateInventory(collision control) = %+v", inventories[0])
	}
	wantAll := activitypool.StateInventory{
		UpstreamID: upstreamA, Fetched: 1, Checking: 1, Available: 1, Suspect: 1,
		Draining: 1, Unhealthy: 1, Extracted: 1,
	}
	if inventories[1] != wantAll {
		t.Fatalf("ReadStateInventory(all states) = %+v, want %+v", inventories[1], wantAll)
	}

	transition := activitypool.HealthUpdate{
		ProxyID: "state-0", CheckedAt: now.Add(time.Second), NextState: proxyDomain.StateChecking,
	}
	if _, err := store.ApplyHealth(context.Background(), transition); err != nil {
		t.Fatalf("ApplyHealth(state inventory transition): %v", err)
	}
	if _, err := store.ApplyHealth(context.Background(), transition); err != nil {
		t.Fatalf("ApplyHealth(state inventory replay): %v", err)
	}
	afterTransition, err := store.ReadStateInventory(context.Background(), []string{upstreamA}, now.Add(time.Second))
	if err != nil || len(afterTransition) != 1 || afterTransition[0].Fetched != 0 || afterTransition[0].Checking != 2 {
		t.Fatalf("ReadStateInventory(after transition) = %+v, %v", afterTransition, err)
	}

	extractCommand := extractionDomain.Command{
		RequestID: "state-inventory-extract", ClientID: "client-a", Requested: 1,
		IdempotencyKey: "state-inventory-extract", IdempotencyTTL: time.Minute,
		Fulfillment: extractionDomain.Partial, Now: now.Add(2 * time.Second),
		Upstreams: []string{upstreamA},
	}
	result, err := store.Extract(context.Background(), extractCommand)
	if err != nil || result.Returned != 1 {
		t.Fatalf("Extract(state inventory) = %+v, %v", result, err)
	}
	extractCommand.RequestID = "state-inventory-extract-replay"
	if replayed, err := store.Extract(context.Background(), extractCommand); err != nil || replayed.Returned != 1 ||
		replayed.Items[0].ID != result.Items[0].ID {
		t.Fatalf("Extract(state inventory replay) = %+v, %v", replayed, err)
	}
	afterExtract, err := store.ReadStateInventory(context.Background(), []string{upstreamA}, now.Add(2*time.Second))
	if err != nil || len(afterExtract) != 1 || afterExtract[0].Available != 0 || afterExtract[0].Extracted != 2 {
		t.Fatalf("ReadStateInventory(after extract) = %+v, %v", afterExtract, err)
	}

	afterExpiry, err := store.ReadStateInventory(context.Background(), []string{upstreamA, upstreamB}, now.Add(2*time.Minute))
	if err != nil || len(afterExpiry) != 2 || afterExpiry[0] != (activitypool.StateInventory{UpstreamID: upstreamA}) ||
		afterExpiry[1] != (activitypool.StateInventory{UpstreamID: upstreamB}) {
		t.Fatalf("ReadStateInventory(after expiry) = %+v, %v", afterExpiry, err)
	}

	empty, err := store.ReadStateInventory(context.Background(), nil, now)
	if err != nil || len(empty) != 0 {
		t.Fatalf("ReadStateInventory(empty) = %+v, %v", empty, err)
	}
}

func runConcurrencyContract(t *testing.T, factory Factory) {
	t.Helper()
	now := contractNow()
	extractStore := newStore(t, factory)
	for iteration := range 100 {
		proxyID := fmt.Sprintf("extract-race-%d", iteration)
		seedAvailable(t, extractStore, "provider-a", now, 10*time.Minute,
			contractProxy(proxyID, fmt.Sprintf("198.51.100.%d", iteration+1), proxyDomain.StateAvailable))
		results := make(chan extractionDomain.Result, 2)
		errorsCh := make(chan error, 2)
		var workers sync.WaitGroup
		for worker := range 2 {
			workers.Add(1)
			go func(worker int) {
				defer workers.Done()
				result, err := extractStore.Extract(context.Background(), extractionDomain.Command{
					RequestID: fmt.Sprintf("extract-race-%d-%d", iteration, worker),
					ClientID:  fmt.Sprintf("client-%d", worker), Requested: 1,
					Fulfillment: extractionDomain.Partial, Now: now.Add(time.Second),
				})
				results <- result
				errorsCh <- err
			}(worker)
		}
		workers.Wait()
		for range 2 {
			if err := <-errorsCh; err != nil {
				t.Fatalf("iteration %d Extract(): %v", iteration, err)
			}
		}
		returned := (<-results).Returned + (<-results).Returned
		if returned != 1 {
			t.Fatalf("iteration %d extracted = %d, want 1", iteration, returned)
		}
	}

	ownershipStore := newStore(t, factory)
	for iteration := range 100 {
		proxyID := fmt.Sprintf("ownership-race-%d", iteration)
		seedAvailable(t, ownershipStore, "provider-a", now, 10*time.Minute,
			contractProxy(proxyID, fmt.Sprintf("203.0.113.%d", iteration+1), proxyDomain.StateAvailable))
		var workers sync.WaitGroup
		assigned := make(chan bool, 1)
		extracted := make(chan bool, 1)
		errorsCh := make(chan error, 2)
		workers.Add(2)
		go func() {
			defer workers.Done()
			_, err := ownershipStore.Assign(context.Background(), now.Add(time.Second), proxyID, "worker-a", time.Minute)
			if err != nil && !errors.Is(err, ownershipDomain.ErrOwnershipUnavailable) {
				errorsCh <- err
			}
			assigned <- err == nil
		}()
		go func() {
			defer workers.Done()
			result, err := ownershipStore.Extract(context.Background(), extractionDomain.Command{
				RequestID: fmt.Sprintf("ownership-race-%d", iteration), ClientID: "client-a", Requested: 1,
				Fulfillment: extractionDomain.Partial, Now: now.Add(time.Second),
			})
			if err != nil {
				errorsCh <- err
			}
			extracted <- err == nil && result.Returned == 1
		}()
		workers.Wait()
		close(errorsCh)
		for err := range errorsCh {
			t.Fatalf("iteration %d ownership race: %v", iteration, err)
		}
		wins := 0
		if <-assigned {
			wins++
		}
		if <-extracted {
			wins++
		}
		if wins != 1 {
			t.Fatalf("iteration %d winners = %d, want 1", iteration, wins)
		}
	}
}

func runCancellationContract(t *testing.T, store Store) {
	t.Helper()
	now := contractNow()
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	checks := []struct {
		name string
		call func() error
	}{
		{name: "upsert", call: func() error {
			_, err := store.UpsertFetched(ctx, "provider-a", activitypool.FetchedBatch{ObservedAt: now, ConfiguredTTL: time.Minute, MaxSize: 1})
			return err
		}},
		{name: "health", call: func() error {
			_, err := store.ApplyHealth(ctx, activitypool.HealthUpdate{ProxyID: "proxy-a", CheckedAt: now, NextState: proxyDomain.StateChecking})
			return err
		}},
		{name: "inventory", call: func() error { _, err := store.Inventory(ctx, "provider-a", now); return err }},
		{name: "state inventory", call: func() error {
			_, err := store.ReadStateInventory(ctx, []string{"provider-a"}, now)
			return err
		}},
		{name: "sweep", call: func() error { _, err := store.SweepExpired(ctx, now, 1); return err }},
		{name: "extract", call: func() error {
			_, err := store.Extract(ctx, extractionDomain.Command{Requested: 1, Fulfillment: extractionDomain.Partial, Now: now})
			return err
		}},
		{name: "assign", call: func() error { _, err := store.Assign(ctx, now, "proxy-a", "worker-a", time.Minute); return err }},
		{name: "renew", call: func() error { _, err := store.Renew(ctx, now, "proxy-a", "worker-a", 1, time.Minute); return err }},
		{name: "begin drain", call: func() error { _, err := store.BeginDrain(ctx, "proxy-a", "worker-a", 1); return err }},
		{name: "pending drains", call: func() error { _, err := store.PendingDrains(ctx, "worker-a", 1); return err }},
		{name: "acknowledge drain", call: func() error { return store.AcknowledgeDrain(ctx, "proxy-a", "worker-a", 1, 0, 0) }},
		{name: "get", call: func() error { _, _, err := store.Get(ctx, "proxy-a"); return err }},
		{name: "expire", call: func() error { _, err := store.Expire(ctx, now, 1); return err }},
	}
	for _, check := range checks {
		t.Run(check.name, func(t *testing.T) {
			if err := check.call(); !errors.Is(err, context.Canceled) {
				t.Fatalf("error = %v, want context.Canceled", err)
			}
		})
	}
}

func newStore(t *testing.T, factory Factory) Store {
	t.Helper()
	store, cleanup := factory(t)
	if store == nil {
		t.Fatal("contract factory returned nil store")
	}
	if cleanup != nil {
		t.Cleanup(cleanup)
	}
	return store
}

func upsertOne(t *testing.T, store Store, upstreamID string, now time.Time, ttl time.Duration, candidate proxyDomain.Proxy) {
	t.Helper()
	result, err := store.UpsertFetched(context.Background(), upstreamID, activitypool.FetchedBatch{
		ObservedAt: now, ConfiguredTTL: ttl, MaxSize: 500, Proxies: []proxyDomain.Proxy{candidate},
	})
	if err != nil || result.Inserted != 1 {
		t.Fatalf("UpsertFetched(%s) = %+v, %v", candidate.ID, result, err)
	}
}

func seedAvailable(t *testing.T, store Store, upstreamID string, now time.Time, ttl time.Duration, candidate proxyDomain.Proxy) {
	t.Helper()
	candidate.State = proxyDomain.StateAvailable
	upsertOne(t, store, upstreamID, now, ttl, candidate)
}

func assertInventory(t *testing.T, store Store, upstreamID string, now time.Time, want int) {
	t.Helper()
	inventory, err := store.Inventory(context.Background(), upstreamID, now)
	if err != nil || inventory.UpstreamID != upstreamID || inventory.Managed != want {
		t.Fatalf("Inventory(%s) = %+v, %v; want %d", upstreamID, inventory, err, want)
	}
}

func contractProxy(id, host string, state proxyDomain.State) proxyDomain.Proxy {
	return proxyDomain.Proxy{
		ID: id, Scheme: proxyDomain.SchemeHTTP, Host: host, Port: 8080, State: state,
		Tags: map[string]string{"region": "cn", "carrier": "ct"},
	}
}

func contractNow() time.Time {
	return time.Now().UTC().Add(time.Hour).Truncate(time.Millisecond)
}