easyai-ai-gateway/apps/api/internal/store/model_identity_integration_test.go

package store

import (
	"context"
	"errors"
	"os"
	"strings"
	"testing"
	"time"
)

func TestModelIdentityRoutingAndDeleteProtection(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 model identity PostgreSQL integration tests")
	}

	ctx := context.Background()
	db, err := Connect(ctx, databaseURL)
	if err != nil {
		t.Fatalf("connect test database: %v", err)
	}
	defer db.Close()

	suffix := time.Now().UTC().Format("20060102150405.000000000")
	invocationName := "official-model-" + suffix
	legacyAlias := "Legacy Model " + suffix
	displayName := "Model Card " + suffix
	baseModel, err := db.CreateBaseModel(ctx, BaseModelInput{
		ProviderKey:       "openai",
		CanonicalModelKey: "openai:" + invocationName,
		InvocationName:    invocationName,
		ProviderModelName: "provider/" + invocationName,
		ModelType:         StringList{"text_generate"},
		DisplayName:       displayName,
		LegacyAliases:     StringList{legacyAlias},
		Status:            "active",
	})
	if err != nil {
		t.Fatalf("create base model: %v", err)
	}
	defer func() {
		_, _ = db.pool.Exec(ctx, `DELETE FROM platform_models WHERE base_model_id = $1::uuid`, baseModel.ID)
		_, _ = db.pool.Exec(ctx, `DELETE FROM base_model_catalog WHERE id = $1::uuid`, baseModel.ID)
	}()

	var platformID string
	if err := db.pool.QueryRow(ctx, `
SELECT id::text
FROM integration_platforms
WHERE status = 'enabled' AND deleted_at IS NULL
ORDER BY created_at ASC
LIMIT 1`).Scan(&platformID); err != nil {
		t.Fatalf("find enabled platform: %v", err)
	}
	var platformModelID string
	if err := db.pool.QueryRow(ctx, `
INSERT INTO platform_models (
  platform_id, base_model_id, model_name, provider_model_name, model_alias,
  model_type, display_name, enabled
)
VALUES ($1::uuid, $2::uuid, $3, $4, $3, '["text_generate"]'::jsonb, $5, true)
RETURNING id::text`, platformID, baseModel.ID, invocationName, "provider/"+invocationName, displayName).Scan(&platformModelID); err != nil {
		t.Fatalf("bind platform model: %v", err)
	}

	if err := db.DeleteBaseModel(ctx, baseModel.ID); !errors.Is(err, ErrBaseModelInUse) {
		t.Fatalf("bound base model deletion should be protected: %v", err)
	}

	officialCandidates, err := db.ListModelCandidates(ctx, invocationName, "text_generate", nil)
	if err != nil || len(officialCandidates) != 1 || officialCandidates[0].PlatformModelID != platformModelID {
		t.Fatalf("official invocation should resolve the bound source: candidates=%+v err=%v", officialCandidates, err)
	}
	legacyCandidates, err := db.ListModelCandidates(ctx, legacyAlias, "text_generate", nil)
	if err != nil || len(legacyCandidates) != 1 || !legacyCandidates[0].LegacyAliasUsed {
		t.Fatalf("registered legacy alias should resolve and be marked: candidates=%+v err=%v", legacyCandidates, err)
	}
	if _, err := db.ListModelCandidates(ctx, displayName, "text_generate", nil); !errors.Is(err, ErrNoModelCandidate) {
		t.Fatalf("display name must not become an invocation alias: %v", err)
	}

	var legacyHits int64
	if err := db.pool.QueryRow(ctx, `
SELECT hit_count
FROM model_alias_usage_metrics
WHERE alias = $1 AND canonical_model_key = $2 AND model_type = 'text_generate'`, legacyAlias, baseModel.CanonicalModelKey).Scan(&legacyHits); err != nil {
		t.Fatalf("read legacy alias metric: %v", err)
	}
	if legacyHits != 1 {
		t.Fatalf("expected one legacy alias hit, got %d", legacyHits)
	}

	deprecatedInvocation := "deprecated-model-" + suffix
	deprecatedBase, err := db.CreateBaseModel(ctx, BaseModelInput{
		ProviderKey:       "openai",
		CanonicalModelKey: "openai:" + deprecatedInvocation,
		InvocationName:    deprecatedInvocation,
		ProviderModelName: deprecatedInvocation,
		ModelType:         StringList{"text_generate"},
		DisplayName:       "Deprecated Model",
		Status:            "deprecated",
	})
	if err != nil {
		t.Fatalf("create deprecated base model: %v", err)
	}
	defer func() {
		_, _ = db.pool.Exec(ctx, `DELETE FROM base_model_catalog WHERE id = $1::uuid`, deprecatedBase.ID)
	}()
	if _, err := db.CreatePlatformModel(ctx, CreatePlatformModelInput{
		PlatformID:        platformID,
		BaseModelID:       deprecatedBase.ID,
		ModelName:         deprecatedInvocation,
		ProviderModelName: deprecatedInvocation,
		ModelType:         StringList{"text_generate"},
	}); !errors.Is(err, ErrInvalidPlatformModelConfiguration) {
		t.Fatalf("deprecated base model must reject new bindings: %v", err)
	}
}