sub2api/backend/internal/handler/openai_models_handler_test.go

package handler

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

	middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
	"github.com/Wei-Shaw/sub2api/internal/service"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
)

func performOrdinaryPinnedModelsRequest(t *testing.T, codex *OpenAIGatewayHandler, group *service.Group, path, etag string) *httptest.ResponseRecorder {
	t.Helper()
	h := &GatewayHandler{openAIGatewayService: codex.gatewayService, maxAccountSwitches: 3}
	recorder := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(recorder)
	c.Request = httptest.NewRequest(http.MethodGet, path, nil)
	c.Request.Header.Set("If-None-Match", etag)
	c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{GroupID: &group.ID, Group: group})
	h.Models(c)
	return recorder
}

func ordinaryPinnedModelIDs(t *testing.T, recorder *httptest.ResponseRecorder) []string {
	t.Helper()
	var response gatewayModelsResponseForTest
	require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &response))
	require.NotNil(t, response.Data, "an empty catalog must be an array, not null")
	return modelIDsForTest(response.Data)
}

func TestOrdinaryPinnedModelsUsesSelectedAccountsAndFinalETag(t *testing.T) {
	accounts := []service.Account{
		newPinnedCodexAccount(1, service.StatusActive, true, false),
		newPinnedCodexAccount(2, service.StatusActive, true, true),
		newPinnedCodexAccount(3, service.StatusActive, true, false),
	}
	upstream := &codexModelsPinnedHTTPUpstream{bodies: map[int64]string{
		1: `{"data":[{"id":"unselected"}]}`,
		2: `{"data":[{"id":"special-model","owned_by":"first","created":123},{"id":"gpt-image-1"}]}`,
		3: `{"data":[{"id":"special-model","owned_by":"second"},{"id":"text-embedding-3-large"}]}`,
	}}
	h := newPinnedCodexTestHandler(accounts, upstream, 3)
	group := &service.Group{ID: 91, Platform: service.PlatformOpenAI,
		CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{Enabled: true, AccountIDs: []int64{2, 3}}}
	first := performOrdinaryPinnedModelsRequest(t, h, group, "/v1/models", "")
	require.Equal(t, http.StatusOK, first.Code, first.Body.String())
	require.Equal(t, []string{"special-model", "gpt-image-1", "text-embedding-3-large"}, ordinaryPinnedModelIDs(t, first))
	var response gatewayModelsResponseForTest
	require.NoError(t, json.Unmarshal(first.Body.Bytes(), &response))
	require.Equal(t, "model", response.Data[0].Object)
	require.Equal(t, "first", response.Data[0].OwnedBy)
	require.EqualValues(t, 123, response.Data[0].Created)
	require.Equal(t, []int64{2, 3}, upstream.accountIDs())
	alias := performOrdinaryPinnedModelsRequest(t, h, group, "/models", first.Header().Get("ETag"))
	require.Equal(t, http.StatusNotModified, alias.Code)
	require.Empty(t, alias.Body.String())
	require.Equal(t, []int64{2, 3}, upstream.accountIDs(), "fresh requests must not hit upstream")

	otherGroup := *group
	otherGroup.ID++
	otherGroup.ModelAllowlist = service.GroupModelAllowlist{Enabled: true, Models: []string{"text-embedding-3-large", "special-model"}}
	filtered := performOrdinaryPinnedModelsRequest(t, h, &otherGroup, "/v1/models", first.Header().Get("ETag"))
	require.Equal(t, http.StatusOK, filtered.Code)
	require.Equal(t, []string{"text-embedding-3-large", "special-model"}, ordinaryPinnedModelIDs(t, filtered))
	require.NotEqual(t, first.Header().Get("ETag"), filtered.Header().Get("ETag"))
	again := performOrdinaryPinnedModelsRequest(t, h, group, "/v1/models", "")
	require.Equal(t, first.Body.String(), again.Body.String(), "group filters must not modify account cache")
	require.Equal(t, []int64{2, 3}, upstream.accountIDs())
}

