104 lines
4.5 KiB
Go
104 lines
4.5 KiB
Go
package worker
|
|
|
|
import (
|
|
"context"
|
|
"net"
|
|
"testing"
|
|
"time"
|
|
|
|
controlplanev1 "proxy-pool/gen/controlplane/v1"
|
|
"proxy-pool/internal/domain/workerruntime"
|
|
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/codes"
|
|
"google.golang.org/grpc/status"
|
|
"google.golang.org/grpc/test/bufconn"
|
|
"google.golang.org/protobuf/types/known/timestamppb"
|
|
)
|
|
|
|
func TestGRPCHandlerMapsWorkerRequests(t *testing.T) {
|
|
service := &grpcServiceStub{registration: Registration{WorkerID: "worker-a", SessionID: "session-a", OwnershipEpoch: 9, HeartbeatInterval: time.Second, MaxStaleAge: 2 * time.Second}}
|
|
client, cleanup := grpcWorkerClient(t, service, allowIdentity{})
|
|
defer cleanup()
|
|
registered, err := client.RegisterWorker(context.Background(), &controlplanev1.RegisterWorkerRequest{
|
|
WorkerId: "worker-a", InstanceId: "instance-a", Zone: "zone-a", SupportedProtocolVersion: 1,
|
|
})
|
|
if err != nil || registered.GetSessionId() != "session-a" || registered.GetHeartbeatInterval().AsDuration() != time.Second {
|
|
t.Fatalf("RegisterWorker() = %+v, %v", registered, err)
|
|
}
|
|
if err := service.acknowledgeErr; err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
_, err = client.AcknowledgeSnapshot(context.Background(), &controlplanev1.AcknowledgeSnapshotRequest{
|
|
WorkerId: "worker-a", SessionId: "session-a", Version: 7, OwnershipEpoch: 9, Checksum: make([]byte, 32),
|
|
})
|
|
if err != nil || service.acknowledgement.Version != 7 {
|
|
t.Fatalf("AcknowledgeSnapshot() error = %v; command=%+v", err, service.acknowledgement)
|
|
}
|
|
response, err := client.ReportRuntime(context.Background(), &controlplanev1.ReportRuntimeRequest{
|
|
WorkerId: "worker-a", SessionId: "session-a", SnapshotVersion: 7, OwnershipEpoch: 9, ReportSequence: 1,
|
|
ObservedAt: timestamppb.New(time.Now()), Counters: []*controlplanev1.ProxyRuntime{{ProxyId: "proxy-a", Active: 2, Reserved: 1}},
|
|
})
|
|
if err != nil || response.GetAcceptedOwnershipEpoch() != 9 || service.report.Counters[0].Active != 2 {
|
|
t.Fatalf("ReportRuntime() = %+v, %v; report=%+v", response, err, service.report)
|
|
}
|
|
}
|
|
|
|
func TestGRPCHandlerMapsErrorsAndLeavesStreamsUnimplemented(t *testing.T) {
|
|
service := &grpcServiceStub{registerErr: ErrProtocolVersion}
|
|
client, cleanup := grpcWorkerClient(t, service, allowIdentity{})
|
|
defer cleanup()
|
|
_, err := client.RegisterWorker(context.Background(), &controlplanev1.RegisterWorkerRequest{WorkerId: "worker-a"})
|
|
if status.Code(err) != codes.FailedPrecondition {
|
|
t.Fatalf("RegisterWorker() code = %s, want FailedPrecondition", status.Code(err))
|
|
}
|
|
stream, err := client.WatchSnapshots(context.Background(), &controlplanev1.WatchSnapshotsRequest{})
|
|
_, streamErr := stream.Recv()
|
|
if err != nil || status.Code(streamErr) != codes.Unimplemented {
|
|
t.Fatalf("WatchSnapshots() = %v, %v", err, streamErr)
|
|
}
|
|
outcomes, err := client.ReportOutcomes(context.Background())
|
|
_, outcomesErr := outcomes.CloseAndRecv()
|
|
if err != nil || status.Code(outcomesErr) != codes.Unimplemented {
|
|
t.Fatalf("ReportOutcomes() = %v, %v", err, outcomesErr)
|
|
}
|
|
}
|
|
|
|
type grpcServiceStub struct {
|
|
registration Registration
|
|
registerErr error
|
|
acknowledgement SnapshotAcknowledgement
|
|
acknowledgeErr error
|
|
report workerruntime.Report
|
|
reportErr error
|
|
}
|
|
|
|
func (stub *grpcServiceStub) Register(context.Context, RegisterCommand) (Registration, error) {
|
|
return stub.registration, stub.registerErr
|
|
}
|
|
func (stub *grpcServiceStub) Acknowledge(_ context.Context, acknowledgement SnapshotAcknowledgement) error {
|
|
stub.acknowledgement = acknowledgement
|
|
return stub.acknowledgeErr
|
|
}
|
|
func (stub *grpcServiceStub) ReportRuntime(_ context.Context, report workerruntime.Report) (RuntimeDecision, error) {
|
|
stub.report = report
|
|
return RuntimeDecision{AcceptedOwnershipEpoch: 9}, stub.reportErr
|
|
}
|
|
|
|
type allowIdentity struct{}
|
|
|
|
func (allowIdentity) Authorize(context.Context, string) error { return nil }
|
|
|
|
func grpcWorkerClient(t *testing.T, service Service, identity IdentityAuthorizer) (controlplanev1.WorkerControlPlaneClient, func()) {
|
|
t.Helper()
|
|
listener := bufconn.Listen(1 << 20)
|
|
server := grpc.NewServer()
|
|
controlplanev1.RegisterWorkerControlPlaneServer(server, NewGRPCHandler(service, identity))
|
|
go func() { _ = server.Serve(listener) }()
|
|
connection, err := grpc.NewClient("passthrough:///bufnet", grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return listener.Dial() }), grpc.WithInsecure())
|
|
if err != nil {
|
|
t.Fatalf("grpc.NewClient(): %v", err)
|
|
}
|
|
return controlplanev1.NewWorkerControlPlaneClient(connection), func() { _ = connection.Close(); server.Stop(); _ = listener.Close() }
|
|
}
|