105 lines
3.7 KiB
Go
105 lines
3.7 KiB
Go
package jobs
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestWorkerUsesUniqueLeaseOwnerForEachTick(t *testing.T) {
|
|
store := &leaseTestStore{}
|
|
worker := NewWorker(store, nil, nil, nil, nil, WorkerConfig{}, nil)
|
|
tokens := []string{"worker-1:lease-a", "worker-1:lease-b"}
|
|
worker.leaseOwner = func(workerID string) (string, error) {
|
|
token := tokens[0]
|
|
tokens = tokens[1:]
|
|
return token, nil
|
|
}
|
|
|
|
first, err := worker.Tick(context.Background(), "worker-1")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
second, err := worker.Tick(context.Background(), "worker-1")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
if first.WorkerID != "worker-1" || second.WorkerID != "worker-1" {
|
|
t.Fatalf("public worker IDs = %q, %q", first.WorkerID, second.WorkerID)
|
|
}
|
|
if len(store.claimOwners) != 2 || store.claimOwners[0] == store.claimOwners[1] {
|
|
t.Fatalf("claim owners = %#v, want two unique lease tokens", store.claimOwners)
|
|
}
|
|
if store.claimOwners[0] != "worker-1:lease-a" || store.claimOwners[1] != "worker-1:lease-b" {
|
|
t.Fatalf("claim owners = %#v", store.claimOwners)
|
|
}
|
|
}
|
|
|
|
func TestWorkerSettleUsesClaimedJobLeaseToken(t *testing.T) {
|
|
now := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC)
|
|
store := &leaseTestStore{job: Job{ID: "job-1", Status: StatusRunning, LockedBy: "worker-1:lease-current"}}
|
|
worker := NewWorker(store, nil, nil, nil, nil, WorkerConfig{}, func() time.Time { return now })
|
|
|
|
settled, action, err := worker.settle(context.Background(), store.job)
|
|
if err != nil || action != "released" {
|
|
t.Fatalf("settle = (%#v, %q, %v)", settled, action, err)
|
|
}
|
|
if len(store.expectedLeaseOwners) != 1 || store.expectedLeaseOwners[0] != "worker-1:lease-current" {
|
|
t.Fatalf("expected lease owners = %#v", store.expectedLeaseOwners)
|
|
}
|
|
}
|
|
|
|
func TestWorkerStaleLeaseCannotSettleNewLeaseFromSameWorkerID(t *testing.T) {
|
|
now := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC)
|
|
store := &leaseTestStore{job: Job{ID: "job-1", Status: StatusRunning, LockedBy: "worker-1:lease-new"}}
|
|
worker := NewWorker(store, nil, nil, nil, nil, WorkerConfig{}, func() time.Time { return now })
|
|
stale := store.job
|
|
stale.LockedBy = "worker-1:lease-old"
|
|
|
|
_, _, err := worker.settle(context.Background(), stale)
|
|
if !errors.Is(err, ErrTransitionConflict) {
|
|
t.Fatalf("settle stale lease error = %v, want transition conflict", err)
|
|
}
|
|
if store.job.LockedBy != "worker-1:lease-new" {
|
|
t.Fatalf("fresh lease was overwritten: %#v", store.job)
|
|
}
|
|
}
|
|
|
|
type leaseTestStore struct {
|
|
job Job
|
|
claimOwners []string
|
|
expectedLeaseOwners []string
|
|
}
|
|
|
|
func (store *leaseTestStore) ListJobs(context.Context, ListFilter) ([]Job, error) { return nil, nil }
|
|
func (store *leaseTestStore) FindJob(context.Context, Scope, string) (Job, bool, error) {
|
|
return Job{}, false, nil
|
|
}
|
|
func (store *leaseTestStore) FindIdempotentJob(context.Context, Scope, string) (Job, bool, error) {
|
|
return Job{}, false, nil
|
|
}
|
|
func (store *leaseTestStore) CreateJob(context.Context, Job) (Job, error) { return Job{}, nil }
|
|
func (store *leaseTestStore) DeleteJob(context.Context, string) error { return nil }
|
|
func (store *leaseTestStore) ClaimJobs(_ context.Context, owner string, _ int, _ int) ([]Job, error) {
|
|
store.claimOwners = append(store.claimOwners, owner)
|
|
return nil, nil
|
|
}
|
|
func (store *leaseTestStore) UpdateJob(_ context.Context, _ string, patch Patch) (Job, error) {
|
|
if patch.ExpectedLockedBy != nil {
|
|
store.expectedLeaseOwners = append(store.expectedLeaseOwners, *patch.ExpectedLockedBy)
|
|
if store.job.LockedBy != *patch.ExpectedLockedBy {
|
|
return Job{}, ErrTransitionConflict
|
|
}
|
|
}
|
|
if patch.ScheduledAt != nil {
|
|
store.job.ScheduledAt = *patch.ScheduledAt
|
|
}
|
|
if patch.ClearLease {
|
|
store.job.LockedBy = ""
|
|
store.job.LockedAt = nil
|
|
}
|
|
return store.job, nil
|
|
}
|