proxy-pool/internal/checker/controlplane/runner_test.go

package controlplane

import (
	"context"
	"errors"
	"io"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	controlplanev1 "proxy-pool/gen/controlplane/v1"
	"proxy-pool/internal/checker/probe"

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

func TestRunnerExecutesBoundedTaskBatchAndReportsFacts(t *testing.T) {
	now := time.Date(2026, 8, 1, 13, 0, 0, 0, time.UTC)
	client := &clientStub{stream: &taskStreamStub{tasks: []*controlplanev1.CheckTask{
		checkerTask("task-a", now), checkerTask("task-b", now),
	}}}
	executor := &executorStub{}
	runner, err := NewRunner(client, executor, Options{
		CheckerID: "checker-a", InstanceID: "instance-a", MaxInFlight: 2,
		SupportedLevels: []controlplanev1.CheckLevel{controlplanev1.CheckLevel_CHECK_LEVEL_BASIC},
		ReportBatchSize: 2, RetryDelay: time.Millisecond, Now: func() time.Time { return now },
	})
	if err != nil {
		t.Fatalf("NewRunner(): %v", err)
	}
	if err := runner.RunOnce(context.Background()); err != nil {
		t.Fatalf("RunOnce(): %v", err)
	}
	if len(client.batches) != 1 || len(client.batches[0].GetObservations()) != 2 || len(executor.tasks) != 3 {
		t.Fatalf("batches=%+v executions=%d", client.batches, len(executor.tasks))
	}
	first, second := client.batches[0].GetObservations()[0], client.batches[0].GetObservations()[1]
	if first.GetTaskId() != "task-a" || !first.GetSuccess() || first.GetFailureClass() != "" ||
		second.GetTaskId() != "task-b" || !second.GetSuccess() || first.GetObservedAt().AsTime() != now {
		t.Fatalf("observations = %+v", client.batches[0].GetObservations())
	}
}

func TestRunnerOmitsEgressProbeURLFromGlobalObservation(t *testing.T) {
	now := time.Date(2026, 8, 1, 13, 30, 0, 0, time.UTC)
	task := checkerTask("task-egress", now)
	task.Level = controlplanev1.CheckLevel_CHECK_LEVEL_EGRESS
	task.TargetUrl = "https://egress.example/identity"
	client := &clientStub{stream: &taskStreamStub{tasks: []*controlplanev1.CheckTask{task}}}
	runner, err := NewRunner(client, &executorStub{}, Options{
		CheckerID: "checker-a", InstanceID: "instance-a", MaxInFlight: 1,
		SupportedLevels: []controlplanev1.CheckLevel{controlplanev1.CheckLevel_CHECK_LEVEL_EGRESS},
		ReportBatchSize: 1, RetryDelay: time.Millisecond, Now: func() time.Time { return now },
	})
	if err != nil {
		t.Fatalf("NewRunner(): %v", err)
	}
	if err := runner.RunOnce(context.Background()); err != nil {
		t.Fatalf("RunOnce(): %v", err)
	}
	observation := client.batches[0].GetObservations()[0]
	if observation.GetLevel() != controlplanev1.CheckLevel_CHECK_LEVEL_EGRESS || observation.GetTargetUrl() != "" || observation.GetRoutingName() != "" {
		t.Fatalf("Observation = %+v, want global EGRESS fact without target profile", observation)
	}
	if observation.GetObservedEgressIp() != "198.51.100.42" {
		t.Fatalf("Observation egress IP = %q", observation.GetObservedEgressIp())
	}
}

func TestRunnerMarksReadyAfterSuccessfulEmptyPull(t *testing.T) {
	var successfulPulls atomic.Int64
	runner, err := NewRunner(&clientStub{stream: &taskStreamStub{}}, &executorStub{}, Options{
		CheckerID: "checker-a", InstanceID: "instance-a", MaxInFlight: 1,
		SupportedLevels: []controlplanev1.CheckLevel{controlplanev1.CheckLevel_CHECK_LEVEL_BASIC},
		Now:             time.Now, OnSuccessfulPull: func() { successfulPulls.Add(1) },
	})
	if err != nil {
		t.Fatalf("NewRunner(): %v", err)
	}
	if err := runner.RunOnce(context.Background()); err != nil {
		t.Fatalf("RunOnce(): %v", err)
	}
	if successfulPulls.Load() != 1 {
		t.Fatalf("successful pulls = %d, want 1", successfulPulls.Load())
	}
}

func TestRunnerMarksUnavailableAfterPullFailure(t *testing.T) {
	failed := make(chan struct{})
	var markOnce sync.Once
	runner, err := NewRunner(&clientStub{streamErr: errors.New("control plane unavailable")}, &executorStub{}, Options{
		CheckerID: "checker-a", InstanceID: "instance-a", MaxInFlight: 1,
		SupportedLevels: []controlplanev1.CheckLevel{controlplanev1.CheckLevel_CHECK_LEVEL_BASIC},
		Now:             time.Now, RetryDelay: time.Millisecond,
		OnFailedPull: func() { markOnce.Do(func() { close(failed) }) },
	})
	if err != nil {
		t.Fatalf("NewRunner(): %v", err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	completed := make(chan error, 1)
	go func() { completed <- runner.Run(ctx) }()
	select {
	case <-failed:
		cancel()
	case <-time.After(time.Second):
		t.Fatal("Runner did not mark failed pull")
	}
	if err := <-completed; !errors.Is(err, context.Canceled) {
		t.Fatalf("Run() = %v, want context.Canceled", err)
	}
}

func checkerTask(id string, now time.Time) *controlplanev1.CheckTask {
	return &controlplanev1.CheckTask{
		TaskId: id, ProxyId: id + "-proxy", Protocol: controlplanev1.ProxyProtocol_PROXY_PROTOCOL_HTTP,
		Host: "proxy.example", Port: 8080, Level: controlplanev1.CheckLevel_CHECK_LEVEL_BASIC,
		Timeout: durationpb.New(time.Second), Attempt: 1, MaxAttempts: 2, Deadline: timestamppb.New(now.Add(time.Minute)),
	}
}

type clientStub struct {
	stream    TaskStream
	streamErr error
	batches   []*controlplanev1.ObservationBatch
}

func (stub *clientStub) StreamCheckTasks(context.Context, *controlplanev1.StreamCheckTasksRequest) (TaskStream, error) {
	return stub.stream, stub.streamErr
}

func (stub *clientStub) ReportObservations(_ context.Context, batch *controlplanev1.ObservationBatch) (*controlplanev1.ReportObservationsResponse, error) {
	stub.batches = append(stub.batches, batch)
	return &controlplanev1.ReportObservationsResponse{Accepted: uint32(len(batch.GetObservations()))}, nil
}

type taskStreamStub struct {
	tasks []*controlplanev1.CheckTask
	next  int
}

func (stub *taskStreamStub) Recv() (*controlplanev1.CheckTask, error) {
	if stub.next >= len(stub.tasks) {
		return nil, io.EOF
	}
	value := stub.tasks[stub.next]
	stub.next++
	return value, nil
}

type executorStub struct {
	mu    sync.Mutex
	calls map[string]int
	tasks []*controlplanev1.CheckTask
}

func (stub *executorStub) Execute(_ context.Context, task *controlplanev1.CheckTask) probe.Result {
	stub.mu.Lock()
	defer stub.mu.Unlock()
	stub.tasks = append(stub.tasks, task)
	if stub.calls == nil {
		stub.calls = make(map[string]int)
	}
	stub.calls[task.GetTaskId()]++
	if task.GetTaskId() == "task-a" && stub.calls[task.GetTaskId()] == 1 {
		return probe.Result{FailureClass: probe.FailureProxyRequest, Latency: time.Millisecond}
	}
	result := probe.Result{Success: true, Latency: 2 * time.Millisecond}
	if task.GetLevel() == controlplanev1.CheckLevel_CHECK_LEVEL_EGRESS {
		result.ObservedEgressIP = "198.51.100.42"
	}
	return result
}