105 lines
4.3 KiB
Go
105 lines
4.3 KiB
Go
package worker
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
|
|
"proxy-pool/internal/domain/workerruntime"
|
|
)
|
|
|
|
func TestServiceRegistersAcknowledgesAndReportsRuntime(t *testing.T) {
|
|
now := time.Date(2026, 7, 31, 10, 0, 0, 0, time.UTC)
|
|
store, err := workerruntime.NewMemoryStore(func() time.Time { return now })
|
|
if err != nil {
|
|
t.Fatalf("NewMemoryStore(): %v", err)
|
|
}
|
|
service, err := NewService(store, Options{
|
|
ProtocolVersion: 1, HeartbeatInterval: 10 * time.Second, SessionTTL: 30 * time.Second,
|
|
MaxStaleAge: 10 * time.Second, MaxRuntimeCounters: 100,
|
|
SessionID: func() (string, error) { return "0123456789abcdef0123456789abcdef", nil },
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewService(): %v", err)
|
|
}
|
|
registered, err := service.Register(context.Background(), RegisterCommand{
|
|
WorkerID: "worker-a", InstanceID: "instance-a", Zone: "zone-a", ProtocolVersion: 1,
|
|
Labels: map[string]string{"region": "test"},
|
|
})
|
|
if err != nil || registered.SessionID == "" || registered.OwnershipEpoch == 0 {
|
|
t.Fatalf("Register() = %+v, %v", registered, err)
|
|
}
|
|
reference := workerruntime.SnapshotReference{
|
|
WorkerID: "worker-a", Version: 7, OwnershipEpoch: registered.OwnershipEpoch,
|
|
Checksum: sha256.Sum256([]byte("snapshot-7")),
|
|
}
|
|
if err := service.IssueSnapshot(context.Background(), registered.SessionID, reference); err != nil {
|
|
t.Fatalf("IssueSnapshot(): %v", err)
|
|
}
|
|
if err := service.Acknowledge(context.Background(), SnapshotAcknowledgement{
|
|
WorkerID: "worker-a", SessionID: registered.SessionID, Version: 7,
|
|
OwnershipEpoch: registered.OwnershipEpoch, Checksum: reference.Checksum[:], Applied: true,
|
|
}); err != nil {
|
|
t.Fatalf("Acknowledge(): %v", err)
|
|
}
|
|
decision, err := service.ReportRuntime(context.Background(), workerruntime.Report{
|
|
WorkerID: "worker-a", SessionID: registered.SessionID, Sequence: 1,
|
|
SnapshotVersion: 7, OwnershipEpoch: registered.OwnershipEpoch, ObservedAt: now,
|
|
})
|
|
if err != nil || decision.RequireFullSnapshot || decision.AcceptedOwnershipEpoch != registered.OwnershipEpoch {
|
|
t.Fatalf("ReportRuntime() = %+v, %v", decision, err)
|
|
}
|
|
decision, err = service.ReportRuntime(context.Background(), workerruntime.Report{
|
|
WorkerID: "worker-a", SessionID: registered.SessionID, Sequence: 2,
|
|
SnapshotVersion: 8, OwnershipEpoch: registered.OwnershipEpoch, ObservedAt: now,
|
|
})
|
|
if err != nil || !decision.RequireFullSnapshot || decision.AcceptedOwnershipEpoch != registered.OwnershipEpoch {
|
|
t.Fatalf("ReportRuntime(mismatch) = %+v, %v", decision, err)
|
|
}
|
|
}
|
|
|
|
func TestServiceRejectsInvalidCommandsAndUnavailableStore(t *testing.T) {
|
|
store := unavailableStore{}
|
|
service, err := NewService(store, Options{
|
|
ProtocolVersion: 1, HeartbeatInterval: time.Second, SessionTTL: 3 * time.Second,
|
|
MaxStaleAge: time.Second, MaxRuntimeCounters: 1,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewService(): %v", err)
|
|
}
|
|
if _, err := service.Register(context.Background(), RegisterCommand{ProtocolVersion: 2}); !errors.Is(err, ErrProtocolVersion) {
|
|
t.Fatalf("Register(protocol) error = %v", err)
|
|
}
|
|
if _, err := service.Register(context.Background(), RegisterCommand{
|
|
WorkerID: "worker-a", InstanceID: "instance-a", Zone: "zone-a", ProtocolVersion: 1,
|
|
}); !errors.Is(err, ErrUnavailable) {
|
|
t.Fatalf("Register(unavailable) error = %v", err)
|
|
}
|
|
if err := service.Acknowledge(context.Background(), SnapshotAcknowledgement{Checksum: make([]byte, 31)}); !errors.Is(err, ErrInvalidCommand) {
|
|
t.Fatalf("Acknowledge(invalid) error = %v", err)
|
|
}
|
|
}
|
|
|
|
type unavailableStore struct{}
|
|
|
|
func (unavailableStore) CurrentOwnershipEpoch(context.Context) (uint64, error) {
|
|
return 0, errors.New("redis unavailable")
|
|
}
|
|
func (unavailableStore) OpenSession(context.Context, workerruntime.Session, time.Duration) error {
|
|
return errors.New("redis unavailable")
|
|
}
|
|
func (unavailableStore) ValidateSession(context.Context, string, string) error {
|
|
return errors.New("redis unavailable")
|
|
}
|
|
func (unavailableStore) RecordIssuedSnapshot(context.Context, string, workerruntime.SnapshotReference, time.Duration) error {
|
|
return errors.New("redis unavailable")
|
|
}
|
|
func (unavailableStore) AcknowledgeSnapshot(context.Context, workerruntime.SnapshotAcknowledgement, time.Duration) error {
|
|
return errors.New("redis unavailable")
|
|
}
|
|
func (unavailableStore) ReplaceRuntime(context.Context, workerruntime.Report, time.Duration) error {
|
|
return errors.New("redis unavailable")
|
|
}
|