sub2api/backend/internal/handler/gateway_models_retrieve_test.go

package handler

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

	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 requestModelForTest(h *GatewayHandler, group *service.Group, modelID, etag string) *httptest.ResponseRecorder {
	rec := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(rec)
	c.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
	c.Request.Header.Set("If-None-Match", etag)
	if modelID != "" {
		c.Params = gin.Params{{Key: "model", Value: modelID}}
	}
	c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{GroupID: &group.ID, Group: group})
	h.Models(c)
	return rec
}

func TestRetrieveModelMatchesVisibleCatalogue(t *testing.T) {
	gin.SetMode(gin.TestMode)
	for _, platform := range []string{service.PlatformOpenAI, service.PlatformAnthropic, service.PlatformGemini, service.PlatformGrok, service.PlatformTypeSafe, service.PlatformComposite} {
		for _, mapped := range []bool{false, true} {
			name := platform + "/fallback"
			if mapped {
				name = platform + "/mapped"
			}
			t.Run(name, func(t *testing.T) {
				accountPlatform := platform
				if platform == service.PlatformComposite {
					accountPlatform = service.PlatformOpenAI
				}
				credentials := map[string]any{}
				if mapped {
					credentials["model_mapping"] = map[string]any{"custom-model": "upstream-only-model"}
				}
				group := &service.Group{ID: 71, Platform: platform}
				h := newGatewayModelsHandlerForTest(&gatewayModelsAccountRepoStub{byGroup: map[int64][]service.Account{
					group.ID: {{ID: 1, Platform: accountPlatform, Type: service.AccountTypeAPIKey, Status: service.StatusActive, Schedulable: true, Credentials: credentials}},
				}})
				list := requestModelForTest(h, group, "", "")
				require.Equal(t, http.StatusOK, list.Code, list.Body.String())
				var catalog struct {
					Data []json.RawMessage `json:"data"`
				}
				require.NoError(t, json.Unmarshal(list.Body.Bytes(), &catalog))
				require.NotEmpty(t, catalog.Data)
				var model struct {
					ID string `json:"id"`
				}
				require.NoError(t, json.Unmarshal(catalog.Data[0], &model))
				retrieved := requestModelForTest(h, group, model.ID, "")
				require.Equal(t, http.StatusOK, retrieved.Code, retrieved.Body.String())
				require.JSONEq(t, string(catalog.Data[0]), retrieved.Body.String())
				require.Equal(t, http.StatusNotFound, requestModelForTest(h, group, "unknown-model", "").Code)
				group.ModelAllowlist = service.GroupModelAllowlist{Enabled: true, Models: []string{"absent-from-source"}}
				hidden := requestModelForTest(h, group, model.ID, "")
				require.Equal(t, http.StatusNotFound, hidden.Code, hidden.Body.String())
				require.Contains(t, hidden.Body.String(), `"code":"model_not_found"`)
			})
		}
	}
}

func TestRetrievePinnedModelPreservesMetadataFilteringAndErrors(t *testing.T) {
	for _, tc := range []struct {
		name       string
		status     int
		selected   []string
		wantStatus int
	}{
		{name: "metadata", wantStatus: http.StatusOK},
		{name: "hidden", selected: []string{"other-model"}, wantStatus: http.StatusNotFound},
		{name: "upstream failure", status: http.StatusServiceUnavailable, wantStatus: http.StatusBadGateway},
	} {
		t.Run(tc.name, func(t *testing.T) {
			upstream := &codexModelsPinnedHTTPUpstream{bodies: map[int64]string{
				2: `{"data":[{"id":"special-model","ID":"not-a-model-id","owned_by":"source-owner","created":123,"extra":{"context":999}}]}`,
			}, statuses: map[int64]int{}}
			if tc.status != 0 {
				upstream.statuses[2] = tc.status
			}
			codex := newPinnedCodexTestHandler([]service.Account{newPinnedCodexAccount(2, service.StatusActive, true, false)}, upstream, 3)
			h := &GatewayHandler{openAIGatewayService: codex.gatewayService, maxAccountSwitches: 3}
			group := &service.Group{ID: 72, Platform: service.PlatformOpenAI,
				ModelAllowlist:            service.GroupModelAllowlist{Enabled: len(tc.selected) > 0, Models: tc.selected},
				CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{Enabled: true, AccountIDs: []int64{2}}}
			got := requestModelForTest(h, group, "special-model", "collection-etag")
			require.Equal(t, tc.wantStatus, got.Code, got.Body.String())
			if tc.wantStatus == http.StatusOK {
				list := requestModelForTest(h, group, "", "")
				var catalog struct {
					Data []json.RawMessage `json:"data"`
				}
				require.NoError(t, json.Unmarshal(list.Body.Bytes(), &catalog))
				require.JSONEq(t, string(catalog.Data[0]), got.Body.String())
				require.Contains(t, got.Body.String(), `"owned_by":"source-owner"`)
				require.Contains(t, got.Body.String(), `"created":123`)
				again := requestModelForTest(h, group, "special-model", list.Header().Get("ETag"))
				require.Equal(t, http.StatusOK, again.Code)
				require.Empty(t, again.Header().Get("ETag"))
				require.Equal(t, http.StatusNotFound, requestModelForTest(h, group, "not-a-model-id", "").Code)
			}
		})
	}
}