wangbo
9a3cc582ce
AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED=false 时仍初始化 River schema 与控制客户端，允许 HTTP 请求持久化并入队；仅跳过执行客户端和动态容量 Worker。关闭上下文时同步清理控制客户端引用。\n\n新增数据库集成测试覆盖无执行 Worker 时成功创建 River job 且 attempted_by 为空。验证：gofmt、go test ./... -count=1、go vet ./...。
75 lines
2.2 KiB
Go
75 lines
2.2 KiB
Go
package runner
|
|
|
|
import (
|
|
"context"
|
|
"io"
|
|
"log/slog"
|
|
"os"
|
|
"strconv"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
|
)
|
|
|
|
func TestAsyncQueueClientEnqueuesWithoutExecutionWorker(t *testing.T) {
|
|
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
|
|
if databaseURL == "" {
|
|
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run the queue client integration test")
|
|
}
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
t.Cleanup(cancel)
|
|
db, err := store.Connect(ctx, databaseURL)
|
|
if err != nil {
|
|
t.Fatalf("connect store: %v", err)
|
|
}
|
|
t.Cleanup(db.Close)
|
|
|
|
suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
|
|
task, err := db.CreateTask(ctx, store.CreateTaskInput{
|
|
Kind: "images.edits",
|
|
Model: "queue-client-" + suffix,
|
|
Request: map[string]any{"prompt": "enqueue without execution worker"},
|
|
Async: true,
|
|
RunMode: "simulation",
|
|
}, &auth.User{ID: "queue-client-" + suffix, Source: "gateway"})
|
|
if err != nil {
|
|
t.Fatalf("create task: %v", err)
|
|
}
|
|
t.Cleanup(func() {
|
|
_, _ = db.Pool().Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id = $1::uuid`, task.ID)
|
|
})
|
|
|
|
service := New(config.Config{AppEnv: "test"}, db, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
|
service.StartAsyncQueueClient(ctx)
|
|
if err := service.EnqueueAsyncTask(ctx, task); err != nil {
|
|
t.Fatalf("enqueue task without execution worker: %v", err)
|
|
}
|
|
|
|
queued, err := db.GetTask(ctx, task.ID)
|
|
if err != nil {
|
|
t.Fatalf("load queued task: %v", err)
|
|
}
|
|
if queued.RiverJobID <= 0 {
|
|
t.Fatal("queue client did not persist a River job ID")
|
|
}
|
|
t.Cleanup(func() {
|
|
_, _ = db.Pool().Exec(context.Background(), `DELETE FROM river_job WHERE id = $1`, queued.RiverJobID)
|
|
})
|
|
|
|
var attemptedByCount int
|
|
if err := db.Pool().QueryRow(ctx, `
|
|
SELECT cardinality(attempted_by)
|
|
FROM river_job
|
|
WHERE id = $1`, queued.RiverJobID).Scan(&attemptedByCount); err != nil {
|
|
t.Fatalf("load River job: %v", err)
|
|
}
|
|
if attemptedByCount != 0 {
|
|
t.Fatalf("control-only queue client executed the job: attempted_by=%d", attemptedByCount)
|
|
}
|
|
}
|