166 lines
5.5 KiB
Go
166 lines
5.5 KiB
Go
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 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, []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, 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 *time.Time:
|
|
*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)...)
|
|
}
|
|
default:
|
|
return errors.New("unsupported scan target")
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func TestJobFixtureJSONIsValid(t *testing.T) {
|
|
if !json.Valid([]byte(`{"status":"queued"}`)) {
|
|
t.Fatal("impossible")
|
|
}
|
|
}
|