package postgres import ( "context" "database/sql" "encoding/json" "errors" "reflect" "strings" "testing" "time" "git.nianxx.cn/wangxuming/NianAIGC/backend/internal/jobs" ) func TestJobsAdapterUsesFullClaimFunctionAndBoundsArguments(t *testing.T) { now := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC) querier := &jobQuerier{rows: &jobRows{rows: [][]any{jobRow(now)}}} database := NewDatabase(Config{Backend: BackendPostgres}, querier) claimed, err := database.ClaimJobs(context.Background(), "worker-1", 99, 0) if err != nil { t.Fatal(err) } if querier.query != ClaimJobsSQL || !reflect.DeepEqual(querier.args, []any{"worker-1", 20, 300}) { t.Fatalf("query=%q args=%#v", querier.query, querier.args) } if len(claimed) != 1 || claimed[0].ID != "job-1" || claimed[0].Status != jobs.StatusQueued || claimed[0].RequestPayload == nil { t.Fatalf("claimed=%#v", claimed) } if strings.Contains(ClaimJobsSQL, "SELECT id FROM") || !strings.Contains(ClaimJobsSQL, "claim_generation_jobs") { t.Fatalf("ClaimJobsSQL does not select the complete contract: %s", ClaimJobsSQL) } } func TestJobsAdapterScopesPublicFindAndListByOwnerAndClient(t *testing.T) { now := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC) querier := &jobQuerier{rows: &jobRows{rows: [][]any{jobRow(now)}}} database := NewDatabase(Config{Backend: BackendPostgres}, querier) scope := jobs.Scope{OwnerID: "api:client", ExternalClientID: "client"} _, found, err := database.FindJob(context.Background(), scope, "job-1") if err != nil || !found { t.Fatalf("FindJob found=%v err=%v", found, err) } if !strings.Contains(querier.query, "owner_id = $2::text") || !strings.Contains(querier.query, "external_client_id = $3::text") || !reflect.DeepEqual(querier.args, []any{"job-1", "api:client", "client"}) { t.Fatalf("query=%q args=%#v", querier.query, querier.args) } querier.rows = &jobRows{} _, err = database.ListJobs(context.Background(), jobs.ListFilter{Scope: scope, Status: jobs.StatusRunning, Capability: "video.generate", Limit: 50}) if err != nil { t.Fatal(err) } if !strings.Contains(querier.query, "owner_id = $1::text") || !strings.Contains(querier.query, "external_client_id = $2::text") || !strings.Contains(querier.query, "status = $3::text") || !strings.Contains(querier.query, "capability = $4::text") { t.Fatalf("list query=%q", querier.query) } } func TestJobsAdapterUpdateClearsLeaseAndProviderTask(t *testing.T) { now := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC) row := jobRow(now) row[11] = nil row[22], row[23] = nil, nil querier := &jobQuerier{rows: &jobRows{rows: [][]any{row}}} database := NewDatabase(Config{Backend: BackendPostgres}, querier) status := jobs.StatusQueued attempts := 2 _, err := database.UpdateJob(context.Background(), "job-1", jobs.Patch{Status: &status, Attempts: &attempts, ClearLease: true, ClearProviderTaskID: true}) if err != nil { t.Fatal(err) } for _, fragment := range []string{"status = $2::text", "attempts = $3::integer", "provider_task_id = NULL", "locked_at = NULL", "locked_by = NULL"} { if !strings.Contains(querier.query, fragment) { t.Fatalf("update query missing %q: %s", fragment, querier.query) } } } func TestJobsAdapterUpdateClearsStaleProviderError(t *testing.T) { now := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC) querier := &jobQuerier{rows: &jobRows{rows: [][]any{jobRow(now)}}} database := NewDatabase(Config{Backend: BackendPostgres}, querier) status := jobs.StatusSucceeded if _, err := database.UpdateJob(context.Background(), "job-1", jobs.Patch{Status: &status, ClearError: true}); err != nil { t.Fatal(err) } if !strings.Contains(querier.query, "error = NULL") { t.Fatalf("update query does not clear stale error: %s", querier.query) } } func TestJobsAdapterFencesWorkerUpdateByExpectedStateAndLeaseOwner(t *testing.T) { now := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC) querier := &jobQuerier{rows: &jobRows{rows: [][]any{jobRow(now)}}} database := NewDatabase(Config{Backend: BackendPostgres}, querier) status := jobs.StatusSucceeded worker := "worker-1" _, err := database.UpdateJob(context.Background(), "job-1", jobs.Patch{Status: &status, ExpectedStatuses: []jobs.Status{jobs.StatusRunning}, ExpectedLockedBy: &worker}) if err != nil { t.Fatal(err) } if !strings.Contains(querier.query, "status = ANY(") || !strings.Contains(querier.query, "locked_by =") { t.Fatalf("fenced update query = %s", querier.query) } querier.rows = &jobRows{} _, err = database.UpdateJob(context.Background(), "job-1", jobs.Patch{Status: &status, ExpectedStatuses: []jobs.Status{jobs.StatusRunning}, ExpectedLockedBy: &worker}) if !errors.Is(err, jobs.ErrTransitionConflict) { t.Fatalf("lost lease error = %v", err) } } func TestJobsAdapterActivatesChargedCreationAtomically(t *testing.T) { querier := &jobQuerier{rows: &jobRows{rows: [][]any{{"job-1"}}}} database := NewDatabase(Config{Backend: BackendPostgres}, querier) if err := database.ActivateCreation(context.Background(), "job-1", json.RawMessage(`{"status":"charged"}`)); err != nil { t.Fatal(err) } if querier.query != ActivateJobCreationSQL || !strings.Contains(querier.query, "dispatch_ready_at") { t.Fatalf("activation query = %s", querier.query) } } func TestJobsAdapterActivationReplayIsIdempotentForSameBilling(t *testing.T) { querier := &jobQuerier{rows: &jobRows{rows: [][]any{{"job-1"}}}} database := NewDatabase(Config{Backend: BackendPostgres}, querier) billingState := json.RawMessage(`{"status":"charged","ledgerEntryId":"ledger-1"}`) if err := database.ActivateCreation(context.Background(), "job-1", billingState); err != nil { t.Fatal(err) } querier.rows = &jobRows{rows: [][]any{{"job-1"}}} if err := database.ActivateCreation(context.Background(), "job-1", billingState); err != nil { t.Fatalf("replay: %v", err) } if !strings.Contains(querier.query, "dispatch_ready_at IS NULL OR billing = $2::jsonb") { t.Fatalf("activation is not idempotent for matching billing: %s", querier.query) } } func TestJobsAdapterWritesBillingAndOutputAssetsThroughNarrowStateSeam(t *testing.T) { querier := &jobQuerier{rows: &jobRows{rows: [][]any{{"job-1"}}}} database := NewDatabase(Config{Backend: BackendPostgres}, querier) if err := database.WriteBilling(context.Background(), "job-1", json.RawMessage(`{"status":"refunded"}`)); err != nil { t.Fatal(err) } if querier.query != WriteJobBillingSQL || !reflect.DeepEqual(querier.args, []any{"job-1", json.RawMessage(`{"status":"refunded"}`)}) { t.Fatalf("billing query=%q args=%#v", querier.query, querier.args) } querier.rows = &jobRows{rows: [][]any{{"job-1"}}} if err := database.WriteOutputAssetIDs(context.Background(), "job-1", []string{"asset-1", "asset-2"}); err != nil { t.Fatal(err) } if querier.query != WriteJobOutputAssetIDsSQL || !reflect.DeepEqual(querier.args, []any{"job-1", []string{"asset-1", "asset-2"}}) { t.Fatalf("outputs query=%q args=%#v", querier.query, querier.args) } } func TestJobsAdapterFencesOrchestrationStateWrites(t *testing.T) { querier := &jobQuerier{rows: &jobRows{rows: [][]any{{"job-1"}}}} database := NewDatabase(Config{Backend: BackendPostgres}, querier) if err := database.WriteBillingFenced(context.Background(), "job-1", json.RawMessage(`{"status":"settled"}`), jobs.StatusSucceeded, "worker-1"); err != nil { t.Fatal(err) } if querier.query != WriteJobBillingFencedSQL || !reflect.DeepEqual(querier.args, []any{"job-1", json.RawMessage(`{"status":"settled"}`), "succeeded", "worker-1"}) { t.Fatalf("query=%q args=%#v", querier.query, querier.args) } querier.rows = &jobRows{} if err := database.WriteOutputAssetIDsFenced(context.Background(), "job-1", []string{"asset"}, jobs.StatusSucceeded, "worker-1"); !errors.Is(err, jobs.ErrTransitionConflict) { t.Fatalf("lost lease error=%v", err) } } func TestJobsAdapterPersistsFailedCreationAtomically(t *testing.T) { querier := &jobQuerier{rows: &jobRows{rows: [][]any{{"job-1"}}}} database := NewDatabase(Config{Backend: BackendPostgres}, querier) job := jobs.Job{ ID: "job-1", Status: jobs.StatusFailed, Error: &jobs.JobError{Message: "generation charge failed"}, Billing: json.RawMessage(`{"status":"not_charged"}`), } if err := database.FailCreation(context.Background(), job); err != nil { t.Fatal(err) } if querier.query != FailJobCreationSQL { t.Fatalf("query=%q", querier.query) } if len(querier.args) != 4 || querier.args[0] != "job-1" || querier.args[1] != string(jobs.StatusFailed) || !reflect.DeepEqual(querier.args[2], json.RawMessage(`{"message":"generation charge failed"}`)) || !reflect.DeepEqual(querier.args[3], job.Billing) { t.Fatalf("args=%#v", querier.args) } } func TestJobsAdapterStateWritesFailWhenJobIsMissing(t *testing.T) { database := NewDatabase(Config{Backend: BackendPostgres}, &jobQuerier{rows: &jobRows{}}) if err := database.WriteBilling(context.Background(), "missing", json.RawMessage(`{}`)); err == nil { t.Fatal("WriteBilling error=nil") } if err := database.WriteOutputAssetIDs(context.Background(), "missing", []string{}); err == nil { t.Fatal("WriteOutputAssetIDs error=nil") } } func TestJobsAdapterFailsClosedWithoutPostgres(t *testing.T) { database := NewDatabase(Config{Backend: BackendLocal}, nil) if _, err := database.ClaimJobs(context.Background(), "worker", 1, 300); err == nil { t.Fatal("ClaimJobs error=nil") } if _, found, err := database.FindJob(context.Background(), jobs.Scope{OwnerID: "o"}, "j"); err == nil || found { t.Fatalf("FindJob found=%v err=%v", found, err) } } func jobRow(now time.Time) []any { return []any{ "job-1", "api:client", "client", "image.generate", "mock", "fixture", "queued", "prompt", []string{}, []string{}, []string{}, nil, nil, []byte(`{"input":true}`), nil, nil, nil, "idem-1", "fingerprint", 5, 0, 3, now, nil, nil, nil, nil, "https://hooks.example.test", 0, nil, []byte(`{"source":"api","accountId":"client","displayName":"client"}`), nil, now, nil, now, now, } } type jobQuerier struct { rows *jobRows err error query string args []any } func (querier *jobQuerier) Query(_ context.Context, query string, args ...any) (Rows, error) { querier.query, querier.args = query, args return querier.rows, querier.err } type jobRows struct { rows [][]any index int err error } func (rows *jobRows) Close() {} func (rows *jobRows) Err() error { return rows.err } func (rows *jobRows) Next() bool { return rows.index < len(rows.rows) } func (rows *jobRows) Scan(dest ...any) error { if rows.index >= len(rows.rows) { return errors.New("scan past end") } row := rows.rows[rows.index] rows.index++ if len(row) != len(dest) { return errors.New("scan arity mismatch") } for index, value := range row { switch target := dest[index].(type) { case *string: if value != nil { *target = value.(string) } case *int: *target = value.(int) case *int64: *target = value.(int64) case *time.Time: if value != nil { *target = value.(time.Time) } case *sql.NullString: if value != nil { *target = sql.NullString{String: value.(string), Valid: true} } case *sql.NullTime: if value != nil { *target = sql.NullTime{Time: value.(time.Time), Valid: true} } case *[]string: *target = append([]string(nil), value.([]string)...) case *[]byte: if value != nil { *target = append([]byte(nil), value.([]byte)...) } case *any: *target = value default: return errors.New("unsupported scan target") } } return nil } func TestJobFixtureJSONIsValid(t *testing.T) { if !json.Valid([]byte(`{"status":"queued"}`)) { t.Fatal("impossible") } }