internal/app/admin_models_test.go

package app

import (
	"context"
	"fmt"
	"net/http"
	"testing"
	"time"

	"ccLoad/internal/model"
	"ccLoad/internal/storage"

	"github.com/gin-gonic/gin"
)

func TestAdminModels_FetchModelsPreview(t *testing.T) {
	var gotAuth string
	upstream := newTestHTTPServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/v1/models" {
			http.NotFound(w, r)
			return
		}
		gotAuth = r.Header.Get("Authorization")
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"id":"gpt-4o"},{"id":"gpt-4o-mini"}]}`))
	}))
	t.Cleanup(upstream.Close)

	server, _, cleanup := setupAdminTestServer(t)
	defer cleanup()

	t.Run("invalid request", func(t *testing.T) {
		c, w := newTestContext(t, newJSONRequestBytes(http.MethodPost, "/admin/channels/models/fetch", []byte(`{}`)))

		server.HandleFetchModelsPreview(c)
		if w.Code != http.StatusBadRequest {
			t.Fatalf("status=%d, want %d", w.Code, http.StatusBadRequest)
		}
	})

	t.Run("success", func(t *testing.T) {
		payload := map[string]any{
			"channel_type": " openai ",
			"url":          upstream.URL,
			"api_key":      "sk-test",
		}
		c, w := newTestContext(t, newJSONRequest(t, http.MethodPost, "/admin/channels/models/fetch", payload))

		server.HandleFetchModelsPreview(c)
		if w.Code != http.StatusOK {
			t.Fatalf("status=%d, want %d, body=%s", w.Code, http.StatusOK, w.Body.String())
		}

		var resp struct {
			Success bool                `json:"success"`
			Data    FetchModelsResponse `json:"data"`
		}
		mustUnmarshalJSON(t, w.Body.Bytes(), &resp)
		if !resp.Success || resp.Data.Source != "api" || len(resp.Data.Models) != 2 {
			t.Fatalf("unexpected resp: %+v", resp)
		}
		if resp.Data.Models[0].RedirectModel != resp.Data.Models[0].Model {
			t.Fatalf("expected redirect_model filled, got %+v", resp.Data.Models[0])
		}
		if gotAuth != "Bearer sk-test" {
			t.Fatalf("Authorization=%q, want %q", gotAuth, "Bearer sk-test")
		}
	})

	t.Run("multi url fallback", func(t *testing.T) {
		failCalls := 0
		okCalls := 0

		failUpstream := newTestHTTPServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			failCalls++
			http.Error(w, "boom", http.StatusBadGateway)
		}))
		t.Cleanup(failUpstream.Close)

		okUpstream := newTestHTTPServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			okCalls++
			time.Sleep(15 * time.Millisecond)
			if r.URL.Path != "/v1/models" {
				http.NotFound(w, r)
				return
			}
			w.Header().Set("Content-Type", "application/json")
			_, _ = w.Write([]byte(`{"data":[{"id":"gpt-4.1-mini"}]}`))
		}))
		t.Cleanup(okUpstream.Close)

		payload := map[string]any{
			"channel_type": "openai",
			"url":          failUpstream.URL + "\n" + okUpstream.URL,
			"api_key":      "sk-test",
		}
		c, w := newTestContext(t, newJSONRequest(t, http.MethodPost, "/admin/channels/models/fetch", payload))

		server.HandleFetchModelsPreview(c)
		if w.Code != http.StatusOK {
			t.Fatalf("status=%d, want %d, body=%s", w.Code, http.StatusOK, w.Body.String())
		}

		var resp struct {
			Success bool                `json:"success"`
			Data    FetchModelsResponse `json:"data"`
		}
		mustUnmarshalJSON(t, w.Body.Bytes(), &resp)
		if !resp.Success || len(resp.Data.Models) != 1 || resp.Data.Models[0].Model != "gpt-4.1-mini" {
			t.Fatalf("unexpected resp: %+v", resp)
		}
		if failCalls < 1 || okCalls < 1 {
			t.Fatalf("expected fallback attempts, failCalls=%d okCalls=%d", failCalls, okCalls)
		}
	})
}

func TestAdminModels_HandleFetchModels(t *testing.T) {
	// upstream: 先返回成功，再返回错误
	var call int
	upstream := newTestHTTPServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/v1/models" {
			http.NotFound(w, r)
			return
		}
		call++
		if call == 1 {
			w.Header().Set("Content-Type", "application/json")
			_, _ = w.Write([]byte(`{"data":[{"id":"gpt-4o"}]}`))
			return
		}
		http.Error(w, "boom", http.StatusBadGateway)
	}))
	t.Cleanup(upstream.Close)

	server, store, cleanup := setupAdminTestServer(t)
	defer cleanup()

	// 需要 channelCache
	server.channelCache = storage.NewChannelCache(store, time.Minute)

	ctx := context.Background()
	cfg, err := store.CreateConfig(ctx, &model.Config{
		Name:         "c1",
		URL:          upstream.URL,
		Priority:     1,
		ChannelType:  "openai",
		ModelEntries: []model.ModelEntry{{Model: "m1"}},
		Enabled:      true,
	})
	if err != nil {
		t.Fatalf("CreateConfig failed: %v", err)
	}
	if err := store.CreateAPIKeysBatch(ctx, []*model.APIKey{
		{ChannelID: cfg.ID, KeyIndex: 0, APIKey: "sk-test", KeyStrategy: model.KeyStrategySequential},
	}); err != nil {
		t.Fatalf("CreateAPIKeysBatch failed: %v", err)
	}

	t.Run("success", func(t *testing.T) {
		c, w := newTestContext(t, newRequest(http.MethodGet, "/admin/channels/1/models/fetch", nil))
		c.Params = gin.Params{{Key: "id", Value: "1"}}

		server.HandleFetchModels(c)
		if w.Code != http.StatusOK {
			t.Fatalf("status=%d, want %d, body=%s", w.Code, http.StatusOK, w.Body.String())
		}
		var resp struct {
			Success bool                `json:"success"`
			Data    FetchModelsResponse `json:"data"`
		}
		mustUnmarshalJSON(t, w.Body.Bytes(), &resp)
		if !resp.Success || len(resp.Data.Models) != 1 || resp.Data.Models[0].Model != "gpt-4o" {
			t.Fatalf("unexpected resp: %+v", resp)
		}
	})

	t.Run("upstream error returns 200 with success=false", func(t *testing.T) {
		c, w := newTestContext(t, newRequest(http.MethodGet, "/admin/channels/1/models/fetch", nil))
		c.Params = gin.Params{{Key: "id", Value: "1"}}

		server.HandleFetchModels(c)
		if w.Code != http.StatusOK {
			t.Fatalf("status=%d, want %d", w.Code, http.StatusOK)
		}
		var resp struct {
			Success bool   `json:"success"`
			Error   string `json:"error"`
		}
		mustUnmarshalJSON(t, w.Body.Bytes(), &resp)
		if resp.Success || resp.Error == "" {
			t.Fatalf("expected success=false with error, got %+v", resp)
		}
	})
}

func TestAdminModels_HandleFetchModels_MultiURL(t *testing.T) {
	failCalls := 0
	failUpstream := newTestHTTPServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		failCalls++
		http.Error(w, "boom", http.StatusBadGateway)
	}))
	t.Cleanup(failUpstream.Close)

	okCalls := 0
	okUpstream := newTestHTTPServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		okCalls++
		time.Sleep(15 * time.Millisecond)
		if r.URL.Path != "/v1/models" {
			http.NotFound(w, r)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"id":"gpt-4.1"}]}`))
	}))
	t.Cleanup(okUpstream.Close)

	server, store, cleanup := setupAdminTestServer(t)
	defer cleanup()
	server.channelCache = storage.NewChannelCache(store, time.Minute)
	server.urlSelector = NewURLSelector()

	ctx := context.Background()
	cfg, err := store.CreateConfig(ctx, &model.Config{
		Name:         "multi-url-channel",
		URL:          failUpstream.URL + "\n" + okUpstream.URL,
		Priority:     1,
		ChannelType:  "openai",
		ModelEntries: []model.ModelEntry{{Model: "m1"}},
		Enabled:      true,
	})
	if err != nil {
		t.Fatalf("CreateConfig failed: %v", err)
	}
	if err := store.CreateAPIKeysBatch(ctx, []*model.APIKey{
		{ChannelID: cfg.ID, KeyIndex: 0, APIKey: "sk-test", KeyStrategy: model.KeyStrategySequential},
	}); err != nil {
		t.Fatalf("CreateAPIKeysBatch failed: %v", err)
	}
	// 强制第一跳命中失败URL，确保触发fallback与反馈逻辑
	server.urlSelector.CooldownURL(cfg.ID, okUpstream.URL)

	c, w := newTestContext(t, newRequest(http.MethodGet, "/admin/channels/1/models/fetch", nil))
	c.Params = gin.Params{{Key: "id", Value: fmt.Sprintf("%d", cfg.ID)}}

	server.HandleFetchModels(c)
	if w.Code != http.StatusOK {
		t.Fatalf("status=%d, want %d, body=%s", w.Code, http.StatusOK, w.Body.String())
	}

	var resp struct {
		Success bool                `json:"success"`
		Data    FetchModelsResponse `json:"data"`
	}
	mustUnmarshalJSON(t, w.Body.Bytes(), &resp)
	if !resp.Success {
		t.Fatalf("expected success=true, body=%s", w.Body.String())
	}
	if len(resp.Data.Models) != 1 || resp.Data.Models[0].Model != "gpt-4.1" {
		t.Fatalf("unexpected models: %+v", resp.Data.Models)
	}
	if failCalls < 1 || okCalls < 1 {
		t.Fatalf("expected both URLs attempted, failCalls=%d okCalls=%d", failCalls, okCalls)
	}
	if !server.urlSelector.IsCooledDown(cfg.ID, failUpstream.URL) {
		t.Fatalf("expected failed URL cooled down, url=%s", failUpstream.URL)
	}
	latency, exists := server.urlSelector.latencies[urlKey{channelID: cfg.ID, url: okUpstream.URL}]
	if !exists || latency == nil || latency.value <= 0 {
		t.Fatalf("expected success URL latency recorded, got=%v", latency)
	}
}

