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

package controlplane

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

	controlplanev1 "proxy-pool/gen/controlplane/v1"
	domain "proxy-pool/internal/domain/outcome"
	gatewayOutcome "proxy-pool/internal/gateway/outcome"
)

func TestOutcomeReporterRetriesExactPendingBatchAfterTransportFailure(t *testing.T) {
	queue, err := gatewayOutcome.NewQueue(gatewayOutcome.QueueOptions{Capacity: 4, MaxBatch: 2})
	if err != nil {
		t.Fatalf("NewQueue() = %v", err)
	}
	queue.Record(domain.Event{ProxyID: "proxy-a", Stage: domain.StageDial, Success: true, ObservedAt: time.Now()})
	client := &outcomeClientStub{closeErr: errors.New("connection dropped")}
	reporter, err := NewOutcomeReporter(client, queue, OutcomeReporterOptions{WorkerID: "worker-a"})
	if err != nil {
		t.Fatalf("NewOutcomeReporter() = %v", err)
	}
	registration := Registration{SessionID: "session-a"}
	reporter.sessionID = registration.SessionID
	if err := reporter.Report(context.Background(), registration.SessionID); !errors.Is(err, client.closeErr) {
		t.Fatalf("Report(first) = %v, want transport failure", err)
	}
	client.closeErr = nil
	if err := reporter.Report(context.Background(), registration.SessionID); err != nil {
		t.Fatalf("Report(retry) = %v", err)
	}
	if len(client.batches) != 2 || client.batches[0].GetSequence() != 1 || client.batches[1].GetSequence() != 1 ||
		client.batches[0].GetOutcomes()[0].GetProxyId() != "proxy-a" {
		t.Fatalf("sent batches = %+v", client.batches)
	}
}

func TestOutcomeReporterCombinesAvailableBatchesAndRequiresConfirmation(t *testing.T) {
	queue, err := gatewayOutcome.NewQueue(gatewayOutcome.QueueOptions{Capacity: 8, MaxBatch: 1})
	if err != nil {
		t.Fatalf("NewQueue() = %v", err)
	}
	for _, proxyID := range []string{"proxy-a", "proxy-b"} {
		queue.Record(domain.Event{ProxyID: proxyID, Stage: domain.StageDial, Success: true, ObservedAt: time.Now()})
	}
	client := &outcomeClientStub{}
	reporter, err := NewOutcomeReporter(client, queue, OutcomeReporterOptions{WorkerID: "worker-a", MaxBatchesPerReportRPC: 2})
	if err != nil {
		t.Fatalf("NewOutcomeReporter() = %v", err)
	}
	reporter.sessionID = "session-a"
	if err := reporter.Report(context.Background(), "session-a"); err != nil {
		t.Fatalf("Report() = %v", err)
	}
	if len(client.batches) != 2 || client.batches[0].GetSequence() != 1 || client.batches[1].GetSequence() != 2 {
		t.Fatalf("sent batches = %+v", client.batches)
	}
}

