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

package controlplane

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

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

	"google.golang.org/protobuf/types/known/durationpb"
)

func TestRuntimeReporterRegistersAndReportsSnapshotCounters(t *testing.T) {
	store := snapshot.NewStore("cluster-a", "worker-a")
	applySnapshot(t, store)
	client := &clientStub{registration: &controlplanev1.RegisterWorkerResponse{
		WorkerId: "worker-a", SessionId: "session-a", OwnershipEpoch: 7,
		HeartbeatInterval: durationpb.New(10 * time.Second), MaxStaleAge: durationpb.New(30 * time.Second),
	}}
	now := time.Date(2026, 7, 31, 10, 0, 0, 0, time.UTC)
	reporter, err := NewRuntimeReporter(client, store, Options{
		WorkerID: "worker-a", InstanceID: "instance-a", Zone: "zone-a", ProtocolVersion: 1,
		Now: func() time.Time { return now },
	})
	if err != nil {
		t.Fatalf("NewRuntimeReporter(): %v", err)
	}
	registration, err := reporter.Register(context.Background())
	if err != nil || registration.SessionID != "session-a" || client.register.GetZone() != "zone-a" {
		t.Fatalf("Register() = %+v, %v; request=%+v", registration, err, client.register)
	}
	if err := reporter.Report(context.Background()); err != nil {
		t.Fatalf("Report(): %v", err)
	}
	if client.runtime.GetSessionId() != "session-a" || client.runtime.GetReportSequence() != 1 ||
		client.runtime.GetSnapshotVersion() != 1 || !client.runtime.GetObservedAt().AsTime().Equal(now) {
		t.Fatalf("ReportRuntime request = %+v", client.runtime)
	}
	if err := reporter.Report(context.Background()); err != nil {
		t.Fatalf("second Report(): %v", err)
	}
	if client.runtime.GetReportSequence() != 2 {
		t.Fatalf("second report sequence = %d, want 2", client.runtime.GetReportSequence())
	}
}

func TestRuntimeReporterRetainsSequenceOnTransportErrorAndSurfacesResync(t *testing.T) {
	store := snapshot.NewStore("cluster-a", "worker-a")
	applySnapshot(t, store)
	transportErr := errors.New("temporary transport failure")
	client := &clientStub{registration: &controlplanev1.RegisterWorkerResponse{
		WorkerId: "worker-a", SessionId: "session-a", OwnershipEpoch: 7,
		HeartbeatInterval: durationpb.New(time.Second), MaxStaleAge: durationpb.New(3 * time.Second),
	}, runtimeErr: transportErr}
	reporter, err := NewRuntimeReporter(client, store, Options{WorkerID: "worker-a", InstanceID: "instance-a", Zone: "zone-a", ProtocolVersion: 1, Now: time.Now})
	if err != nil {
		t.Fatalf("NewRuntimeReporter(): %v", err)
	}
	if _, err := reporter.Register(context.Background()); err != nil {
		t.Fatalf("Register(): %v", err)
	}
	if err := reporter.Report(context.Background()); !errors.Is(err, transportErr) {
		t.Fatalf("Report() error = %v, want transport error", err)
	}
	client.runtimeErr = nil
	client.runtimeResponse = &controlplanev1.ReportRuntimeResponse{RequireFullSnapshot: true}
	if err := reporter.Report(context.Background()); !errors.Is(err, ErrFullSnapshotRequired) {
		t.Fatalf("Report() error = %v, want ErrFullSnapshotRequired", err)
	}
	if client.runtime.GetReportSequence() != 1 {
		t.Fatalf("retried report sequence = %d, want 1", client.runtime.GetReportSequence())
	}
}

