479 lines
14 KiB
Go
479 lines
14 KiB
Go
package runnerregistry
|
|
|
|
import (
|
|
"testing"
|
|
"time"
|
|
|
|
otopb "github.com/toki/oto/services/core/oto"
|
|
)
|
|
|
|
func TestRegistryAcceptsRunnerRegistration(t *testing.T) {
|
|
acceptedAt := time.Date(2026, 6, 5, 12, 0, 0, 0, time.UTC)
|
|
registry := NewWithClock(func() time.Time { return acceptedAt })
|
|
|
|
response := registry.Register(&otopb.RegisterRunnerRequest{
|
|
EnrollmentToken: "token-123",
|
|
RunnerId: "runner-123",
|
|
Alias: "build-linux",
|
|
ProtocolVersion: "oto.runner.v1",
|
|
Capability: &otopb.RunnerCapability{
|
|
Name: "oto-runner",
|
|
Version: "1.0.0",
|
|
},
|
|
CommandCatalog: &otopb.CommandCatalogSummary{
|
|
CommandTypes: []string{"Shell", "Git"},
|
|
},
|
|
})
|
|
|
|
if !response.GetAccepted() {
|
|
t.Fatalf("Register accepted=false, reason=%q", response.GetRejectReason())
|
|
}
|
|
if response.GetRunnerId() != "runner-123" {
|
|
t.Fatalf("response runner id = %q, want runner-123", response.GetRunnerId())
|
|
}
|
|
|
|
record, ok := registry.Snapshot("runner-123")
|
|
if !ok {
|
|
t.Fatal("registered runner was not stored")
|
|
}
|
|
if record.Alias != "build-linux" {
|
|
t.Fatalf("record alias = %q, want build-linux", record.Alias)
|
|
}
|
|
if record.ProtocolVersion != "oto.runner.v1" {
|
|
t.Fatalf("record protocol = %q, want oto.runner.v1", record.ProtocolVersion)
|
|
}
|
|
if record.Capability.Name != "oto-runner" || record.Capability.Version != "1.0.0" {
|
|
t.Fatalf("record capability = %#v", record.Capability)
|
|
}
|
|
if record.Status != StatusAccepted {
|
|
t.Fatalf("record status = %q, want %q", record.Status, StatusAccepted)
|
|
}
|
|
if !record.AcceptedAt.Equal(acceptedAt) {
|
|
t.Fatalf("record accepted at = %v, want %v", record.AcceptedAt, acceptedAt)
|
|
}
|
|
if got := record.CommandTypes; len(got) != 2 || got[0] != "Shell" || got[1] != "Git" {
|
|
t.Fatalf("record command types = %#v, want [Shell Git]", got)
|
|
}
|
|
}
|
|
|
|
func TestRegistryRejectsIncompleteRegistration(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
req *otopb.RegisterRunnerRequest
|
|
want string
|
|
}{
|
|
{
|
|
name: "missing token",
|
|
req: &otopb.RegisterRunnerRequest{
|
|
RunnerId: "runner-123",
|
|
ProtocolVersion: "oto.runner.v1",
|
|
},
|
|
want: "missing enrollment token",
|
|
},
|
|
{
|
|
name: "missing runner id",
|
|
req: &otopb.RegisterRunnerRequest{
|
|
EnrollmentToken: "token-123",
|
|
ProtocolVersion: "oto.runner.v1",
|
|
},
|
|
want: "missing runner id",
|
|
},
|
|
{
|
|
name: "missing protocol",
|
|
req: &otopb.RegisterRunnerRequest{
|
|
EnrollmentToken: "token-123",
|
|
RunnerId: "runner-123",
|
|
},
|
|
want: "missing protocol version",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
response := New().Register(tt.req)
|
|
if response.GetAccepted() {
|
|
t.Fatal("Register accepted incomplete request")
|
|
}
|
|
if response.GetRejectReason() != tt.want {
|
|
t.Fatalf("reject reason = %q, want %q", response.GetRejectReason(), tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestRegistryHeartbeatTransitionsStatus(t *testing.T) {
|
|
var mockTime time.Time
|
|
registry := NewWithClock(func() time.Time { return mockTime })
|
|
|
|
// 1. Heartbeat on unknown runner should fail
|
|
hbRes := registry.Heartbeat("runner-123", otopb.HeartbeatStatus_HEARTBEAT_STATUS_HEALTHY)
|
|
if hbRes.GetSuccess() {
|
|
t.Fatal("expected heartbeat on unknown runner to fail")
|
|
}
|
|
|
|
// 2. Register runner
|
|
mockTime = time.Date(2026, 6, 5, 12, 0, 0, 0, time.UTC)
|
|
registry.Register(&otopb.RegisterRunnerRequest{
|
|
EnrollmentToken: "token-123",
|
|
RunnerId: "runner-123",
|
|
ProtocolVersion: "oto.runner.v1",
|
|
Capability: &otopb.RunnerCapability{
|
|
Name: "oto-runner",
|
|
Version: "1.0.0",
|
|
},
|
|
})
|
|
|
|
record, _ := registry.Snapshot("runner-123")
|
|
if record.Status != StatusAccepted {
|
|
t.Fatalf("status = %q, want %q", record.Status, StatusAccepted)
|
|
}
|
|
|
|
// 3. First heartbeat transitions to Online
|
|
mockTime = mockTime.Add(5 * time.Second)
|
|
hbRes = registry.Heartbeat("runner-123", otopb.HeartbeatStatus_HEARTBEAT_STATUS_HEALTHY)
|
|
if !hbRes.GetSuccess() {
|
|
t.Fatalf("heartbeat failed: %s", hbRes.GetErrorMessage())
|
|
}
|
|
|
|
record, _ = registry.Snapshot("runner-123")
|
|
if record.Status != StatusOnline {
|
|
t.Fatalf("status = %q, want %q", record.Status, StatusOnline)
|
|
}
|
|
if !record.FirstHeartbeatAt.Equal(mockTime) {
|
|
t.Fatalf("FirstHeartbeatAt = %v, want %v", record.FirstHeartbeatAt, mockTime)
|
|
}
|
|
if !record.LastHeartbeatAt.Equal(mockTime) {
|
|
t.Fatalf("LastHeartbeatAt = %v, want %v", record.LastHeartbeatAt, mockTime)
|
|
}
|
|
|
|
// 4. Subsequent heartbeat updates LastHeartbeatAt but not FirstHeartbeatAt
|
|
mockTime = mockTime.Add(10 * time.Second)
|
|
hbRes = registry.Heartbeat("runner-123", otopb.HeartbeatStatus_HEARTBEAT_STATUS_HEALTHY)
|
|
if !hbRes.GetSuccess() {
|
|
t.Fatalf("heartbeat failed: %s", hbRes.GetErrorMessage())
|
|
}
|
|
|
|
record, _ = registry.Snapshot("runner-123")
|
|
if record.Status != StatusOnline {
|
|
t.Fatalf("status = %q, want %q", record.Status, StatusOnline)
|
|
}
|
|
if !record.FirstHeartbeatAt.Equal(mockTime.Add(-10 * time.Second)) {
|
|
t.Fatalf("FirstHeartbeatAt should remain unchanged")
|
|
}
|
|
if !record.LastHeartbeatAt.Equal(mockTime) {
|
|
t.Fatalf("LastHeartbeatAt = %v, want %v", record.LastHeartbeatAt, mockTime)
|
|
}
|
|
}
|
|
|
|
func TestRegistryDisconnect(t *testing.T) {
|
|
registry := New()
|
|
registry.Register(&otopb.RegisterRunnerRequest{
|
|
EnrollmentToken: "token-123",
|
|
RunnerId: "runner-123",
|
|
ProtocolVersion: "oto.runner.v1",
|
|
Capability: &otopb.RunnerCapability{
|
|
Name: "oto-runner",
|
|
Version: "1.0.0",
|
|
},
|
|
})
|
|
|
|
if ok := registry.Disconnect("unknown-runner"); ok {
|
|
t.Fatal("expected disconnect on unknown runner to fail")
|
|
}
|
|
|
|
if ok := registry.Disconnect("runner-123"); !ok {
|
|
t.Fatal("expected disconnect to succeed")
|
|
}
|
|
|
|
record, _ := registry.Snapshot("runner-123")
|
|
if record.Status != StatusDisconnected {
|
|
t.Fatalf("status = %q, want %q", record.Status, StatusDisconnected)
|
|
}
|
|
}
|
|
|
|
func TestRegistryUpdatesAliasAndPreservesItOnReregister(t *testing.T) {
|
|
registry := New()
|
|
registry.Register(&otopb.RegisterRunnerRequest{
|
|
EnrollmentToken: "token-123",
|
|
RunnerId: "runner-123",
|
|
Alias: "device-default",
|
|
ProtocolVersion: "oto.runner.v1",
|
|
Capability: &otopb.RunnerCapability{
|
|
Name: "oto-runner",
|
|
Version: "1.0.0",
|
|
},
|
|
})
|
|
|
|
record, ok := registry.UpdateAlias("runner-123", "friendly-name")
|
|
if !ok {
|
|
t.Fatal("expected alias update to succeed")
|
|
}
|
|
if record.Alias != "friendly-name" {
|
|
t.Fatalf("alias = %q, want friendly-name", record.Alias)
|
|
}
|
|
|
|
registry.Register(&otopb.RegisterRunnerRequest{
|
|
EnrollmentToken: "token-123",
|
|
RunnerId: "runner-123",
|
|
Alias: "device-default",
|
|
ProtocolVersion: "oto.runner.v1",
|
|
Capability: &otopb.RunnerCapability{
|
|
Name: "oto-runner",
|
|
Version: "1.0.0",
|
|
},
|
|
})
|
|
|
|
record, ok = registry.Snapshot("runner-123")
|
|
if !ok {
|
|
t.Fatal("expected runner to exist")
|
|
}
|
|
if record.Alias != "friendly-name" {
|
|
t.Fatalf("alias after reregister = %q, want friendly-name", record.Alias)
|
|
}
|
|
|
|
if _, ok := registry.UpdateAlias("unknown-runner", "name"); ok {
|
|
t.Fatal("expected alias update on unknown runner to fail")
|
|
}
|
|
}
|
|
|
|
func TestRegistryCheckTimeouts(t *testing.T) {
|
|
var mockTime time.Time
|
|
registry := NewWithClock(func() time.Time { return mockTime })
|
|
|
|
// Register 2 runners:
|
|
// runner-accepted: remains in accepted state (no heartbeat)
|
|
// runner-online: transitions to online state
|
|
mockTime = time.Date(2026, 6, 5, 12, 0, 0, 0, time.UTC)
|
|
registry.Register(&otopb.RegisterRunnerRequest{
|
|
EnrollmentToken: "token-123",
|
|
RunnerId: "runner-accepted",
|
|
ProtocolVersion: "oto.runner.v1",
|
|
Capability: &otopb.RunnerCapability{
|
|
Name: "oto-runner",
|
|
Version: "1.0.0",
|
|
},
|
|
})
|
|
registry.Register(&otopb.RegisterRunnerRequest{
|
|
EnrollmentToken: "token-123",
|
|
RunnerId: "runner-online",
|
|
ProtocolVersion: "oto.runner.v1",
|
|
Capability: &otopb.RunnerCapability{
|
|
Name: "oto-runner",
|
|
Version: "1.0.0",
|
|
},
|
|
})
|
|
|
|
// transition runner-online to online
|
|
mockTime = mockTime.Add(1 * time.Second)
|
|
registry.Heartbeat("runner-online", otopb.HeartbeatStatus_HEARTBEAT_STATUS_HEALTHY)
|
|
|
|
// verify initial states
|
|
recAccepted, _ := registry.Snapshot("runner-accepted")
|
|
recOnline, _ := registry.Snapshot("runner-online")
|
|
if recAccepted.Status != StatusAccepted || recOnline.Status != StatusOnline {
|
|
t.Fatal("initial status setup mismatch")
|
|
}
|
|
|
|
// Let 30 seconds pass.
|
|
// CheckTimeouts with 10s threshold. Both should timeout.
|
|
mockTime = mockTime.Add(30 * time.Second)
|
|
registry.CheckTimeouts(10 * time.Second)
|
|
|
|
recAccepted, _ = registry.Snapshot("runner-accepted")
|
|
recOnline, _ = registry.Snapshot("runner-online")
|
|
|
|
if recAccepted.Status != StatusHeartbeatTimeout {
|
|
t.Fatalf("runner-accepted status = %q, want %q", recAccepted.Status, StatusHeartbeatTimeout)
|
|
}
|
|
if recOnline.Status != StatusHeartbeatTimeout {
|
|
t.Fatalf("runner-online status = %q, want %q", recOnline.Status, StatusHeartbeatTimeout)
|
|
}
|
|
if recAccepted.FailureReason != "heartbeat timeout before first heartbeat" {
|
|
t.Fatalf("runner-accepted failure reason: %q", recAccepted.FailureReason)
|
|
}
|
|
if recOnline.FailureReason != "heartbeat timeout" {
|
|
t.Fatalf("runner-online failure reason: %q", recOnline.FailureReason)
|
|
}
|
|
}
|
|
|
|
func TestRegistryHeartbeatTerminalStateDefense(t *testing.T) {
|
|
registry := New()
|
|
registry.Register(&otopb.RegisterRunnerRequest{
|
|
EnrollmentToken: "token-123",
|
|
RunnerId: "runner-123",
|
|
ProtocolVersion: "oto.runner.v1",
|
|
Capability: &otopb.RunnerCapability{
|
|
Name: "oto-runner",
|
|
Version: "1.0.0",
|
|
},
|
|
})
|
|
|
|
// Move to Online
|
|
registry.Heartbeat("runner-123", otopb.HeartbeatStatus_HEARTBEAT_STATUS_HEALTHY)
|
|
|
|
// Disconnect
|
|
registry.Disconnect("runner-123")
|
|
rec, _ := registry.Snapshot("runner-123")
|
|
if rec.Status != StatusDisconnected {
|
|
t.Fatalf("expected disconnected, got %s", rec.Status)
|
|
}
|
|
|
|
// Heartbeat should fail and keep state as disconnected
|
|
hbRes := registry.Heartbeat("runner-123", otopb.HeartbeatStatus_HEARTBEAT_STATUS_HEALTHY)
|
|
if hbRes.GetSuccess() {
|
|
t.Fatal("expected heartbeat to be rejected on disconnected runner")
|
|
}
|
|
rec, _ = registry.Snapshot("runner-123")
|
|
if rec.Status != StatusDisconnected {
|
|
t.Fatalf("expected status to remain disconnected, got %s", rec.Status)
|
|
}
|
|
|
|
// Re-register to make it accepted again
|
|
registry.Register(&otopb.RegisterRunnerRequest{
|
|
EnrollmentToken: "token-123",
|
|
RunnerId: "runner-123",
|
|
ProtocolVersion: "oto.runner.v1",
|
|
Capability: &otopb.RunnerCapability{
|
|
Name: "oto-runner",
|
|
Version: "1.0.0",
|
|
},
|
|
})
|
|
rec, _ = registry.Snapshot("runner-123")
|
|
if rec.Status != StatusAccepted {
|
|
t.Fatalf("expected status to be accepted, got %s", rec.Status)
|
|
}
|
|
|
|
// Revive to online via heartbeat
|
|
registry.Heartbeat("runner-123", otopb.HeartbeatStatus_HEARTBEAT_STATUS_HEALTHY)
|
|
|
|
// Force timeout
|
|
var mockTime time.Time
|
|
registry = NewWithClock(func() time.Time { return mockTime })
|
|
mockTime = time.Date(2026, 6, 5, 12, 0, 0, 0, time.UTC)
|
|
registry.Register(&otopb.RegisterRunnerRequest{
|
|
EnrollmentToken: "token-123",
|
|
RunnerId: "runner-123",
|
|
ProtocolVersion: "oto.runner.v1",
|
|
Capability: &otopb.RunnerCapability{
|
|
Name: "oto-runner",
|
|
Version: "1.0.0",
|
|
},
|
|
})
|
|
mockTime = mockTime.Add(1 * time.Second)
|
|
registry.Heartbeat("runner-123", otopb.HeartbeatStatus_HEARTBEAT_STATUS_HEALTHY)
|
|
|
|
// pass timeout threshold
|
|
mockTime = mockTime.Add(30 * time.Second)
|
|
registry.CheckTimeouts(10 * time.Second)
|
|
rec, _ = registry.Snapshot("runner-123")
|
|
if rec.Status != StatusHeartbeatTimeout {
|
|
t.Fatalf("expected status heartbeat_timeout, got %s", rec.Status)
|
|
}
|
|
|
|
// Heartbeat should fail and keep state as heartbeat_timeout
|
|
hbRes = registry.Heartbeat("runner-123", otopb.HeartbeatStatus_HEARTBEAT_STATUS_HEALTHY)
|
|
if hbRes.GetSuccess() {
|
|
t.Fatal("expected heartbeat to be rejected on timed out runner")
|
|
}
|
|
rec, _ = registry.Snapshot("runner-123")
|
|
if rec.Status != StatusHeartbeatTimeout {
|
|
t.Fatalf("expected status to remain heartbeat_timeout, got %s", rec.Status)
|
|
}
|
|
}
|
|
|
|
func TestRegistryRejectsIncompatibleRunnerContract(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
req *otopb.RegisterRunnerRequest
|
|
wantMsg string
|
|
wantCode string
|
|
}{
|
|
{
|
|
name: "unsupported protocol version",
|
|
req: &otopb.RegisterRunnerRequest{
|
|
EnrollmentToken: "token-123",
|
|
RunnerId: "runner-123",
|
|
ProtocolVersion: "oto.runner.v2",
|
|
Capability: &otopb.RunnerCapability{
|
|
Name: "oto-runner",
|
|
Version: "1.0.0",
|
|
},
|
|
},
|
|
wantMsg: `unsupported protocol version: "oto.runner.v2"`,
|
|
wantCode: "incompatible_runner",
|
|
},
|
|
{
|
|
name: "missing capability",
|
|
req: &otopb.RegisterRunnerRequest{
|
|
EnrollmentToken: "token-123",
|
|
RunnerId: "runner-123",
|
|
ProtocolVersion: "oto.runner.v1",
|
|
},
|
|
wantMsg: "missing runner capability",
|
|
wantCode: "incompatible_runner",
|
|
},
|
|
{
|
|
name: "unsupported capability name",
|
|
req: &otopb.RegisterRunnerRequest{
|
|
EnrollmentToken: "token-123",
|
|
RunnerId: "runner-123",
|
|
ProtocolVersion: "oto.runner.v1",
|
|
Capability: &otopb.RunnerCapability{
|
|
Name: "other-runner",
|
|
Version: "1.0.0",
|
|
},
|
|
},
|
|
wantMsg: `unsupported capability name: "other-runner"`,
|
|
wantCode: "incompatible_runner",
|
|
},
|
|
{
|
|
name: "missing capability version",
|
|
req: &otopb.RegisterRunnerRequest{
|
|
EnrollmentToken: "token-123",
|
|
RunnerId: "runner-123",
|
|
ProtocolVersion: "oto.runner.v1",
|
|
Capability: &otopb.RunnerCapability{
|
|
Name: "oto-runner",
|
|
Version: "",
|
|
},
|
|
},
|
|
wantMsg: "missing capability version",
|
|
wantCode: "incompatible_runner",
|
|
},
|
|
{
|
|
name: "incompatible capability major version",
|
|
req: &otopb.RegisterRunnerRequest{
|
|
EnrollmentToken: "token-123",
|
|
RunnerId: "runner-123",
|
|
ProtocolVersion: "oto.runner.v1",
|
|
Capability: &otopb.RunnerCapability{
|
|
Name: "oto-runner",
|
|
Version: "2.0.0",
|
|
},
|
|
},
|
|
wantMsg: `incompatible capability major version: "2.0.0", want major version 1`,
|
|
wantCode: "incompatible_runner",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
response := New().Register(tt.req)
|
|
if response.GetAccepted() {
|
|
t.Fatal("Register accepted incompatible contract")
|
|
}
|
|
if response.GetRejectReason() != tt.wantMsg {
|
|
t.Fatalf("reject reason = %q, want %q", response.GetRejectReason(), tt.wantMsg)
|
|
}
|
|
if response.GetError() == nil {
|
|
t.Fatal("missing structured error")
|
|
}
|
|
if response.GetError().GetCode() != tt.wantCode {
|
|
t.Fatalf("error code = %q, want %q", response.GetError().GetCode(), tt.wantCode)
|
|
}
|
|
if response.GetError().GetMessage() != tt.wantMsg {
|
|
t.Fatalf("error message = %q, want %q", response.GetError().GetMessage(), tt.wantMsg)
|
|
}
|
|
})
|
|
}
|
|
}
|