315 lines
13 KiB
Go
315 lines
13 KiB
Go
package health
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"net"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
controlplanev1 "proxy-pool/gen/controlplane/v1"
|
|
healthDomain "proxy-pool/internal/domain/health"
|
|
proxyDomain "proxy-pool/internal/domain/proxy"
|
|
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/codes"
|
|
"google.golang.org/grpc/credentials/insecure"
|
|
"google.golang.org/grpc/status"
|
|
"google.golang.org/grpc/test/bufconn"
|
|
"google.golang.org/protobuf/types/known/durationpb"
|
|
"google.golang.org/protobuf/types/known/timestamppb"
|
|
)
|
|
|
|
func TestGRPCHandlerReportsAcceptedAndRejectedObservations(t *testing.T) {
|
|
global := &recordingGlobalStore{}
|
|
target := &recordingTargetStore{}
|
|
reducer, err := NewReducer(global, target, func(context.Context, healthDomain.Observation) (int, error) { return 2, nil })
|
|
if err != nil {
|
|
t.Fatalf("NewReducer(): %v", err)
|
|
}
|
|
identity := &recordingCheckerIdentity{}
|
|
handler, err := NewGRPCHandler(reducer, identity, GRPCHandlerOptions{MaxObservationsPerBatch: 3})
|
|
if err != nil {
|
|
t.Fatalf("NewGRPCHandler(): %v", err)
|
|
}
|
|
now := time.Date(2026, 7, 31, 14, 0, 0, 0, time.UTC)
|
|
response, err := handler.ReportObservations(context.Background(), &controlplanev1.ObservationBatch{
|
|
CheckerId: "checker-a",
|
|
Observations: []*controlplanev1.HealthObservation{
|
|
grpcHealthObservation("task-basic", controlplanev1.CheckLevel_CHECK_LEVEL_BASIC, now),
|
|
grpcHealthObservation("task-target", controlplanev1.CheckLevel_CHECK_LEVEL_TARGET, now.Add(time.Second)),
|
|
{TaskId: "invalid", ProxyId: "proxy-a", Level: controlplanev1.CheckLevel_CHECK_LEVEL_UNSPECIFIED},
|
|
},
|
|
})
|
|
if err != nil || response.GetAccepted() != 2 || response.GetRejected() != 1 {
|
|
t.Fatalf("ReportObservations() = (%+v, %v)", response, err)
|
|
}
|
|
if identity.checkerID != "checker-a" || len(global.commands) != 1 || len(target.commands) != 1 {
|
|
t.Fatalf("identity=%q global=%d target=%d", identity.checkerID, len(global.commands), len(target.commands))
|
|
}
|
|
}
|
|
|
|
func TestGRPCHandlerRejectsInvalidCallsBeforeStoreMutation(t *testing.T) {
|
|
global := &recordingGlobalStore{}
|
|
target := &recordingTargetStore{}
|
|
reducer, err := NewReducer(global, target, func(context.Context, healthDomain.Observation) (int, error) { return 2, nil })
|
|
if err != nil {
|
|
t.Fatalf("NewReducer(): %v", err)
|
|
}
|
|
identity := &recordingCheckerIdentity{err: errors.New("not allowed")}
|
|
handler, err := NewGRPCHandler(reducer, identity, GRPCHandlerOptions{MaxObservationsPerBatch: 1})
|
|
if err != nil {
|
|
t.Fatalf("NewGRPCHandler(): %v", err)
|
|
}
|
|
_, err = handler.ReportObservations(context.Background(), &controlplanev1.ObservationBatch{
|
|
CheckerId: "checker-a", Observations: []*controlplanev1.HealthObservation{
|
|
grpcHealthObservation("task-basic", controlplanev1.CheckLevel_CHECK_LEVEL_BASIC, time.Now()),
|
|
},
|
|
})
|
|
if status.Code(err) != codes.PermissionDenied || len(global.commands) != 0 || len(target.commands) != 0 {
|
|
t.Fatalf("ReportObservations(denied) = %v; global=%d target=%d", err, len(global.commands), len(target.commands))
|
|
}
|
|
_, err = handler.ReportObservations(context.Background(), &controlplanev1.ObservationBatch{
|
|
CheckerId: "checker-a", Observations: []*controlplanev1.HealthObservation{
|
|
grpcHealthObservation("task-one", controlplanev1.CheckLevel_CHECK_LEVEL_BASIC, time.Now()),
|
|
grpcHealthObservation("task-two", controlplanev1.CheckLevel_CHECK_LEVEL_BASIC, time.Now()),
|
|
},
|
|
})
|
|
if status.Code(err) != codes.InvalidArgument {
|
|
t.Fatalf("ReportObservations(too large) = %v, want InvalidArgument", err)
|
|
}
|
|
}
|
|
|
|
func TestGRPCHandlerReturnsUnavailableForStoreFailure(t *testing.T) {
|
|
global := &recordingGlobalStore{err: errors.New("redis unavailable")}
|
|
target := &recordingTargetStore{}
|
|
reducer, err := NewReducer(global, target, func(context.Context, healthDomain.Observation) (int, error) { return 2, nil })
|
|
if err != nil {
|
|
t.Fatalf("NewReducer(): %v", err)
|
|
}
|
|
handler, err := NewGRPCHandler(reducer, &recordingCheckerIdentity{}, GRPCHandlerOptions{MaxObservationsPerBatch: 1})
|
|
if err != nil {
|
|
t.Fatalf("NewGRPCHandler(): %v", err)
|
|
}
|
|
_, err = handler.ReportObservations(context.Background(), &controlplanev1.ObservationBatch{
|
|
CheckerId: "checker-a", Observations: []*controlplanev1.HealthObservation{
|
|
grpcHealthObservation("task-basic", controlplanev1.CheckLevel_CHECK_LEVEL_BASIC, time.Now()),
|
|
},
|
|
})
|
|
if status.Code(err) != codes.Unavailable {
|
|
t.Fatalf("ReportObservations(store failure) = %v, want Unavailable", err)
|
|
}
|
|
}
|
|
|
|
func TestGRPCHandlerStreamsLeasedTasksAndFencesReportedFacts(t *testing.T) {
|
|
now := time.Date(2026, 8, 1, 12, 0, 0, 0, time.UTC)
|
|
broker, err := NewMemoryTaskBroker(MemoryTaskBrokerOptions{
|
|
LeaseTTL: time.Minute, Now: func() time.Time { return now },
|
|
Material: TaskMaterialResolverFunc(func(context.Context, Candidate) (TaskMaterial, error) {
|
|
return TaskMaterial{
|
|
Protocol: proxyDomain.SchemeHTTP, Host: "proxy.example", Port: 8080,
|
|
SecretRef: "cred-a", CredentialVersion: "v1", Username: "user", Password: "secret",
|
|
}, nil
|
|
}),
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewMemoryTaskBroker(): %v", err)
|
|
}
|
|
if offered, offerErr := broker.Offer(context.Background(), []PlannedTask{{
|
|
Candidate: Candidate{ProxyID: "proxy-a", State: proxyDomain.StateFetched, Level: healthDomain.LevelBasic, DueAt: now},
|
|
Deadline: now.Add(10 * time.Second), Attempts: 2,
|
|
}}); offerErr != nil || offered != 1 {
|
|
t.Fatalf("Offer() = (%d, %v)", offered, offerErr)
|
|
}
|
|
global := &recordingGlobalStore{}
|
|
reducer, err := NewReducer(global, &recordingTargetStore{}, func(context.Context, healthDomain.Observation) (int, error) { return 2, nil })
|
|
if err != nil {
|
|
t.Fatalf("NewReducer(): %v", err)
|
|
}
|
|
metrics := &recordingTaskMetrics{}
|
|
handler, err := NewGRPCHandler(reducer, &recordingCheckerIdentity{}, GRPCHandlerOptions{
|
|
MaxObservationsPerBatch: 2, MaxTasksPerClaim: 2, TaskBroker: broker, Metrics: metrics, Now: func() time.Time { return now },
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewGRPCHandler(): %v", err)
|
|
}
|
|
client, closeClient := newCheckerGRPCClient(t, handler)
|
|
defer closeClient()
|
|
|
|
stream, err := client.StreamCheckTasks(context.Background(), &controlplanev1.StreamCheckTasksRequest{
|
|
CheckerId: "checker-a", InstanceId: "instance-a", MaxInFlight: 2,
|
|
SupportedLevels: []controlplanev1.CheckLevel{controlplanev1.CheckLevel_CHECK_LEVEL_BASIC},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("StreamCheckTasks(): %v", err)
|
|
}
|
|
task, err := stream.Recv()
|
|
if err != nil || task.GetProxyId() != "proxy-a" || task.GetPassword() != "secret" || task.GetCredentialVersion() != "v1" ||
|
|
task.GetMaxAttempts() != 2 || task.GetAttempt() != 1 || task.GetTimeout().AsDuration() != 10*time.Second {
|
|
t.Fatalf("Recv() = (%+v, %v)", task, err)
|
|
}
|
|
if _, err := stream.Recv(); !errors.Is(err, io.EOF) {
|
|
t.Fatalf("second Recv() error = %v, want EOF", err)
|
|
}
|
|
acceptedObservation := grpcHealthObservation(task.GetTaskId(), controlplanev1.CheckLevel_CHECK_LEVEL_BASIC, now.Add(time.Second))
|
|
acceptedObservation.LeaseToken = task.GetLeaseToken()
|
|
response, err := client.ReportObservations(context.Background(), &controlplanev1.ObservationBatch{
|
|
CheckerId: "checker-a", Observations: []*controlplanev1.HealthObservation{
|
|
acceptedObservation,
|
|
},
|
|
})
|
|
if err != nil || response.GetAccepted() != 1 || len(global.commands) != 1 {
|
|
t.Fatalf("ReportObservations() = (%+v, %v), global=%d", response, err, len(global.commands))
|
|
}
|
|
rejectedObservation := grpcHealthObservation(task.GetTaskId(), controlplanev1.CheckLevel_CHECK_LEVEL_BASIC, now.Add(time.Second))
|
|
rejectedObservation.LeaseToken = task.GetLeaseToken()
|
|
rejected, err := client.ReportObservations(context.Background(), &controlplanev1.ObservationBatch{
|
|
CheckerId: "checker-b", Observations: []*controlplanev1.HealthObservation{
|
|
rejectedObservation,
|
|
},
|
|
})
|
|
if err != nil || rejected.GetRejected() != 1 || len(global.commands) != 1 {
|
|
t.Fatalf("ReportObservations(other checker) = (%+v, %v), global=%d", rejected, err, len(global.commands))
|
|
}
|
|
if metrics.dispatches(healthDomain.LevelBasic) != 1 ||
|
|
metrics.observations(healthDomain.LevelBasic, healthDomain.ObservationMetricAccepted) != 1 ||
|
|
metrics.observations(healthDomain.LevelBasic, healthDomain.ObservationMetricRejected) != 1 {
|
|
t.Fatalf("Checker metrics = %+v", metrics)
|
|
}
|
|
}
|
|
|
|
func TestGRPCHandlerStreamsEgressProbeURLAndAcceptsGlobalFact(t *testing.T) {
|
|
now := time.Date(2026, 8, 1, 12, 30, 0, 0, time.UTC)
|
|
broker, err := NewMemoryTaskBroker(MemoryTaskBrokerOptions{
|
|
LeaseTTL: time.Minute, Now: func() time.Time { return now },
|
|
Material: TaskMaterialResolverFunc(func(context.Context, Candidate) (TaskMaterial, error) {
|
|
return TaskMaterial{Protocol: proxyDomain.SchemeHTTP, Host: "proxy.example", Port: 8080}, nil
|
|
}),
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewMemoryTaskBroker(): %v", err)
|
|
}
|
|
const egressURL = "https://egress.example/identity"
|
|
if offered, offerErr := broker.Offer(context.Background(), []PlannedTask{{
|
|
Candidate: Candidate{ProxyID: "proxy-a", State: proxyDomain.StateAvailable, Level: healthDomain.LevelEgress, TargetURL: egressURL, DueAt: now},
|
|
Deadline: now.Add(10 * time.Second), Attempts: 1,
|
|
}}); offerErr != nil || offered != 1 {
|
|
t.Fatalf("Offer() = (%d, %v)", offered, offerErr)
|
|
}
|
|
global := &recordingGlobalStore{}
|
|
reducer, err := NewReducer(global, &recordingTargetStore{}, func(context.Context, healthDomain.Observation) (int, error) { return 2, nil })
|
|
if err != nil {
|
|
t.Fatalf("NewReducer(): %v", err)
|
|
}
|
|
handler, err := NewGRPCHandler(reducer, &recordingCheckerIdentity{}, GRPCHandlerOptions{
|
|
MaxObservationsPerBatch: 1, MaxTasksPerClaim: 1, TaskBroker: broker, Now: func() time.Time { return now },
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewGRPCHandler(): %v", err)
|
|
}
|
|
client, closeClient := newCheckerGRPCClient(t, handler)
|
|
defer closeClient()
|
|
|
|
stream, err := client.StreamCheckTasks(context.Background(), &controlplanev1.StreamCheckTasksRequest{
|
|
CheckerId: "checker-a", InstanceId: "instance-a", MaxInFlight: 1,
|
|
SupportedLevels: []controlplanev1.CheckLevel{controlplanev1.CheckLevel_CHECK_LEVEL_EGRESS},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("StreamCheckTasks(): %v", err)
|
|
}
|
|
task, err := stream.Recv()
|
|
if err != nil || task.GetLevel() != controlplanev1.CheckLevel_CHECK_LEVEL_EGRESS || task.GetTargetUrl() != egressURL || task.GetRoutingName() != "" {
|
|
t.Fatalf("Recv() = (%+v, %v)", task, err)
|
|
}
|
|
observation := grpcHealthObservation(task.GetTaskId(), controlplanev1.CheckLevel_CHECK_LEVEL_EGRESS, now.Add(time.Second))
|
|
observation.LeaseToken = task.GetLeaseToken()
|
|
response, err := client.ReportObservations(context.Background(), &controlplanev1.ObservationBatch{
|
|
CheckerId: "checker-a", Observations: []*controlplanev1.HealthObservation{observation},
|
|
})
|
|
if err != nil || response.GetAccepted() != 1 || len(global.commands) != 1 ||
|
|
global.commands[0].Observation.TargetURL != "" || global.commands[0].Observation.RoutingName != "" {
|
|
t.Fatalf("ReportObservations() = (%+v, %v), commands=%+v", response, err, global.commands)
|
|
}
|
|
}
|
|
|
|
func newCheckerGRPCClient(t *testing.T, handler controlplanev1.CheckerControlPlaneServer) (controlplanev1.CheckerControlPlaneClient, func()) {
|
|
t.Helper()
|
|
listener := bufconn.Listen(1 << 20)
|
|
server := grpc.NewServer()
|
|
controlplanev1.RegisterCheckerControlPlaneServer(server, handler)
|
|
go func() { _ = server.Serve(listener) }()
|
|
connection, err := grpc.NewClient("passthrough:///checker-test", grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) {
|
|
return listener.Dial()
|
|
}), grpc.WithTransportCredentials(insecure.NewCredentials()))
|
|
if err != nil {
|
|
server.Stop()
|
|
_ = listener.Close()
|
|
t.Fatalf("grpc.NewClient(): %v", err)
|
|
}
|
|
return controlplanev1.NewCheckerControlPlaneClient(connection), func() {
|
|
_ = connection.Close()
|
|
server.Stop()
|
|
_ = listener.Close()
|
|
}
|
|
}
|
|
|
|
func grpcHealthObservation(taskID string, level controlplanev1.CheckLevel, observedAt time.Time) *controlplanev1.HealthObservation {
|
|
item := &controlplanev1.HealthObservation{
|
|
TaskId: taskID, ProxyId: "proxy-a", Level: level, Success: true,
|
|
Latency: durationpb.New(time.Millisecond), ObservedAt: timestamppb.New(observedAt),
|
|
}
|
|
if level == controlplanev1.CheckLevel_CHECK_LEVEL_TARGET {
|
|
item.RoutingName = "route-a"
|
|
item.TargetUrl = "https://target.example/check"
|
|
}
|
|
return item
|
|
}
|
|
|
|
type recordingCheckerIdentity struct {
|
|
checkerID string
|
|
err error
|
|
}
|
|
|
|
type recordingTaskMetrics struct {
|
|
mu sync.Mutex
|
|
dispatch map[healthDomain.Level]int
|
|
observation map[string]int
|
|
}
|
|
|
|
func (metrics *recordingTaskMetrics) ObserveTaskDispatch(level healthDomain.Level, count int) {
|
|
metrics.mu.Lock()
|
|
defer metrics.mu.Unlock()
|
|
if metrics.dispatch == nil {
|
|
metrics.dispatch = make(map[healthDomain.Level]int)
|
|
}
|
|
metrics.dispatch[level] += count
|
|
}
|
|
|
|
func (metrics *recordingTaskMetrics) ObserveObservation(level healthDomain.Level, result healthDomain.ObservationMetricResult) {
|
|
metrics.mu.Lock()
|
|
defer metrics.mu.Unlock()
|
|
if metrics.observation == nil {
|
|
metrics.observation = make(map[string]int)
|
|
}
|
|
metrics.observation[string(level)+"\x00"+string(result)]++
|
|
}
|
|
|
|
func (metrics *recordingTaskMetrics) dispatches(level healthDomain.Level) int {
|
|
metrics.mu.Lock()
|
|
defer metrics.mu.Unlock()
|
|
return metrics.dispatch[level]
|
|
}
|
|
|
|
func (metrics *recordingTaskMetrics) observations(level healthDomain.Level, result healthDomain.ObservationMetricResult) int {
|
|
metrics.mu.Lock()
|
|
defer metrics.mu.Unlock()
|
|
return metrics.observation[string(level)+"\x00"+string(result)]
|
|
}
|
|
|
|
func (identity *recordingCheckerIdentity) AuthorizeChecker(_ context.Context, checkerID string) error {
|
|
identity.checkerID = checkerID
|
|
return identity.err
|
|
}
|