func TestOutcomeReporterRetriesTransientFailureWithinRegisteredSession(t *testing.T) {
	queue, err := gatewayOutcome.NewQueue(gatewayOutcome.QueueOptions{Capacity: 4, MaxBatch: 2})
	if err != nil {
		t.Fatalf("NewQueue() = %v", err)
	}
	queue.Record(domain.Event{ProxyID: "proxy-a", Stage: domain.StageDial, Success: true, ObservedAt: time.Now()})
	client := &retryingOutcomeClient{failures: 1, confirmed: make(chan struct{}, 1)}
	reporter, err := NewOutcomeReporter(client, queue, OutcomeReporterOptions{WorkerID: "worker-a"})
	if err != nil {
		t.Fatalf("NewOutcomeReporter() = %v", err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	done := make(chan error, 1)
	go func() { done <- reporter.RunRegistered(ctx, Registration{SessionID: "session-a"}) }()

	select {
	case <-client.confirmed:
	case <-time.After(time.Second):
		t.Fatal("outcome batch was not confirmed after transient failure")
	}
	if batches := client.Batches(); len(batches) != 2 || batches[0].GetSequence() != 1 ||
		batches[1].GetSequence() != 1 || batches[0].GetOutcomes()[0].GetProxyId() != "proxy-a" {
		t.Fatalf("sent batches = %+v, want exact retry in the same session", batches)
	}
	cancel()
	if err := <-done; !errors.Is(err, context.Canceled) {
		t.Fatalf("RunRegistered() = %v, want context canceled", err)
	}
}

func TestOutcomeReporterRejectsUnboundedStreamOptions(t *testing.T) {
	queue, err := gatewayOutcome.NewQueue(gatewayOutcome.QueueOptions{Capacity: 1, MaxBatch: 1})
	if err != nil {
		t.Fatalf("NewQueue() = %v", err)
	}
	if _, err := NewOutcomeReporter(&outcomeClientStub{}, queue, OutcomeReporterOptions{
		WorkerID: "worker-a", MaxBatchesPerReportRPC: domain.MaxBatchesPerStream + 1,
	}); !errors.Is(err, ErrInvalidOutcomeReporter) {
		t.Fatalf("NewOutcomeReporter() error = %v, want ErrInvalidOutcomeReporter", err)
	}
}

type outcomeClientStub struct {
	batches  []*controlplanev1.OutcomeBatch
	closeErr error
	accepted uint64
}

func (client *outcomeClientStub) ReportOutcomes(context.Context) (OutcomeStream, error) {
	return &outcomeStreamStub{client: client}, nil
}

type outcomeStreamStub struct{ client *outcomeClientStub }

func (stream *outcomeStreamStub) Send(batch *controlplanev1.OutcomeBatch) error {
	stream.client.batches = append(stream.client.batches, batch)
	return nil
}

func (stream *outcomeStreamStub) CloseAndRecv() (*controlplanev1.ReportOutcomesResponse, error) {
	if stream.client.closeErr != nil {
		return nil, stream.client.closeErr
	}
	accepted := stream.client.accepted
	if accepted == 0 && len(stream.client.batches) > 0 {
		accepted = stream.client.batches[len(stream.client.batches)-1].GetSequence()
	}
	return &controlplanev1.ReportOutcomesResponse{AcceptedThroughSequence: accepted}, nil
}

type retryingOutcomeClient struct {
	mu        sync.Mutex
	batches   []*controlplanev1.OutcomeBatch
	failures  int
	confirmed chan struct{}
}

func (client *retryingOutcomeClient) ReportOutcomes(context.Context) (OutcomeStream, error) {
	return &retryingOutcomeStream{client: client}, nil
}

func (client *retryingOutcomeClient) Batches() []*controlplanev1.OutcomeBatch {
	client.mu.Lock()
	defer client.mu.Unlock()
	return append([]*controlplanev1.OutcomeBatch(nil), client.batches...)
}

type retryingOutcomeStream struct{ client *retryingOutcomeClient }

func (stream *retryingOutcomeStream) Send(batch *controlplanev1.OutcomeBatch) error {
	stream.client.mu.Lock()
	stream.client.batches = append(stream.client.batches, batch)
	stream.client.mu.Unlock()
	return nil
}

func (stream *retryingOutcomeStream) CloseAndRecv() (*controlplanev1.ReportOutcomesResponse, error) {
	stream.client.mu.Lock()
	if stream.client.failures > 0 {
		stream.client.failures--
		stream.client.mu.Unlock()
		return nil, errors.New("temporary delivery failure")
	}
	accepted := stream.client.batches[len(stream.client.batches)-1].GetSequence()
	confirmed := stream.client.confirmed
	stream.client.mu.Unlock()
	select {
	case confirmed <- struct{}{}:
	default:
	}
	return &controlplanev1.ReportOutcomesResponse{AcceptedThroughSequence: accepted}, nil
}