proxy-pool/internal/controller/worker/snapshot_source_test.go

package worker

import (
	"context"
	"errors"
	"testing"
	"time"

	controlplanev1 "proxy-pool/gen/controlplane/v1"
	"proxy-pool/internal/controlplane/snapshotwire"
	ownershipDomain "proxy-pool/internal/domain/ownership"
	proxyDomain "proxy-pool/internal/domain/proxy"
)

func TestInitialSnapshotSourceIssuesNextFullSnapshot(t *testing.T) {
	now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC)
	source, err := NewInitialSnapshotSource(epochReaderStub{epoch: 9}, time.Minute, func() time.Time { return now })
	if err != nil {
		t.Fatalf("NewInitialSnapshotSource(): %v", err)
	}
	updates, err := source.Watch(context.Background(), SnapshotWatchRequest{
		WorkerID: "worker-a", SessionID: "session-a", LastAppliedVersion: 4,
	})
	if err != nil {
		t.Fatalf("Watch(): %v", err)
	}
	full := <-updates
	if full.GetVersion() != 5 || full.GetOwnershipEpoch() != 9 || !full.GetValidUntil().AsTime().Equal(now.Add(time.Minute)) {
		t.Fatalf("snapshot = %+v", full)
	}
	checksum, err := snapshotwire.Checksum(full)
	if err != nil || string(checksum[:]) != string(full.GetChecksum()) {
		t.Fatalf("snapshot checksum = %x, %v; want %x", full.GetChecksum(), err, checksum)
	}
}

func TestOwnedSnapshotSourceBuildsBoundedProxySnapshot(t *testing.T) {
	now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC)
	expiresAt := now.Add(10 * time.Minute)
	usableUntil := now.Add(5 * time.Minute)
	source, err := NewOwnedSnapshotSource(epochReaderStub{epoch: 9}, snapshotReaderStub{proxies: []ownershipDomain.SnapshotProxy{{
		Proxy: proxyDomain.Proxy{
			ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080,
			SourceUpstream: "upstream-a", ExpiresAt: &expiresAt, UsableUntil: &usableUntil,
			MaxConcurrency: 7, State: proxyDomain.StateAvailable, Tags: map[string]string{"region": "cn"},
		},
		OwnershipEpoch: 4, LeaseExpiresAt: now.Add(time.Minute),
	}}}, time.Minute*2, 10, 4096, func() time.Time { return now })
	if err != nil {
		t.Fatalf("NewOwnedSnapshotSource(): %v", err)
	}
	updates, err := source.Watch(context.Background(), SnapshotWatchRequest{WorkerID: "worker-a", SessionID: "session-a"})
	if err != nil {
		t.Fatalf("Watch(): %v", err)
	}
	full := <-updates
	if full.GetVersion() != 1 || full.GetOwnershipEpoch() != 9 || len(full.GetProxies()) != 1 {
		t.Fatalf("snapshot = %+v", full)
	}
	proxy := full.GetProxies()[0]
	if proxy.GetOwnershipEpoch() != 4 || !proxy.GetUsableUntil().AsTime().Equal(now.Add(time.Minute)) ||
		!full.GetValidUntil().AsTime().Equal(now.Add(time.Minute)) {
		t.Fatalf("wire proxy = %+v, valid until = %s", proxy, full.GetValidUntil().AsTime())
	}
	checksum, err := snapshotwire.Checksum(full)
	if err != nil || string(checksum[:]) != string(full.GetChecksum()) {
		t.Fatalf("snapshot checksum = %x, %v; want %x", full.GetChecksum(), err, checksum)
	}
}

