proxy-pool/internal/gateway/controlplane/watcher_test.go

package controlplane

import (
	"context"
	"crypto/sha256"
	"io"
	"testing"
	"time"

	controlplanev1 "proxy-pool/gen/controlplane/v1"
	"proxy-pool/internal/controlplane/snapshotwire"
	"proxy-pool/internal/gateway/snapshot"

	"google.golang.org/protobuf/types/known/emptypb"
	"google.golang.org/protobuf/types/known/timestamppb"
)

func TestSnapshotWatcherAppliesVerifiedFullSnapshotAndAcknowledges(t *testing.T) {
	store := snapshot.NewStore("cluster-a", "worker-a")
	full := &controlplanev1.WorkerSnapshot{
		Version: 1, OwnershipEpoch: 7, GeneratedAt: timestamppb.New(time.Now()), ValidUntil: timestamppb.New(time.Now().Add(time.Minute)),
		Proxies: []*controlplanev1.OwnedProxy{{
			Id: "proxy-a", Upstream: "upstream-a", Protocol: controlplanev1.ProxyProtocol_PROXY_PROTOCOL_HTTP,
			Host: "192.0.2.10", Port: 8080, MaxConcurrency: 3, ExpiresAt: timestamppb.New(time.Now().Add(time.Minute)),
		}},
	}
	setSnapshotChecksum(t, full)
	client := &snapshotClientStub{stream: &snapshotStreamStub{values: []*controlplanev1.SnapshotEnvelope{{Payload: &controlplanev1.SnapshotEnvelope_Full{Full: full}}}}}
	watcher, err := NewSnapshotWatcher(client, store, SnapshotWatcherOptions{ClusterID: "cluster-a", WorkerID: "worker-a"})
	if err != nil {
		t.Fatalf("NewSnapshotWatcher(): %v", err)
	}
	if err := watcher.Watch(context.Background(), "session-a"); err != nil {
		t.Fatalf("Watch(): %v", err)
	}
	view := store.Current()
	if view == nil || view.Version != 1 || view.Epoch != 7 || len(view.Entries) != 1 || view.Entries[0].Proxy.ID != "proxy-a" ||
		!view.ValidUntil.Equal(full.GetValidUntil().AsTime()) {
		t.Fatalf("snapshot view = %+v", view)
	}
	if client.watch.GetSessionId() != "session-a" || client.ack.GetVersion() != 1 || !client.ack.GetApplied() || string(client.ack.GetChecksum()) != string(full.GetChecksum()) {
		t.Fatalf("watch=%+v ack=%+v", client.watch, client.ack)
	}
}

func TestSnapshotWatcherRejectsChecksumAndAcknowledgesFailure(t *testing.T) {
	store := snapshot.NewStore("cluster-a", "worker-a")
	full := &controlplanev1.WorkerSnapshot{Version: 1, OwnershipEpoch: 7, Checksum: make([]byte, sha256.Size)}
	client := &snapshotClientStub{stream: &snapshotStreamStub{values: []*controlplanev1.SnapshotEnvelope{{Payload: &controlplanev1.SnapshotEnvelope_Full{Full: full}}}}}
	watcher, err := NewSnapshotWatcher(client, store, SnapshotWatcherOptions{ClusterID: "cluster-a", WorkerID: "worker-a"})
	if err != nil {
		t.Fatalf("NewSnapshotWatcher(): %v", err)
	}
	if err := watcher.Watch(context.Background(), "session-a"); err == nil {
		t.Fatal("Watch() error = nil, want checksum rejection")
	}
	if client.ack == nil || client.ack.GetApplied() || client.ack.GetVersion() != 1 {
		t.Fatalf("negative acknowledgement = %+v", client.ack)
	}
}

func TestSnapshotWatcherRejectsMissingOverallValidityDeadline(t *testing.T) {
	store := snapshot.NewStore("cluster-a", "worker-a")
	full := &controlplanev1.WorkerSnapshot{Version: 1, OwnershipEpoch: 7, GeneratedAt: timestamppb.New(time.Now())}
	setSnapshotChecksum(t, full)
	client := &snapshotClientStub{stream: &snapshotStreamStub{values: []*controlplanev1.SnapshotEnvelope{{Payload: &controlplanev1.SnapshotEnvelope_Full{Full: full}}}}}
	watcher, err := NewSnapshotWatcher(client, store, SnapshotWatcherOptions{ClusterID: "cluster-a", WorkerID: "worker-a"})
	if err != nil {
		t.Fatalf("NewSnapshotWatcher(): %v", err)
	}
	if err := watcher.Watch(context.Background(), "session-a"); err == nil {
		t.Fatal("Watch() error = nil, want missing validity rejection")
	}
	if client.ack == nil || client.ack.GetApplied() || client.ack.GetErrorCode() != "snapshot_apply_failed" {
		t.Fatalf("negative acknowledgement = %+v", client.ack)
	}
}

func TestSnapshotWatcherRejectsExpiredOverallValidityDeadline(t *testing.T) {
	store := snapshot.NewStore("cluster-a", "worker-a")
	full := &controlplanev1.WorkerSnapshot{
		Version: 1, OwnershipEpoch: 7, GeneratedAt: timestamppb.New(time.Now().Add(-time.Minute)),
		ValidUntil: timestamppb.New(time.Now().Add(-time.Second)),
	}
	setSnapshotChecksum(t, full)
	client := &snapshotClientStub{stream: &snapshotStreamStub{values: []*controlplanev1.SnapshotEnvelope{{Payload: &controlplanev1.SnapshotEnvelope_Full{Full: full}}}}}
	watcher, err := NewSnapshotWatcher(client, store, SnapshotWatcherOptions{ClusterID: "cluster-a", WorkerID: "worker-a"})
	if err != nil {
		t.Fatalf("NewSnapshotWatcher(): %v", err)
	}
	if err := watcher.Watch(context.Background(), "session-a"); err == nil {
		t.Fatal("Watch() error = nil, want expired validity rejection")
	}
	if client.ack == nil || client.ack.GetApplied() {
		t.Fatalf("negative acknowledgement = %+v", client.ack)
	}
}

type snapshotClientStub struct {
	stream SnapshotStream
	watch  *controlplanev1.WatchSnapshotsRequest
	ack    *controlplanev1.AcknowledgeSnapshotRequest
}

func (client *snapshotClientStub) Watch(_ context.Context, request *controlplanev1.WatchSnapshotsRequest) (SnapshotStream, error) {
	client.watch = request
	return client.stream, nil
}

func (client *snapshotClientStub) Acknowledge(_ context.Context, acknowledgement *controlplanev1.AcknowledgeSnapshotRequest) (*emptypb.Empty, error) {
	client.ack = acknowledgement
	return &emptypb.Empty{}, nil
}

type snapshotStreamStub struct {
	values []*controlplanev1.SnapshotEnvelope
	index  int
}

func (stream *snapshotStreamStub) Recv() (*controlplanev1.SnapshotEnvelope, error) {
	if stream.index >= len(stream.values) {
		return nil, io.EOF
	}
	value := stream.values[stream.index]
	stream.index++
	return value, nil
}

func setSnapshotChecksum(t *testing.T, full *controlplanev1.WorkerSnapshot) {
	t.Helper()
	checksum, err := snapshotwire.Checksum(full)
	if err != nil {
		t.Fatalf("workerSnapshotChecksum(): %v", err)
	}
	full.Checksum = checksum[:]
}