func TestOrdinaryPinnedModelsFailureAndEmptyPolicies(t *testing.T) {
	for _, tc := range []struct {
		name       string
		accountIDs []int64
		bodies     map[int64]string
		statuses   map[int64]int
		fallback   bool
		selected   []string
		wantStatus int
		wantIDs    []string
	}{
		{name: "partial failure", accountIDs: []int64{2, 3}, bodies: map[int64]string{2: `{"data":[{"id":"ok"}]}`}, statuses: map[int64]int{3: 503}, wantStatus: 200, wantIDs: []string{"ok"}},
		{name: "all failed", accountIDs: []int64{2, 3}, statuses: map[int64]int{2: 429, 3: 504}, wantStatus: 502},
		{name: "missing members", accountIDs: []int64{99}, wantStatus: 503},
		{name: "empty config", wantStatus: 503},
		{name: "fallback from missing members", accountIDs: []int64{99}, fallback: true, wantStatus: 200, wantIDs: []string{"from-scheduler"}},
		{name: "fallback from failure", accountIDs: []int64{2}, statuses: map[int64]int{2: 503}, fallback: true, wantStatus: 200, wantIDs: []string{"from-scheduler"}},
		{name: "valid empty", accountIDs: []int64{2}, bodies: map[int64]string{2: `{"data":[]}`}, fallback: true, wantStatus: 200},
		{name: "filtered empty", accountIDs: []int64{2}, bodies: map[int64]string{2: `{"data":[{"id":"ok"}]}`}, selected: []string{"absent"}, fallback: true, wantStatus: 200},
		{name: "invalid envelope", accountIDs: []int64{2}, bodies: map[int64]string{2: `{"error":"denied"}`}, wantStatus: 502},
		{name: "null data", accountIDs: []int64{2}, bodies: map[int64]string{2: `{"data":null}`}, wantStatus: 502},
	} {
		t.Run(tc.name, func(t *testing.T) {
			accounts := []service.Account{
				newPinnedCodexAccount(1, service.StatusActive, true, false),
				newPinnedCodexAccount(2, service.StatusActive, true, false),
				newPinnedCodexAccount(3, service.StatusActive, true, false),
			}
			bodies := map[int64]string{1: `{"data":[{"id":"from-scheduler"}]}`}
			for id, body := range tc.bodies {
				bodies[id] = body
			}
			upstream := &codexModelsPinnedHTTPUpstream{bodies: bodies, statuses: tc.statuses}
			h := newPinnedCodexTestHandler(accounts, upstream, 3)
			group := &service.Group{ID: 92, Platform: service.PlatformOpenAI,
				ModelAllowlist:            service.GroupModelAllowlist{Enabled: len(tc.selected) > 0, Models: tc.selected},
				CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{Enabled: true, AccountIDs: tc.accountIDs, FallbackToScheduler: tc.fallback}}
			recorder := performOrdinaryPinnedModelsRequest(t, h, group, "/v1/models", "")
			require.Equal(t, tc.wantStatus, recorder.Code, recorder.Body.String())
			if tc.wantStatus == 200 {
				require.Equal(t, len(tc.wantIDs), len(ordinaryPinnedModelIDs(t, recorder)))
				if len(tc.wantIDs) > 0 {
					require.Equal(t, tc.wantIDs, ordinaryPinnedModelIDs(t, recorder))
				} else {
					require.NotContains(t, upstream.accountIDs(), int64(1), "empty success must not trigger fallback")
				}
			}
		})
	}
}

func TestOrdinaryPinnedModelsSkipsPersistentlyUnavailableAccounts(t *testing.T) {
	accounts := []service.Account{
		newPinnedCodexAccount(1, service.StatusActive, true, true),
		newPinnedCodexAccount(2, service.StatusActive, false, false),
		newPinnedCodexAccount(3, service.StatusDisabled, true, false),
		newPinnedCodexAccount(4, service.StatusActive, true, false),
	}
	expired := time.Now().Add(-time.Hour)
	accounts[3].ExpiresAt = &expired
	accounts[3].AutoPauseOnExpired = true
	overloaded := time.Now().Add(time.Hour)
	accounts[0].OverloadUntil = &overloaded
	accounts[0].TempUnschedulableUntil = &overloaded
	upstream := &codexModelsPinnedHTTPUpstream{bodies: map[int64]string{1: `{"data":[{"id":"available"}]}`}}
	h := newPinnedCodexTestHandler(accounts, upstream, 3)
	group := &service.Group{ID: 93, Platform: service.PlatformOpenAI,
		CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{Enabled: true, AccountIDs: []int64{1, 2, 3, 4, 99}}}
	recorder := performOrdinaryPinnedModelsRequest(t, h, group, "/v1/models", "")
	require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
	require.Equal(t, []int64{1}, upstream.accountIDs())
}

