easyai-ai-gateway/apps/api/internal/runner/admission_test.go

package runner

import (
	"testing"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

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 TestAcceptanceAdmissionScopesEnableBoundedQueueWithoutChangingConcurrency(t *testing.T) {
	input := []store.AdmissionScope{{
		ScopeType:       "platform_model",
		ScopeKey:        "model-1",
		ConcurrentLimit: 10,
		QueueLimit:      0,
		MaxWaitSeconds:  0,
	}}

	got := acceptanceAdmissionScopes(store.GatewayTask{RunMode: "acceptance"}, input)

	if got[0].ConcurrentLimit != 10 {
		t.Fatalf("acceptance must preserve the production concurrency limit, 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 input[0].QueueLimit != 0 || input[0].MaxWaitSeconds != 0 {
		t.Fatalf("acceptance queue overlay must not mutate the production scopes, got %+v", input[0])
	}
}

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 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 TestPinCandidatesToTaskAdmissionIgnoresWaitingOrMissingCandidate(t *testing.T) {
	input := []store.RuntimeModelCandidate{
		{PlatformID: "platform-a", PlatformModelID: "model-a"},
		{PlatformID: "platform-b", PlatformModelID: "model-b"},
	}
	for name, admission := range map[string]*store.TaskAdmission{
		"waiting": {
			Status:          "waiting",
			PlatformID:      "platform-b",
			PlatformModelID: "model-b",
		},
		"missing": {
			Status:          "admitted",
			PlatformID:      "platform-c",
			PlatformModelID: "model-c",
		},
	} {
		t.Run(name, func(t *testing.T) {
			got, pinned := pinCandidatesToTaskAdmission(input, admission)
			if pinned {
				t.Fatalf("unexpected candidate pin for %s: %+v", name, got)
			}
			if got[0].PlatformModelID != "model-a" {
				t.Fatalf("candidate order changed for %s: %+v", name, got)
			}
		})
	}
}