func TestOwnedSnapshotSourceIncludesRoutingInFullSnapshotChecksum(t *testing.T) {
	now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC)
	expiresAt := now.Add(10 * time.Minute)
	usableUntil := now.Add(5 * time.Minute)
	routing := &controlplanev1.RoutingRule{
		Name: "gateway", Enabled: true, HostRegex: ".*", Upstreams: []string{"upstream-a"},
		Strategy:      &controlplanev1.RoutingStrategy{Type: controlplanev1.StrategyType_STRATEGY_TYPE_RANDOM},
		OnUnavailable: controlplanev1.UnavailableAction_UNAVAILABLE_ACTION_REJECT,
	}
	source, err := NewOwnedSnapshotSource(epochReaderStub{epoch: 9}, snapshotReaderStub{proxies: []ownershipDomain.SnapshotProxy{{
		Proxy: proxyDomain.Proxy{
			ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080,
			SourceUpstream: "upstream-a", ExpiresAt: &expiresAt, UsableUntil: &usableUntil,
			MaxConcurrency: 7, State: proxyDomain.StateAvailable,
		},
		OwnershipEpoch: 4, LeaseExpiresAt: now.Add(time.Minute),
	}}}, time.Minute*2, 10, 4096, func() time.Time { return now }, routingSourceStub{rules: []*controlplanev1.RoutingRule{routing}})
	if err != nil {
		t.Fatalf("NewOwnedSnapshotSource(): %v", err)
	}
	updates, err := source.Watch(context.Background(), SnapshotWatchRequest{WorkerID: "worker-a", SessionID: "session-a"})
	if err != nil {
		t.Fatalf("Watch(): %v", err)
	}
	full := <-updates
	if len(full.GetRouting()) != 1 || full.GetRouting()[0].GetName() != "gateway" || full.GetRouting()[0] == routing {
		t.Fatalf("snapshot routing = %+v", full.GetRouting())
	}
	checksum, err := snapshotwire.Checksum(full)
	if err != nil || string(checksum[:]) != string(full.GetChecksum()) {
		t.Fatalf("snapshot checksum = %x, %v; want %x", full.GetChecksum(), err, checksum)
	}
}

func TestOwnedSnapshotSourceRejectsCredentialReferenceUntilMaterialIsAvailable(t *testing.T) {
	now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC)
	expiresAt := now.Add(time.Minute)
	usableUntil := now.Add(time.Minute)
	source, err := NewOwnedSnapshotSource(epochReaderStub{epoch: 1}, snapshotReaderStub{proxies: []ownershipDomain.SnapshotProxy{{
		Proxy: proxyDomain.Proxy{
			ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080,
			SourceUpstream: "upstream-a", ExpiresAt: &expiresAt, UsableUntil: &usableUntil,
			MaxConcurrency: 1, State: proxyDomain.StateAvailable, CredentialVersion: "v1",
		},
		OwnershipEpoch: 1, LeaseExpiresAt: now.Add(time.Minute),
	}}}, time.Minute, 10, 4096, func() time.Time { return now })
	if err != nil {
		t.Fatalf("NewOwnedSnapshotSource(): %v", err)
	}
	_, err = source.Watch(context.Background(), SnapshotWatchRequest{WorkerID: "worker-a", SessionID: "session-a"})
	if !errors.Is(err, ErrSnapshotCredentialsUnavailable) {
		t.Fatalf("Watch() error = %v, want ErrSnapshotCredentialsUnavailable", err)
	}
}

type epochReaderStub struct {
	epoch uint64
	err   error
}

func (reader epochReaderStub) CurrentOwnershipEpoch(context.Context) (uint64, error) {
	return reader.epoch, reader.err
}

type snapshotReaderStub struct {
	proxies []ownershipDomain.SnapshotProxy
	err     error
}

type routingSourceStub struct {
	rules []*controlplanev1.RoutingRule
	err   error
}

func (source routingSourceStub) Read(context.Context) ([]*controlplanev1.RoutingRule, error) {
	return source.rules, source.err
}

func (reader snapshotReaderStub) ReadWorkerSnapshot(context.Context, string, int) ([]ownershipDomain.SnapshotProxy, error) {
	return reader.proxies, reader.err
}