func TestPinnedModelsAllowlistExpandsWildcardsForBothRepresentations(t *testing.T) {
	for _, codex := range []bool{false, true} {
		t.Run(map[bool]string{false: "ordinary", true: "codex"}[codex], func(t *testing.T) {
			accounts := []service.Account{newPinnedCodexAccount(1, service.StatusActive, true, false)}
			accounts[0].Credentials["model_mapping"] = map[string]any{
				"public-a": "gpt-5.5", "public-b": "gpt-5.5", "blocked": "gpt-5.5",
			}
			body := `{"data":[{"id":"gpt-5.5","owned_by":"provider"}]}`
			if codex {
				body = `{"models":[{"slug":"gpt-5.5","context_window":424242}]}`
			}
			upstream := &codexModelsPinnedHTTPUpstream{bodies: map[int64]string{1: body}}
			h := newPinnedCodexTestHandler(accounts, upstream, 3)
			group := &service.Group{ID: 95, Platform: service.PlatformOpenAI,
				ModelAllowlist:            service.GroupModelAllowlist{Enabled: true, Models: []string{"public-b", "public-*"}},
				CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{Enabled: true, AccountIDs: []int64{1}}}
			if codex {
				recorder := performPinnedCodexModelsRequest(t, h, group, "")
				require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
				require.Equal(t, []string{"public-b", "public-a"}, codexHandlerManifestSlugs(t, recorder))
				require.Contains(t, recorder.Body.String(), `"context_window":424242`)
			} else {
				recorder := performOrdinaryPinnedModelsRequest(t, h, group, "/v1/models", "")
				require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
				require.Equal(t, []string{"public-b", "public-a"}, ordinaryPinnedModelIDs(t, recorder))
				require.Contains(t, recorder.Body.String(), `"owned_by":"provider"`)
			}
		})
	}
}

func TestPinnedModelsMappingFollowsUpstreamDiscoveryForBothRepresentations(t *testing.T) {
	for _, codex := range []bool{false, true} {
		t.Run(map[bool]string{false: "ordinary", true: "codex"}[codex], func(t *testing.T) {
			accounts := []service.Account{
				newPinnedCodexAccount(1, service.StatusActive, true, false),
				newPinnedCodexAccount(2, service.StatusActive, true, false),
			}
			accounts[0].Credentials["model_mapping"] = map[string]any{"unselected-alias": "gpt-5.5"}
			accounts[1].Credentials["model_mapping"] = map[string]any{
				"public-alias": "gpt-5.5", "missing-alias": "missing-upstream", "custom-*": "gpt-5.5",
			}
			body := `{"data":[{"id":"gpt-5.5","owned_by":"provider","created":123},{"id":"unmapped-upstream"}]}`
			if codex {
				body = `{"models":[{"slug":"gpt-5.5","context_window":424242},{"slug":"unmapped-upstream"}]}`
			}
			upstream := &codexModelsPinnedHTTPUpstream{bodies: map[int64]string{2: body}}
			h := newPinnedCodexTestHandler(accounts, upstream, 3)
			group := &service.Group{ID: 94, Platform: service.PlatformOpenAI,
				ModelAllowlist:            service.GroupModelAllowlist{Enabled: true, Models: []string{"custom-concrete", "public-alias", "unselected-alias", "missing-alias"}},
				CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{Enabled: true, AccountIDs: []int64{2}}}
			var recorder *httptest.ResponseRecorder
			if codex {
				recorder = performPinnedCodexModelsRequest(t, h, group, "")
			} else {
				recorder = performOrdinaryPinnedModelsRequest(t, h, group, "/v1/models", "")
			}
			require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
			require.Equal(t, []int64{2}, upstream.accountIDs(), "mappings must not bypass pinned discovery")
			if codex {
				require.Equal(t, []string{"custom-concrete", "public-alias"}, codexHandlerManifestSlugs(t, recorder))
				require.Contains(t, recorder.Body.String(), `"context_window":424242`)
			} else {
				require.Equal(t, []string{"custom-concrete", "public-alias"}, ordinaryPinnedModelIDs(t, recorder))
				require.Contains(t, recorder.Body.String(), `"owned_by":"provider"`)
			}
		})
	}
}