129 lines
4.7 KiB
Go
129 lines
4.7 KiB
Go
package controlplane
|
|
|
|
import (
|
|
"context"
|
|
"io"
|
|
"sync"
|
|
"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 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
|
|
batches []*controlplanev1.ObservationBatch
|
|
}
|
|
|
|
func (stub *clientStub) StreamCheckTasks(context.Context, *controlplanev1.StreamCheckTasksRequest) (TaskStream, error) {
|
|
return stub.stream, nil
|
|
}
|
|
|
|
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
|
|
}
|