110 lines
3.8 KiB
Go
110 lines
3.8 KiB
Go
package controlplane
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
controlplanev1 "proxy-pool/gen/controlplane/v1"
|
|
"proxy-pool/internal/gateway/snapshot"
|
|
|
|
"google.golang.org/protobuf/types/known/durationpb"
|
|
"google.golang.org/protobuf/types/known/emptypb"
|
|
"google.golang.org/protobuf/types/known/timestamppb"
|
|
)
|
|
|
|
func TestSessionRunnerRegistersOnceWatchesAndReports(t *testing.T) {
|
|
store := snapshot.NewStore("cluster-a", "worker-a")
|
|
full := &controlplanev1.WorkerSnapshot{
|
|
Version: 1, OwnershipEpoch: 7, GeneratedAt: timestamppb.New(time.Now()), ValidUntil: timestamppb.New(time.Now().Add(time.Minute)),
|
|
}
|
|
setSnapshotChecksum(t, full)
|
|
runtime := make(chan *controlplanev1.ReportRuntimeRequest, 8)
|
|
reporterClient := &sessionReporterClient{runtime: runtime}
|
|
reporter, err := NewRuntimeReporter(reporterClient, store, Options{
|
|
WorkerID: "worker-a", InstanceID: "instance-a", Zone: "zone-a", ProtocolVersion: 1, Now: time.Now,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewRuntimeReporter(): %v", err)
|
|
}
|
|
watcherClient := &sessionWatcherClient{full: full}
|
|
watcher, err := NewSnapshotWatcher(watcherClient, store, SnapshotWatcherOptions{ClusterID: "cluster-a", WorkerID: "worker-a"})
|
|
if err != nil {
|
|
t.Fatalf("NewSnapshotWatcher(): %v", err)
|
|
}
|
|
runner, err := NewSessionRunner(reporter, watcher)
|
|
if err != nil {
|
|
t.Fatalf("NewSessionRunner(): %v", err)
|
|
}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
result := make(chan error, 1)
|
|
go func() { result <- runner.Run(ctx) }()
|
|
select {
|
|
case request := <-runtime:
|
|
if request.GetSessionId() != "session-a" || request.GetSnapshotVersion() != 1 || request.GetReportSequence() != 1 {
|
|
t.Fatalf("runtime request = %+v", request)
|
|
}
|
|
case <-time.After(time.Second):
|
|
t.Fatal("runtime report was not sent")
|
|
}
|
|
cancel()
|
|
if err := <-result; !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("Run() error = %v, want context canceled", err)
|
|
}
|
|
if reporterClient.registers != 1 || watcherClient.ack == nil || !watcherClient.ack.GetApplied() {
|
|
t.Fatalf("registers=%d ack=%+v", reporterClient.registers, watcherClient.ack)
|
|
}
|
|
}
|
|
|
|
type sessionReporterClient struct {
|
|
mu sync.Mutex
|
|
registers int
|
|
runtime chan *controlplanev1.ReportRuntimeRequest
|
|
}
|
|
|
|
func (client *sessionReporterClient) RegisterWorker(context.Context, *controlplanev1.RegisterWorkerRequest) (*controlplanev1.RegisterWorkerResponse, error) {
|
|
client.mu.Lock()
|
|
client.registers++
|
|
client.mu.Unlock()
|
|
return &controlplanev1.RegisterWorkerResponse{
|
|
WorkerId: "worker-a", SessionId: "session-a", OwnershipEpoch: 7,
|
|
HeartbeatInterval: durationpb.New(5 * time.Millisecond), MaxStaleAge: durationpb.New(time.Second),
|
|
}, nil
|
|
}
|
|
|
|
func (client *sessionReporterClient) ReportRuntime(_ context.Context, request *controlplanev1.ReportRuntimeRequest) (*controlplanev1.ReportRuntimeResponse, error) {
|
|
client.runtime <- request
|
|
return &controlplanev1.ReportRuntimeResponse{AcceptedOwnershipEpoch: 7}, nil
|
|
}
|
|
|
|
type sessionWatcherClient struct {
|
|
full *controlplanev1.WorkerSnapshot
|
|
ack *controlplanev1.AcknowledgeSnapshotRequest
|
|
}
|
|
|
|
func (client *sessionWatcherClient) Watch(ctx context.Context, _ *controlplanev1.WatchSnapshotsRequest) (SnapshotStream, error) {
|
|
return &sessionSnapshotStream{ctx: ctx, full: client.full}, nil
|
|
}
|
|
|
|
func (client *sessionWatcherClient) Acknowledge(_ context.Context, acknowledgement *controlplanev1.AcknowledgeSnapshotRequest) (*emptypb.Empty, error) {
|
|
client.ack = acknowledgement
|
|
return &emptypb.Empty{}, nil
|
|
}
|
|
|
|
type sessionSnapshotStream struct {
|
|
ctx context.Context
|
|
full *controlplanev1.WorkerSnapshot
|
|
sent bool
|
|
}
|
|
|
|
func (stream *sessionSnapshotStream) Recv() (*controlplanev1.SnapshotEnvelope, error) {
|
|
if !stream.sent {
|
|
stream.sent = true
|
|
return &controlplanev1.SnapshotEnvelope{Payload: &controlplanev1.SnapshotEnvelope_Full{Full: stream.full}}, nil
|
|
}
|
|
<-stream.ctx.Done()
|
|
return nil, stream.ctx.Err()
|
|
}
|