func TestAdminModels_HandleFetchModels_MultiURL_KeyErrorDoesNotCooldownURL(t *testing.T) {
	keyErrCalls := 0
	keyErrUpstream := newTestHTTPServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		keyErrCalls++
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusUnauthorized)
		_, _ = w.Write([]byte(`{"error":{"type":"authentication_error","message":"invalid api key"}}`))
	}))
	t.Cleanup(keyErrUpstream.Close)

	okCalls := 0
	okUpstream := newTestHTTPServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		okCalls++
		if r.URL.Path != "/v1/models" {
			http.NotFound(w, r)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"id":"gpt-4.1"}]}`))
	}))
	t.Cleanup(okUpstream.Close)

	server, store, cleanup := setupAdminTestServer(t)
	defer cleanup()
	server.channelCache = storage.NewChannelCache(store, time.Minute)
	server.urlSelector = NewURLSelector()

	ctx := context.Background()
	cfg, err := store.CreateConfig(ctx, &model.Config{
		Name:         "multi-url-key-error",
		URL:          keyErrUpstream.URL + "\n" + okUpstream.URL,
		Priority:     1,
		ChannelType:  "openai",
		ModelEntries: []model.ModelEntry{{Model: "m1"}},
		Enabled:      true,
	})
	if err != nil {
		t.Fatalf("CreateConfig failed: %v", err)
	}
	if err := store.CreateAPIKeysBatch(ctx, []*model.APIKey{
		{ChannelID: cfg.ID, KeyIndex: 0, APIKey: "sk-test", KeyStrategy: model.KeyStrategySequential},
	}); err != nil {
		t.Fatalf("CreateAPIKeysBatch failed: %v", err)
	}
	// 强制首跳优先命中 keyErrUpstream，覆盖“先401再fallback”的路径。
	server.urlSelector.CooldownURL(cfg.ID, okUpstream.URL)

	c, w := newTestContext(t, newRequest(http.MethodGet, "/admin/channels/1/models/fetch", nil))
	c.Params = gin.Params{{Key: "id", Value: fmt.Sprintf("%d", cfg.ID)}}

	server.HandleFetchModels(c)
	if w.Code != http.StatusOK {
		t.Fatalf("status=%d, want %d, body=%s", w.Code, http.StatusOK, w.Body.String())
	}

	var resp struct {
		Success bool                `json:"success"`
		Data    FetchModelsResponse `json:"data"`
	}
	mustUnmarshalJSON(t, w.Body.Bytes(), &resp)
	if !resp.Success {
		t.Fatalf("expected success=true, body=%s", w.Body.String())
	}
	if keyErrCalls < 1 || okCalls < 1 {
		t.Fatalf("expected both URLs attempted, keyErrCalls=%d okCalls=%d", keyErrCalls, okCalls)
	}
	if server.urlSelector.IsCooledDown(cfg.ID, keyErrUpstream.URL) {
		t.Fatalf("expected key-error URL not cooled down, url=%s", keyErrUpstream.URL)
	}
}

func TestAdminModels_HandleBatchRefreshModels(t *testing.T) {
	t.Run("merge mode partial success", func(t *testing.T) {
		// channel1: 返回 m1,m2（新增1个）
		upstream1 := newTestHTTPServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if r.URL.Path != "/v1/models" {
				http.NotFound(w, r)
				return
			}
			w.Header().Set("Content-Type", "application/json")
			_, _ = w.Write([]byte(`{"data":[{"id":"m1"},{"id":"m2"}]}`))
		}))
		t.Cleanup(upstream1.Close)

		// channel2: 返回 x1（无变化）
		upstream2 := newTestHTTPServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if r.URL.Path != "/v1/models" {
				http.NotFound(w, r)
				return
			}
			w.Header().Set("Content-Type", "application/json")
			_, _ = w.Write([]byte(`{"data":[{"id":"x1"}]}`))
		}))
		t.Cleanup(upstream2.Close)

		server, store, cleanup := setupAdminTestServer(t)
		defer cleanup()

		ctx := context.Background()
		c1, err := store.CreateConfig(ctx, &model.Config{
			Name:         "c1",
			URL:          upstream1.URL,
			Priority:     1,
			ChannelType:  "openai",
			ModelEntries: []model.ModelEntry{{Model: "m1"}},
			Enabled:      true,
		})
		if err != nil {
			t.Fatalf("CreateConfig c1 failed: %v", err)
		}
		c2, err := store.CreateConfig(ctx, &model.Config{
			Name:         "c2",
			URL:          upstream2.URL,
			Priority:     1,
			ChannelType:  "openai",
			ModelEntries: []model.ModelEntry{{Model: "x1"}},
			Enabled:      true,
		})
		if err != nil {
			t.Fatalf("CreateConfig c2 failed: %v", err)
		}
		c3, err := store.CreateConfig(ctx, &model.Config{
			Name:         "c3-no-key",
			URL:          upstream2.URL,
			Priority:     1,
			ChannelType:  "openai",
			ModelEntries: []model.ModelEntry{{Model: "y1"}},
			Enabled:      true,
		})
		if err != nil {
			t.Fatalf("CreateConfig c3 failed: %v", err)
		}

		if err := store.CreateAPIKeysBatch(ctx, []*model.APIKey{
			{ChannelID: c1.ID, KeyIndex: 0, APIKey: "k1", KeyStrategy: model.KeyStrategySequential},
			{ChannelID: c2.ID, KeyIndex: 0, APIKey: "k2", KeyStrategy: model.KeyStrategySequential},
		}); err != nil {
			t.Fatalf("CreateAPIKeysBatch failed: %v", err)
		}

		c, w := newTestContext(t, newJSONRequest(t, http.MethodPost, "/admin/channels/models/refresh-batch", map[string]any{
			"channel_ids": []int64{c1.ID, c2.ID, c3.ID},
			"mode":        "merge",
		}))
		server.HandleBatchRefreshModels(c)

		if w.Code != http.StatusOK {
			t.Fatalf("status=%d, want %d, body=%s", w.Code, http.StatusOK, w.Body.String())
		}

		var resp struct {
			Success bool `json:"success"`
			Data    struct {
				Updated   int `json:"updated"`
				Unchanged int `json:"unchanged"`
				Failed    int `json:"failed"`
			} `json:"data"`
		}
		mustUnmarshalJSON(t, w.Body.Bytes(), &resp)
		if !resp.Success {
			t.Fatalf("expected success=true, body=%s", w.Body.String())
		}
		if resp.Data.Updated != 1 || resp.Data.Unchanged != 1 || resp.Data.Failed != 1 {
			t.Fatalf("unexpected summary: %+v", resp.Data)
		}

		got1, err := store.GetConfig(ctx, c1.ID)
		if err != nil {
			t.Fatalf("GetConfig c1 failed: %v", err)
		}
		got2, err := store.GetConfig(ctx, c2.ID)
		if err != nil {
			t.Fatalf("GetConfig c2 failed: %v", err)
		}
		if len(got1.ModelEntries) != 2 {
			t.Fatalf("c1 model count=%d, want 2", len(got1.ModelEntries))
		}
		if len(got2.ModelEntries) != 1 {
			t.Fatalf("c2 model count=%d, want 1", len(got2.ModelEntries))
		}
	})

	t.Run("replace mode", func(t *testing.T) {
		upstream := newTestHTTPServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if r.URL.Path != "/v1/models" {
				http.NotFound(w, r)
				return
			}
			w.Header().Set("Content-Type", "application/json")
			_, _ = w.Write([]byte(`{"data":[{"id":"new-1"}]}`))
		}))
		t.Cleanup(upstream.Close)

		server, store, cleanup := setupAdminTestServer(t)
		defer cleanup()

		ctx := context.Background()
		cfg, err := store.CreateConfig(ctx, &model.Config{
			Name:        "replace-channel",
			URL:         upstream.URL,
			Priority:    1,
			ChannelType: "openai",
			ModelEntries: []model.ModelEntry{
				{Model: "old-1"},
				{Model: "old-2"},
			},
			Enabled: true,
		})
		if err != nil {
			t.Fatalf("CreateConfig failed: %v", err)
		}
		if err := store.CreateAPIKeysBatch(ctx, []*model.APIKey{
			{ChannelID: cfg.ID, KeyIndex: 0, APIKey: "k", KeyStrategy: model.KeyStrategySequential},
		}); err != nil {
			t.Fatalf("CreateAPIKeysBatch failed: %v", err)
		}

		c, w := newTestContext(t, newJSONRequest(t, http.MethodPost, "/admin/channels/models/refresh-batch", map[string]any{
			"channel_ids": []int64{cfg.ID},
			"mode":        "replace",
		}))
		server.HandleBatchRefreshModels(c)
		if w.Code != http.StatusOK {
			t.Fatalf("status=%d, want %d, body=%s", w.Code, http.StatusOK, w.Body.String())
		}

		got, err := store.GetConfig(ctx, cfg.ID)
		if err != nil {
			t.Fatalf("GetConfig failed: %v", err)
		}
		if len(got.ModelEntries) != 1 || got.ModelEntries[0].Model != "new-1" {
			t.Fatalf("unexpected models after replace: %#v", got.ModelEntries)
		}
	})

	t.Run("invalid mode", func(t *testing.T) {
		server, _, cleanup := setupAdminTestServer(t)
		defer cleanup()

		c, w := newTestContext(t, newJSONRequestBytes(http.MethodPost, "/admin/channels/models/refresh-batch", []byte(`{"channel_ids":[1],"mode":"xxx"}`)))
		server.HandleBatchRefreshModels(c)
		if w.Code != http.StatusBadRequest {
			t.Fatalf("status=%d, want %d", w.Code, http.StatusBadRequest)
		}
	})
}