351 lines
13 KiB
Go
351 lines
13 KiB
Go
package jobs
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestFingerprintMatchesSharedStableJSONVector(t *testing.T) {
|
|
body := map[string]any{
|
|
"settings": map[string]any{"width": float64(1024), "height": float64(768)},
|
|
"prompt": "hello", "idempotencyKey": "body-key", "capability": "image.generate",
|
|
}
|
|
got, err := Fingerprint(body)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
const want = "3ac3128766c0bdac36627b996cb5b3b1ed752794989d559c24f6621988e51a03"
|
|
if got != want {
|
|
t.Fatalf("Fingerprint = %q, want %q", got, want)
|
|
}
|
|
}
|
|
|
|
func TestServiceReplaysMatchingIdempotencyAndRejectsDrift(t *testing.T) {
|
|
clock := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC)
|
|
store := newMemoryJobStore()
|
|
service := NewService(store, func() time.Time { return clock })
|
|
command := fixtureCreateCommand("hello")
|
|
|
|
created, reused, err := service.Create(context.Background(), command)
|
|
if err != nil || reused {
|
|
t.Fatalf("first Create = (%#v,%v,%v)", created, reused, err)
|
|
}
|
|
replayed, reused, err := service.Create(context.Background(), command)
|
|
if err != nil || !reused || replayed.ID != created.ID {
|
|
t.Fatalf("replay Create = (%#v,%v,%v)", replayed, reused, err)
|
|
}
|
|
_, _, err = service.Create(context.Background(), fixtureCreateCommand("hello!"))
|
|
var conflict *Error
|
|
if !errors.As(err, &conflict) || conflict.Kind != ErrorConflict || conflict.Status != 409 {
|
|
t.Fatalf("drift error = %#v, want conflict/409", err)
|
|
}
|
|
}
|
|
|
|
func TestServiceRecoversUniqueInsertRaceAsIdempotentReplay(t *testing.T) {
|
|
store := newMemoryJobStore()
|
|
command := fixtureCreateCommand("hello")
|
|
fingerprint, _ := Fingerprint(command.IdempotencyBody)
|
|
existing := command.Job
|
|
existing.IdempotencyFingerprint = fingerprint
|
|
store.race = &existing
|
|
service := NewService(store, func() time.Time { return time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC) })
|
|
|
|
job, reused, err := service.Create(context.Background(), command)
|
|
if err != nil || !reused || job.ID != existing.ID {
|
|
t.Fatalf("Create race = (%#v,%v,%v)", job, reused, err)
|
|
}
|
|
}
|
|
|
|
func TestCancelIsOwnerScopedIdempotentAndClearsLease(t *testing.T) {
|
|
store := newMemoryJobStore()
|
|
locked := time.Date(2026, 8, 13, 7, 59, 0, 0, time.UTC)
|
|
store.jobs["job-1"] = Job{ID: "job-1", OwnerID: "api:client", ExternalClientID: "client", Status: StatusRunning, LockedAt: &locked, LockedBy: "worker"}
|
|
service := NewService(store, func() time.Time { return time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC) })
|
|
refunds := &refundStub{}
|
|
|
|
cancelled, err := service.Cancel(context.Background(), Scope{OwnerID: "api:client", ExternalClientID: "client"}, "job-1", refunds)
|
|
if err != nil || cancelled.Status != StatusCancelled || cancelled.LockedAt != nil || cancelled.LockedBy != "" || refunds.calls != 1 {
|
|
t.Fatalf("Cancel = %#v, %v; refunds=%d", cancelled, err, refunds.calls)
|
|
}
|
|
again, err := service.Cancel(context.Background(), Scope{OwnerID: "api:client", ExternalClientID: "client"}, "job-1", refunds)
|
|
if err != nil || again.Status != StatusCancelled || refunds.calls != 1 {
|
|
t.Fatalf("second Cancel = %#v, %v; refunds=%d", again, err, refunds.calls)
|
|
}
|
|
_, err = service.Cancel(context.Background(), Scope{OwnerID: "api:other", ExternalClientID: "other"}, "job-1", refunds)
|
|
var notFound *Error
|
|
if !errors.As(err, ¬Found) || notFound.Status != 404 {
|
|
t.Fatalf("cross-owner error = %#v", err)
|
|
}
|
|
}
|
|
|
|
func TestWorkerSchedulesRetryReleasesRunningAndSettlesTerminal(t *testing.T) {
|
|
now := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC)
|
|
tests := []struct {
|
|
name string
|
|
job Job
|
|
want string
|
|
status Status
|
|
attempts int
|
|
scheduled time.Time
|
|
}{
|
|
{name: "retryable failure", job: Job{ID: "retry", Status: StatusFailed, Error: &JobError{Message: "temporary", Retryable: true}, MaxAttempts: 3}, want: "retry_scheduled", status: StatusQueued, attempts: 1, scheduled: now.Add(10 * time.Second)},
|
|
{name: "running release", job: Job{ID: "running", Status: StatusRunning}, want: "released", status: StatusRunning, scheduled: now.Add(5 * time.Second)},
|
|
{name: "terminal failure", job: Job{ID: "terminal", Status: StatusFailed, Error: &JobError{Message: "fatal"}, MaxAttempts: 3}, want: "processed", status: StatusFailed, attempts: 1},
|
|
}
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
store := newMemoryJobStore()
|
|
test.job.OwnerID = "owner"
|
|
test.job.Capability = "image.generate"
|
|
test.job.UpdatedAt = now
|
|
store.claimed = []Job{test.job}
|
|
store.jobs[test.job.ID] = test.job
|
|
processor := &processorStub{job: test.job}
|
|
refunds := &refundStub{}
|
|
worker := NewWorker(store, processor, refunds, nil, nil, WorkerConfig{}, func() time.Time { return now })
|
|
result, err := worker.Tick(context.Background(), "worker-1")
|
|
if err != nil || len(result.Jobs) != 1 || result.Jobs[0].Action != test.want {
|
|
t.Fatalf("Tick = %#v, %v", result, err)
|
|
}
|
|
stored := store.jobs[test.job.ID]
|
|
if stored.Status != test.status || stored.Attempts != test.attempts {
|
|
t.Fatalf("stored = %#v, want status=%s attempts=%d", stored, test.status, test.attempts)
|
|
}
|
|
if !test.scheduled.IsZero() && !stored.ScheduledAt.Equal(test.scheduled) {
|
|
t.Fatalf("scheduled = %s, want %s", stored.ScheduledAt, test.scheduled)
|
|
}
|
|
if test.want == "retry_scheduled" && stored.Error != nil {
|
|
t.Fatalf("retry must clear the transient error, got %#v", stored.Error)
|
|
}
|
|
if test.want == "processed" && refunds.calls != 1 {
|
|
t.Fatalf("refund calls = %d, want 1", refunds.calls)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestWorkerRecoversTerminalProviderResultWithoutResubmittingAndFinalizesAfterEffects(t *testing.T) {
|
|
now := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC)
|
|
locked := now.Add(-time.Minute)
|
|
job := Job{ID: "terminal", OwnerID: "owner", Capability: "image.generate", Provider: "fixture", Status: StatusSucceeded, LockedAt: &locked, LockedBy: "worker-1"}
|
|
store := newMemoryJobStore()
|
|
store.claimed = []Job{job}
|
|
store.jobs[job.ID] = job
|
|
processor := &processorStub{job: job}
|
|
usage := &usageStub{}
|
|
webhooks := &webhookStub{result: WebhookResult{Attempts: 1, LastStatus: map[string]any{"ok": true}}}
|
|
worker := NewWorker(store, processor, nil, usage, webhooks, WorkerConfig{}, func() time.Time { return now })
|
|
|
|
result, err := worker.Tick(context.Background(), "worker-1")
|
|
if err != nil || len(result.Jobs) != 1 || result.Jobs[0].Action != "processed" {
|
|
t.Fatalf("Tick = %#v, %v", result, err)
|
|
}
|
|
stored := store.jobs[job.ID]
|
|
if processor.calls != 1 || usage.calls != 1 || webhooks.calls != 1 || stored.FinalizedAt == nil || stored.LockedBy != "" {
|
|
t.Fatalf("processor=%d usage=%d webhook=%d stored=%#v", processor.calls, usage.calls, webhooks.calls, stored)
|
|
}
|
|
}
|
|
|
|
func TestWorkerLostLeaseCannotOverwriteCancellation(t *testing.T) {
|
|
now := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC)
|
|
locked := now.Add(-time.Minute)
|
|
claimed := Job{ID: "race", OwnerID: "owner", Capability: "image.generate", Provider: "fixture", Status: StatusRunning, LockedAt: &locked, LockedBy: "worker-1"}
|
|
store := newMemoryJobStore()
|
|
store.claimed = []Job{claimed}
|
|
cancelled := claimed
|
|
cancelled.Status, cancelled.LockedAt, cancelled.LockedBy = StatusCancelled, nil, ""
|
|
store.jobs[claimed.ID] = cancelled
|
|
worker := NewWorker(store, &processorStub{job: Job{ID: claimed.ID, Status: StatusSucceeded}}, nil, nil, nil, WorkerConfig{}, func() time.Time { return now })
|
|
|
|
_, err := worker.Tick(context.Background(), "worker-1")
|
|
if !errors.Is(err, ErrTransitionConflict) {
|
|
t.Fatalf("Tick error = %v, want transition conflict", err)
|
|
}
|
|
if got := store.jobs[claimed.ID]; got.Status != StatusCancelled {
|
|
t.Fatalf("stored = %#v", got)
|
|
}
|
|
}
|
|
|
|
func TestWorkerCountsEachFailedAdvanceExactlyOnce(t *testing.T) {
|
|
now := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC)
|
|
store := newMemoryJobStore()
|
|
job := Job{ID: "retry", OwnerID: "owner", Capability: "image.generate", Provider: "fixture", Status: StatusRunning, LockedBy: "worker-1", MaxAttempts: 3}
|
|
store.claimed, store.jobs[job.ID] = []Job{job}, job
|
|
worker := NewWorker(store, &processorStub{err: errors.New("temporary")}, nil, nil, nil, WorkerConfig{RetryBase: time.Second}, func() time.Time { return now })
|
|
|
|
if _, err := worker.Tick(context.Background(), "worker-1"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
got := store.jobs[job.ID]
|
|
if got.Attempts != 1 || got.Status != StatusQueued {
|
|
t.Fatalf("after one failed advance = %#v", got)
|
|
}
|
|
}
|
|
|
|
func TestWorkerRetryPreservesKnownProviderTaskToAvoidDuplicateSubmission(t *testing.T) {
|
|
now := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC)
|
|
store := newMemoryJobStore()
|
|
job := Job{ID: "poll", OwnerID: "owner", Capability: "image.generate", Provider: "fixture", ProviderTaskID: "provider-task", Status: StatusRunning, LockedBy: "worker-1", MaxAttempts: 3}
|
|
store.claimed, store.jobs[job.ID] = []Job{job}, job
|
|
worker := NewWorker(store, &processorStub{err: errors.New("query timeout")}, nil, nil, nil, WorkerConfig{}, func() time.Time { return now })
|
|
if _, err := worker.Tick(context.Background(), "worker-1"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if got := store.jobs[job.ID]; got.ProviderTaskID != "provider-task" || got.Status != StatusQueued {
|
|
t.Fatalf("retry lost provider recovery handle: %#v", got)
|
|
}
|
|
}
|
|
|
|
func fixtureCreateCommand(prompt string) CreateCommand {
|
|
body := map[string]any{
|
|
"capability": "image.generate", "prompt": prompt,
|
|
"settings": map[string]any{"height": float64(768), "width": float64(1024)},
|
|
"idempotencyKey": "body-key",
|
|
}
|
|
return CreateCommand{Job: Job{
|
|
ID: "job-1", OwnerID: "api:client", ExternalClientID: "client", Capability: "image.generate",
|
|
Provider: "mock", ReqKey: "fixture", Status: StatusQueued, IdempotencyKey: "idem-1",
|
|
RequestPayload: json.RawMessage(`{"prompt":"` + prompt + `"}`),
|
|
}, IdempotencyBody: body}
|
|
}
|
|
|
|
type memoryJobStore struct {
|
|
jobs map[string]Job
|
|
claimed []Job
|
|
race *Job
|
|
}
|
|
|
|
func (store *memoryJobStore) DeleteJob(_ context.Context, id string) error {
|
|
delete(store.jobs, id)
|
|
return nil
|
|
}
|
|
|
|
func newMemoryJobStore() *memoryJobStore { return &memoryJobStore{jobs: map[string]Job{}} }
|
|
func (store *memoryJobStore) ListJobs(context.Context, ListFilter) ([]Job, error) { return nil, nil }
|
|
func (store *memoryJobStore) FindJob(_ context.Context, scope Scope, id string) (Job, bool, error) {
|
|
job, ok := store.jobs[id]
|
|
return job, ok && scope.Owns(job), nil
|
|
}
|
|
func (store *memoryJobStore) FindIdempotentJob(_ context.Context, scope Scope, key string) (Job, bool, error) {
|
|
for _, job := range store.jobs {
|
|
if scope.Owns(job) && job.IdempotencyKey == key {
|
|
return job, true, nil
|
|
}
|
|
}
|
|
return Job{}, false, nil
|
|
}
|
|
func (store *memoryJobStore) CreateJob(_ context.Context, job Job) (Job, error) {
|
|
if store.race != nil {
|
|
store.jobs[store.race.ID] = *store.race
|
|
store.race = nil
|
|
return Job{}, ErrUniqueIdempotency
|
|
}
|
|
store.jobs[job.ID] = job
|
|
return job, nil
|
|
}
|
|
func (store *memoryJobStore) UpdateJob(_ context.Context, id string, patch Patch) (Job, error) {
|
|
job := store.jobs[id]
|
|
if len(patch.ExpectedStatuses) != 0 {
|
|
matched := false
|
|
for _, status := range patch.ExpectedStatuses {
|
|
matched = matched || job.Status == status
|
|
}
|
|
if !matched {
|
|
return Job{}, ErrTransitionConflict
|
|
}
|
|
}
|
|
if patch.ExpectedLockedBy != nil && job.LockedBy != *patch.ExpectedLockedBy {
|
|
return Job{}, ErrTransitionConflict
|
|
}
|
|
if patch.Status != nil {
|
|
job.Status = *patch.Status
|
|
}
|
|
if patch.Error != nil {
|
|
job.Error = patch.Error
|
|
}
|
|
if patch.ClearError {
|
|
job.Error = nil
|
|
}
|
|
if patch.Attempts != nil {
|
|
job.Attempts = *patch.Attempts
|
|
}
|
|
if patch.ScheduledAt != nil {
|
|
job.ScheduledAt = *patch.ScheduledAt
|
|
}
|
|
if patch.CompletedAt != nil {
|
|
job.CompletedAt = patch.CompletedAt
|
|
}
|
|
if patch.ClearLease {
|
|
job.LockedAt = nil
|
|
job.LockedBy = ""
|
|
}
|
|
if patch.ClearProviderTaskID {
|
|
job.ProviderTaskID = ""
|
|
}
|
|
if patch.ProviderTaskID != nil {
|
|
job.ProviderTaskID = *patch.ProviderTaskID
|
|
}
|
|
if patch.ProviderDispatchStartedAt != nil {
|
|
value := *patch.ProviderDispatchStartedAt
|
|
job.ProviderDispatchStartedAt = &value
|
|
}
|
|
if patch.ClearProviderDispatch {
|
|
job.ProviderDispatchStartedAt = nil
|
|
}
|
|
if patch.SetResponsePayload {
|
|
job.ResponsePayload = append([]byte(nil), patch.ResponsePayload...)
|
|
}
|
|
if patch.WebhookAttempts != nil {
|
|
job.WebhookAttempts = *patch.WebhookAttempts
|
|
}
|
|
if patch.SetWebhookStatus {
|
|
job.WebhookLastStatus = patch.WebhookLastStatus
|
|
}
|
|
if patch.FinalizedAt != nil {
|
|
job.FinalizedAt = patch.FinalizedAt
|
|
}
|
|
store.jobs[id] = job
|
|
return job, nil
|
|
}
|
|
func (store *memoryJobStore) ClaimJobs(context.Context, string, int, int) ([]Job, error) {
|
|
return store.claimed, nil
|
|
}
|
|
|
|
type refundStub struct{ calls int }
|
|
|
|
func (stub *refundStub) Refund(_ context.Context, job Job, _ string) (Job, error) {
|
|
stub.calls++
|
|
return job, nil
|
|
}
|
|
|
|
type processorStub struct {
|
|
job Job
|
|
err error
|
|
calls int
|
|
}
|
|
|
|
func (stub *processorStub) Advance(_ context.Context, input Job) (Job, error) {
|
|
stub.calls++
|
|
if stub.job.ID == "" {
|
|
stub.job = input
|
|
}
|
|
return stub.job, stub.err
|
|
}
|
|
|
|
type usageStub struct{ calls int }
|
|
|
|
func (stub *usageStub) Record(context.Context, Job) error { stub.calls++; return nil }
|
|
|
|
type webhookStub struct {
|
|
calls int
|
|
result WebhookResult
|
|
}
|
|
|
|
func (stub *webhookStub) Deliver(context.Context, Job) (WebhookResult, error) {
|
|
stub.calls++
|
|
return stub.result, nil
|
|
}
|