ai-gateway-go/internal/scheduler/scheduler_integration_test.go

package scheduler

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"
	"time"

	"aigateway.local/core/internal/platform/config"
	"aigateway.local/core/internal/platform/cryptox"
	"aigateway.local/core/internal/platform/database"
)

func TestSchedulerPostgreSQLLifecycle(t *testing.T) {
	databaseURL := os.Getenv("SCHEDULER_TEST_DATABASE_URL")
	if databaseURL == "" {
		t.Skip("SCHEDULER_TEST_DATABASE_URL is not set")
	}
	ctx := context.Background()
	pool, err := database.Open(ctx, config.Database{URL: databaseURL, MaxConns: 8})
	if err != nil {
		t.Fatal(err)
	}
	defer pool.Close()
	adminID := "64444444-4444-4444-8444-444444444444"
	appID := "65555555-5555-4555-8555-555555555555"
	versionID := "66666666-6666-4666-8666-666666666666"
	cleanup := func() {
		_, _ = pool.Exec(ctx, `DELETE FROM gateway.scheduled_tasks WHERE code='scheduler_test_task'`)
		_, _ = pool.Exec(ctx, `DELETE FROM gateway.applications WHERE id=$1`, appID)
		_, _ = pool.Exec(ctx, `DELETE FROM gateway.admin_accounts WHERE id=$1`, adminID)
	}
	cleanup()
	defer cleanup()
	if _, err = pool.Exec(ctx, `INSERT INTO gateway.admin_accounts(id,username,password_hash,role,active) VALUES($1,'scheduler-test-admin','test','superadmin',true)`, adminID); err != nil {
		t.Fatal(err)
	}
	tx, err := pool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	configJSON := `{"model":"test-model","knowledge_base_ids":[],"tool_ids":[],"retrieval_top_k":4,"temperature":0.2,"max_tool_rounds":1}`
	if _, err = tx.Exec(ctx, `INSERT INTO gateway.applications(id,code,name,status,draft_config,created_by) VALUES($1,'scheduler_test_app','Scheduler Test App','active',$2,$3)`, appID, configJSON, adminID); err != nil {
		t.Fatal(err)
	}
	if _, err = tx.Exec(ctx, `INSERT INTO gateway.application_versions(id,application_id,version,config,published_by) VALUES($1,$2,1,$3,$4)`, versionID, appID, configJSON, adminID); err != nil {
		t.Fatal(err)
	}
	if _, err = tx.Exec(ctx, `UPDATE gateway.applications SET published_version=1 WHERE id=$1`, appID); err != nil {
		t.Fatal(err)
	}
	if err = tx.Commit(ctx); err != nil {
		t.Fatal(err)
	}

	requestCount := 0
	failRequests := false
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		requestCount++
		if r.URL.Path != "/v1/applications/scheduler_test_app/chat/completions" || r.Header.Get("X-Gateway-API-Key") != "gw_scheduler_test" {
			http.Error(w, "unexpected request", http.StatusUnauthorized)
			return
		}
		if failRequests {
			http.Error(w, `{"error":"temporary failure"}`, http.StatusServiceUnavailable)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"scheduled answer"}}]}`))
	}))
	defer server.Close()
	cipher, err := cryptox.NewKeyring("AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=", 1, "", "scheduled-task-api-key")
	if err != nil {
		t.Fatal(err)
	}
	service := NewService(pool, cipher)
	task, err := service.Save(ctx, "", TaskInput{Code: "scheduler_test_task", Name: "Scheduler Test", CronExpression: "*/5 * * * *", Timezone: "UTC", TargetType: "application", TargetCode: "scheduler_test_app", Prompt: "create report", Variables: json.RawMessage(`{"scope":"daily"}`), APIKey: "gw_scheduler_test"}, adminID)
	if err != nil {
		t.Fatal(err)
	}
	if !task.HasAPIKey || task.Enabled {
		t.Fatalf("unexpected task: %+v", task)
	}
	if _, err = service.QueueManual(ctx, task.ID); err != nil {
		t.Fatal(err)
	}
	engine := NewEngine(service, server.URL, "integration-worker", 10, 3, time.Minute, nil)
	processed, err := engine.Tick(ctx, time.Now())
	if err != nil || processed != 1 || requestCount != 1 {
		t.Fatalf("tick processed=%d requests=%d err=%v", processed, requestCount, err)
	}
	runs, err := service.Runs(ctx, task.ID, 10)
	if err != nil || len(runs) != 1 || runs[0].Status != "success" || responseAnswer(runs[0].Response) != "scheduled answer" {
		t.Fatalf("runs=%+v err=%v", runs, err)
	}
	var completed bool
	if err = pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM gateway.outbox_events WHERE aggregate_type='scheduled_task' AND aggregate_id=$1 AND event_type='scheduled_task.completed')`, runs[0].ID).Scan(&completed); err != nil || !completed {
		t.Fatalf("completion event=%v err=%v", completed, err)
	}

	// Force a due schedule in the past. Tick must enqueue it once, execute it,
	// and advance next_run_at beyond now rather than replay every missed slot.
	if _, err = service.SetEnabled(ctx, task.ID, true); err != nil {
		t.Fatal(err)
	}
	forcedNow := time.Now().UTC()
	if _, err = pool.Exec(ctx, `UPDATE gateway.scheduled_tasks SET next_run_at=$2 WHERE id=$1`, task.ID, forcedNow.Add(-time.Minute)); err != nil {
		t.Fatal(err)
	}
	processed, err = engine.Tick(ctx, forcedNow)
	if err != nil || processed != 1 || requestCount != 2 {
		t.Fatalf("due tick processed=%d requests=%d err=%v", processed, requestCount, err)
	}
	task, err = service.Get(ctx, task.ID)
	if err != nil || task.NextRunAt == nil || !task.NextRunAt.After(forcedNow) {
		t.Fatalf("next run was not advanced: %+v err=%v", task.NextRunAt, err)
	}

	// Ordinary gateway failures remain pending until the configured attempt
	// limit, then become failed and emit exactly one terminal event.
	if _, err = service.SetEnabled(ctx, task.ID, false); err != nil {
		t.Fatal(err)
	}
	failRequests = true
	failedRun, err := service.QueueManual(ctx, task.ID)
	if err != nil {
		t.Fatal(err)
	}
	for attempt := 1; attempt <= 3; attempt++ {
		processed, err = engine.Tick(ctx, time.Now())
		if err != nil || processed != 1 {
			t.Fatalf("retry tick attempt=%d processed=%d err=%v", attempt, processed, err)
		}
		current, getErr := service.getRun(ctx, failedRun.ID)
		wantStatus := "pending"
		if attempt == 3 {
			wantStatus = "failed"
		}
		if getErr != nil || current.Status != wantStatus || current.Attempts != attempt {
			t.Fatalf("retry attempt=%d run=%+v err=%v", attempt, current, getErr)
		}
	}
	var failedEvents int
	if err = pool.QueryRow(ctx, `SELECT count(*) FROM gateway.outbox_events WHERE aggregate_type='scheduled_task' AND aggregate_id=$1 AND event_type='scheduled_task.failed'`, failedRun.ID).Scan(&failedEvents); err != nil || failedEvents != 1 {
		t.Fatalf("failure events=%d err=%v", failedEvents, err)
	}

	// A worker lease that stays running past its timeout is finalized through
	// the same failed-run and outbox path once it reaches the attempt limit.
	staleRun, err := service.QueueManual(ctx, task.ID)
	if err != nil {
		t.Fatal(err)
	}
	if _, err = pool.Exec(ctx, `UPDATE gateway.scheduled_task_runs SET status='running',attempts=3,worker_id='dead-worker',started_at=$2 WHERE id=$1`, staleRun.ID, time.Now().Add(-2*time.Minute)); err != nil {
		t.Fatal(err)
	}
	processed, err = engine.Tick(ctx, time.Now())
	if err != nil || processed != 0 {
		t.Fatalf("stale tick processed=%d err=%v", processed, err)
	}
	staleRun, err = service.getRun(ctx, staleRun.ID)
	if err != nil || staleRun.Status != "failed" || staleRun.Error != "执行超时且达到最大重试次数" {
		t.Fatalf("stale run=%+v err=%v", staleRun, err)
	}
}