func TestRuntimeReporterRejectsInvalidState(t *testing.T) {
	store := snapshot.NewStore("cluster-a", "worker-a")
	if _, err := NewRuntimeReporter(&clientStub{}, store, Options{}); !errors.Is(err, ErrInvalidOptions) {
		t.Fatalf("NewRuntimeReporter() error = %v, want ErrInvalidOptions", err)
	}
	reporter, err := NewRuntimeReporter(&clientStub{}, store, Options{WorkerID: "worker-a", InstanceID: "instance-a", Zone: "zone-a", ProtocolVersion: 1, Now: time.Now})
	if err != nil {
		t.Fatalf("NewRuntimeReporter(): %v", err)
	}
	if err := reporter.Report(context.Background()); !errors.Is(err, ErrNotRegistered) {
		t.Fatalf("Report() error = %v, want ErrNotRegistered", err)
	}
}

func TestRuntimeReporterRunWaitsForSnapshotThenReports(t *testing.T) {
	store := snapshot.NewStore("cluster-a", "worker-a")
	client := &runClient{runtime: make(chan *controlplanev1.ReportRuntimeRequest, 8)}
	reporter, err := NewRuntimeReporter(client, store, Options{
		WorkerID: "worker-a", InstanceID: "instance-a", Zone: "zone-a", ProtocolVersion: 1, Now: time.Now,
	})
	if err != nil {
		t.Fatalf("NewRuntimeReporter(): %v", err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	result := make(chan error, 1)
	go func() { result <- reporter.Run(ctx) }()
	time.Sleep(15 * time.Millisecond)
	select {
	case request := <-client.runtime:
		t.Fatalf("ReportRuntime() before snapshot = %+v", request)
	default:
	}
	applySnapshot(t, store)
	select {
	case request := <-client.runtime:
		if request.GetReportSequence() != 1 {
			t.Fatalf("report sequence = %d, want 1", request.GetReportSequence())
		}
	case <-time.After(time.Second):
		t.Fatal("ReportRuntime() was not called after applying a snapshot")
	}
	cancel()
	if err := <-result; !errors.Is(err, context.Canceled) {
		t.Fatalf("Run() error = %v, want context canceled", err)
	}
}

type clientStub struct {
	register        *controlplanev1.RegisterWorkerRequest
	registration    *controlplanev1.RegisterWorkerResponse
	registerErr     error
	runtime         *controlplanev1.ReportRuntimeRequest
	runtimeResponse *controlplanev1.ReportRuntimeResponse
	runtimeErr      error
}

func (client *clientStub) RegisterWorker(_ context.Context, request *controlplanev1.RegisterWorkerRequest) (*controlplanev1.RegisterWorkerResponse, error) {
	client.register = request
	return client.registration, client.registerErr
}

func (client *clientStub) ReportRuntime(_ context.Context, request *controlplanev1.ReportRuntimeRequest) (*controlplanev1.ReportRuntimeResponse, error) {
	client.runtime = request
	if client.runtimeResponse == nil {
		client.runtimeResponse = &controlplanev1.ReportRuntimeResponse{AcceptedOwnershipEpoch: 7}
	}
	return client.runtimeResponse, client.runtimeErr
}

func applySnapshot(t *testing.T, store *snapshot.Store) {
	t.Helper()
	envelope := snapshot.Envelope{ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 7, Version: 1, Full: true}
	envelope.Checksum = snapshot.Checksum(nil)
	if err := store.Apply(envelope); err != nil {
		t.Fatalf("Apply(): %v", err)
	}
}

type runClient struct {
	runtime chan *controlplanev1.ReportRuntimeRequest
}

func (*runClient) RegisterWorker(context.Context, *controlplanev1.RegisterWorkerRequest) (*controlplanev1.RegisterWorkerResponse, error) {
	return &controlplanev1.RegisterWorkerResponse{
		WorkerId: "worker-a", SessionId: "session-a", OwnershipEpoch: 7,
		HeartbeatInterval: durationpb.New(10 * time.Millisecond), MaxStaleAge: durationpb.New(time.Second),
	}, nil
}

func (client *runClient) ReportRuntime(_ context.Context, request *controlplanev1.ReportRuntimeRequest) (*controlplanev1.ReportRuntimeResponse, error) {
	client.runtime <- request
	return &controlplanev1.ReportRuntimeResponse{AcceptedOwnershipEpoch: 7}, nil
}