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 }