package runner import ( "errors" "testing" "time" "github.com/easyai/easyai-ai-gateway/apps/api/internal/clients" "github.com/easyai/easyai-ai-gateway/apps/api/internal/store" "github.com/jackc/pgx/v5/pgconn" ) func TestAsyncAdmissionTaskErrorDisposition(t *testing.T) { for _, testCase := range []struct { name string err error terminal bool code string }{ {name: "bare no candidate", err: store.ErrNoModelCandidate, terminal: true, code: "no_model_candidate"}, {name: "cooldown", err: &store.ModelCandidateUnavailableError{Code: "model_cooling_down", Message: "cooling", RetryAfter: time.Minute}, code: "model_cooling_down"}, {name: "retryable rate limit", err: &store.RateLimitExceededError{Retryable: true}, code: "gateway_rate_limited"}, {name: "invalid parameter", err: &clients.ClientError{Code: "invalid_parameter", Retryable: false}, terminal: true, code: "invalid_parameter"}, {name: "retryable client", err: &clients.ClientError{Code: "network", Retryable: true}, code: "network"}, {name: "unknown", err: errors.New("unknown dispatcher failure"), code: "client_error"}, } { t.Run(testCase.name, func(t *testing.T) { if got := asyncAdmissionTaskErrorTerminal(testCase.err); got != testCase.terminal { t.Fatalf("terminal=%v, want %v", got, testCase.terminal) } if got := asyncAdmissionTaskErrorCode(testCase.err); got != testCase.code { t.Fatalf("code=%q, want %q", got, testCase.code) } }) } } func TestAsyncAdmissionSystemErrorClassification(t *testing.T) { for _, testCase := range []struct { name string code string system bool }{ {name: "connection", code: "08006", system: true}, {name: "deadlock", code: "40P01", system: true}, {name: "resource exhaustion", code: "53300", system: true}, {name: "task data violation", code: "22023", system: false}, {name: "task unique violation", code: "23505", system: false}, } { t.Run(testCase.name, func(t *testing.T) { err := &pgconn.PgError{Code: testCase.code} if got := isAsyncAdmissionSystemError(err); got != testCase.system { t.Fatalf("SQLSTATE %s system=%v, want %v", testCase.code, got, testCase.system) } }) } } func TestAsyncAdmissionRetryAtUsesRecoveryAndBoundedBackoff(t *testing.T) { now := time.Unix(1000, 0) recoveryAt := now.Add(7 * time.Minute) if got := asyncAdmissionRetryAt(now, "task-recovery", 9, &store.ModelCandidateUnavailableError{ Message: "cooling", RetryAfter: time.Minute, RecoveryAt: recoveryAt, }); !got.Equal(recoveryAt) { t.Fatalf("recovery retry=%s, want %s", got, recoveryAt) } previous := time.Duration(0) for failureCount := 0; failureCount < 10; failureCount++ { delay := asyncAdmissionRetryAt(now, "stable-task", failureCount, errors.New("boom")).Sub(now) if delay < time.Second || delay > time.Minute { t.Fatalf("failure %d delay=%s outside jittered 1s..60s bounds", failureCount, delay) } if delay < previous { t.Fatalf("failure %d delay=%s decreased from %s", failureCount, delay, previous) } previous = delay } } func TestAdmissionInputFromSnapshotRefreshesWorkerCapacity(t *testing.T) { admission := store.TaskAdmission{ TaskID: "task-1", PlatformID: "platform-1", PlatformModelID: "model-1", UserGroupID: "group-1", QueueKey: "queue-1", Mode: "async", Priority: 7, Scopes: []store.AdmissionScope{ {ScopeType: "platform_model", ScopeKey: "model-1", ConcurrentLimit: 10}, {ScopeType: "worker_capacity", ScopeKey: "global", ConcurrentLimit: 48}, }, ReselectRequestedAt: time.Time{}, } input := admissionInputFromSnapshot(admission, 24) if input.TaskID != admission.TaskID || input.PlatformModelID != admission.PlatformModelID || len(input.Scopes) != 2 { t.Fatalf("unexpected admission input: %+v", input) } if input.Scopes[1].ConcurrentLimit != 24 { t.Fatalf("worker capacity=%v, want 24", input.Scopes[1].ConcurrentLimit) } if input.Scopes[0].ConcurrentLimit != 10 { t.Fatalf("business scope changed: %+v", input.Scopes[0]) } if admission.Scopes[1].ConcurrentLimit != 48 { t.Fatalf("snapshot was mutated: %+v", admission.Scopes[1]) } } func TestDistributedAdmissionModelTypeBoundary(t *testing.T) { for _, modelType := range []string{ "image_generate", "image_edit", "image_analysis", "image_vectorize", "video_generate", "video_enhance", "image_to_video", "text_to_video", "video_edit", "video_reference", "video_first_last_frame", "video_understanding", "omni_video", "omni", "audio_generate", "audio_understanding", "text_to_speech", "voice_clone", } { if !distributedAdmissionModelType(modelType) { t.Errorf("%s should use distributed admission", modelType) } } for _, modelType := range []string{"text_generate", "embedding", "rerank", "", "unknown"} { if distributedAdmissionModelType(modelType) { t.Errorf("%s should preserve immediate execution/rate-limit behavior", modelType) } } } func TestAcceptanceAdmissionScopesIsolateRunWithoutChangingLimits(t *testing.T) { input := []store.AdmissionScope{{ ScopeType: "platform_model", ScopeKey: "model-1", ConcurrentLimit: 10, QueueLimit: 0, MaxWaitSeconds: 0, }, { ScopeType: "worker_capacity", ScopeKey: "global", ConcurrentLimit: 48, }} got := acceptanceAdmissionScopes(store.GatewayTask{RunMode: "acceptance", AcceptanceRunID: "run-1"}, input) if got[0].ConcurrentLimit != 10 || got[0].ScopeKey != "acceptance:run-1:model-1" { t.Fatalf("protocol-emulated acceptance must preserve the limit in an isolated scope, got %+v", got[0]) } if got[0].QueueLimit != acceptanceQueueLimit || got[0].MaxWaitSeconds != acceptanceQueueMaxWait { t.Fatalf("acceptance must enable a bounded queue, got %+v", got[0]) } if got[1].ConcurrentLimit != 48 || got[1].ScopeKey != "acceptance:run-1:global" { t.Fatalf("acceptance worker capacity changed, got %+v", got[1]) } if input[0].QueueLimit != 0 || input[0].MaxWaitSeconds != 0 { t.Fatalf("acceptance queue overlay must not mutate the production scopes, got %+v", input[0]) } } func TestAcceptanceCanaryPreservesProductionConcurrency(t *testing.T) { input := []store.AdmissionScope{{ ScopeType: "platform_model", ScopeKey: "model-1", ConcurrentLimit: 10, }} got := acceptanceAdmissionScopes(store.GatewayTask{RunMode: "acceptance_canary"}, input) if got[0].ConcurrentLimit != 10 { t.Fatalf("real acceptance canary must preserve production concurrency, got %+v", got[0]) } } func TestAcceptanceInfrastructureReservationsIsolateEveryMetric(t *testing.T) { input := []store.RateLimitReservation{ {ScopeType: "platform_model", ScopeKey: "model-1", Metric: "concurrent", Limit: 10}, {ScopeType: "platform_model", ScopeKey: "model-1", Metric: "rpm", Limit: 600}, {ScopeType: "user_group", ScopeKey: "group-1", Metric: "concurrent", Limit: 20}, } got := acceptanceInfrastructureReservations(store.GatewayTask{RunMode: "acceptance", AcceptanceRunID: "run-1"}, input) if len(got) != 3 || got[0].ScopeKey != "acceptance:run-1:model-1" || got[1].ScopeKey != "acceptance:run-1:model-1" || got[2].ScopeKey != "acceptance:run-1:group-1" { t.Fatalf("unexpected acceptance reservations: %+v", got) } canary := acceptanceInfrastructureReservations(store.GatewayTask{RunMode: "acceptance_canary"}, input) if len(canary) != len(input) { t.Fatalf("real canary reservations changed: %+v", canary) } } func TestAcceptanceQuotaScopePrefixMatchesIsolatedCounters(t *testing.T) { acceptanceTask := store.GatewayTask{RunMode: "acceptance", AcceptanceRunID: "run-1"} if got := acceptanceQuotaScopePrefix(acceptanceTask); got != "acceptance:run-1:" { t.Fatalf("acceptance quota scope prefix=%q", got) } if got := acceptanceQuotaScopePrefix(store.GatewayTask{RunMode: "acceptance_canary", AcceptanceRunID: "run-1"}); got != "" { t.Fatalf("canary must use production quota scope, got %q", got) } if got := acceptanceQuotaScopePrefix(store.GatewayTask{RunMode: "production"}); got != "" { t.Fatalf("production quota scope changed, got %q", got) } } func TestAcceptanceAdmissionScopesLeaveProductionPolicyUnchanged(t *testing.T) { input := []store.AdmissionScope{{ ScopeType: "platform_model", ScopeKey: "model-1", ConcurrentLimit: 10, QueueLimit: 3, MaxWaitSeconds: 7, }} got := acceptanceAdmissionScopes(store.GatewayTask{RunMode: "production"}, input) if got[0].QueueLimit != 3 || got[0].MaxWaitSeconds != 7 || got[0].ConcurrentLimit != 10 { t.Fatalf("production admission policy changed: %+v", got[0]) } } func TestAsyncAdmissionBatchCoversCertifiedGlobalCapacity(t *testing.T) { for _, testCase := range []struct { globalHardLimit int want int }{ {globalHardLimit: 48, want: 48}, {globalHardLimit: 64, want: 64}, {globalHardLimit: 128, want: 64}, {globalHardLimit: 0, want: 64}, } { if got := asyncAdmissionBatchLimit(testCase.globalHardLimit); got != testCase.want { t.Fatalf("async admission batch limit for %d = %d, want %d", testCase.globalHardLimit, got, testCase.want) } } } func TestPinCandidatesToTaskAdmissionPreservesAdmittedCandidate(t *testing.T) { input := []store.RuntimeModelCandidate{ {PlatformID: "platform-a", PlatformModelID: "model-a"}, {PlatformID: "platform-b", PlatformModelID: "model-b"}, {PlatformID: "platform-c", PlatformModelID: "model-c"}, } admission := &store.TaskAdmission{ Status: "admitted", PlatformID: "platform-b", PlatformModelID: "model-b", } got, pinned := pinCandidatesToTaskAdmission(input, admission) if !pinned { t.Fatal("expected admitted candidate to be pinned") } if got[0].PlatformModelID != "model-b" || got[1].PlatformModelID != "model-a" || got[2].PlatformModelID != "model-c" { t.Fatalf("unexpected candidate order: %+v", got) } if input[0].PlatformModelID != "model-a" { t.Fatalf("candidate pinning mutated the caller slice: %+v", input) } } func TestPinCandidatesToTaskAdmissionPreservesWaitingCandidate(t *testing.T) { input := []store.RuntimeModelCandidate{ {PlatformID: "platform-a", PlatformModelID: "model-a"}, {PlatformID: "platform-b", PlatformModelID: "model-b"}, } admission := &store.TaskAdmission{ Status: "waiting", PlatformID: "platform-b", PlatformModelID: "model-b", } got, pinned := pinCandidatesToTaskAdmission(input, admission) if !pinned || got[0].PlatformModelID != "model-b" { t.Fatalf("waiting candidate was not pinned: pinned=%v candidates=%+v", pinned, got) } } func TestPinCandidatesToTaskAdmissionAllowsRequestedReselection(t *testing.T) { input := []store.RuntimeModelCandidate{ {PlatformID: "platform-a", PlatformModelID: "model-a"}, {PlatformID: "platform-b", PlatformModelID: "model-b"}, } admission := &store.TaskAdmission{ Status: "waiting", PlatformID: "platform-b", PlatformModelID: "model-b", ReselectRequestedAt: time.Now(), } got, pinned := pinCandidatesToTaskAdmission(input, admission) if pinned { t.Fatalf("reselection request unexpectedly pinned the old candidate: %+v", got) } if got[0].PlatformModelID != "model-a" { t.Fatalf("reselection changed the sorted candidate order: %+v", got) } } func TestPinCandidatesToTaskAdmissionAllowsProviderRateQuotaMigration(t *testing.T) { input := []store.RuntimeModelCandidate{ { PlatformID: "platform-b", PlatformModelID: "model-b", PlatformPriority: 200, LoadMetrics: store.RuntimeCandidateLoadMetrics{ RPMLimit: 12, RPMCurrent: 0, }, }, { PlatformID: "platform-a", PlatformModelID: "model-a", PlatformPriority: 100, LoadAvoided: true, LoadMetrics: store.RuntimeCandidateLoadMetrics{ RPMLimit: 4, RPMCurrent: 4, }, }, } admission := &store.TaskAdmission{ Status: "admitted", PlatformID: "platform-a", PlatformModelID: "model-a", } got, pinned := pinCandidatesToTaskAdmission(input, admission) if pinned { t.Fatalf("rate-full provider remained pinned: %+v", got) } if got[0].PlatformModelID != "model-b" { t.Fatalf("available fallback was not preserved: %+v", got) } } func TestPinCandidatesToTaskAdmissionKeepsConcurrencyLeaseBinding(t *testing.T) { input := []store.RuntimeModelCandidate{ {PlatformID: "platform-b", PlatformModelID: "model-b"}, { PlatformID: "platform-a", PlatformModelID: "model-a", LoadAvoided: true, LoadMetrics: store.RuntimeCandidateLoadMetrics{ ConcurrentLimit: 2, ConcurrentCurrent: 2, }, }, } admission := &store.TaskAdmission{ Status: "admitted", PlatformID: "platform-a", PlatformModelID: "model-a", } got, pinned := pinCandidatesToTaskAdmission(input, admission) if !pinned || got[0].PlatformModelID != "model-a" { t.Fatalf("valid concurrency lease binding was not preserved: pinned=%v candidates=%+v", pinned, got) } } func TestPinCandidatesToTaskAdmissionIgnoresMissingCandidate(t *testing.T) { input := []store.RuntimeModelCandidate{ {PlatformID: "platform-a", PlatformModelID: "model-a"}, {PlatformID: "platform-b", PlatformModelID: "model-b"}, } admission := &store.TaskAdmission{ Status: "admitted", PlatformID: "platform-c", PlatformModelID: "model-c", } got, pinned := pinCandidatesToTaskAdmission(input, admission) if pinned { t.Fatalf("unexpected candidate pin for missing binding: %+v", got) } if got[0].PlatformModelID != "model-a" { t.Fatalf("candidate order changed for missing binding: %+v", got) } }