internal/app/admin_cooldown_test.go

package app

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

	"ccLoad/internal/model"

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

// TestHandleSetChannelCooldown 测试设置渠道冷却
func TestHandleSetChannelCooldown(t *testing.T) {
	tests := []struct {
		name           string
		channelID      string
		requestBody    map[string]any
		setupChannel   bool
		expectedStatus int
		expectSuccess  bool
	}{
		{
			name:      "成功设置渠道冷却",
			channelID: "1",
			requestBody: map[string]any{
				"duration_ms": 60000, // 60秒
			},
			setupChannel:   true,
			expectedStatus: http.StatusOK,
			expectSuccess:  true,
		},
		{
			name:      "无效的渠道ID",
			channelID: "invalid",
			requestBody: map[string]any{
				"duration_ms": 60000,
			},
			setupChannel:   false,
			expectedStatus: http.StatusBadRequest,
			expectSuccess:  false,
		},
		{
			name:      "无效的请求体",
			channelID: "1",
			requestBody: map[string]any{
				"invalid_field": "value",
			},
			setupChannel:   true,
			expectedStatus: http.StatusBadRequest,
			expectSuccess:  false,
		},
		{
			name:      "大冷却时长",
			channelID: "1",
			requestBody: map[string]any{
				"duration_ms": 1800000, // 30分钟
			},
			setupChannel:   true,
			expectedStatus: http.StatusOK,
			expectSuccess:  true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// 创建测试服务器
			srv := newInMemoryServer(t)

			// 设置渠道(如果需要)
			if tt.setupChannel {
				cfg := &model.Config{
					ID:           1,
					Name:         "test-channel",
					URL:          "http://test.example.com",
					Priority:     1,
					ModelEntries: []model.ModelEntry{{Model: "test-model", RedirectModel: ""}},
					Enabled:      true,
				}
				_, err := srv.store.CreateConfig(context.Background(), cfg)
				if err != nil {
					t.Fatalf("创建测试渠道失败: %v", err)
				}
			}

			c, w := newTestContext(t, newJSONRequest(t, http.MethodPost, "/admin/channels/"+tt.channelID+"/cooldown", tt.requestBody))
			c.Params = gin.Params{{Key: "id", Value: tt.channelID}}

			// 调用处理函数
			srv.HandleSetChannelCooldown(c)

			// 验证响应状态码
			if w.Code != tt.expectedStatus {
				t.Errorf("期望状态码 %d, 实际 %d", tt.expectedStatus, w.Code)
			}

			resp := mustParseAPIResponse[json.RawMessage](t, w.Body.Bytes())
			if resp.Success != tt.expectSuccess {
				t.Errorf("期望 success=%v, 实际=%v, error=%q", tt.expectSuccess, resp.Success, resp.Error)
			}
		})
	}
}

// TestHandleSetKeyCooldown 测试设置Key冷却
func TestHandleSetKeyCooldown(t *testing.T) {
	tests := []struct {
		name           string
		channelID      string
		keyIndex       string
		requestBody    map[string]any
		setupData      bool
		expectedStatus int
		expectSuccess  bool
	}{
		{
			name:      "成功设置Key冷却",
			channelID: "1",
			keyIndex:  "0",
			requestBody: map[string]any{
				"duration_ms": 30000, // 30秒
			},
			setupData:      true,
			expectedStatus: http.StatusOK,
			expectSuccess:  true,
		},
		{
			name:      "无效的渠道ID",
			channelID: "invalid",
			keyIndex:  "0",
			requestBody: map[string]any{
				"duration_ms": 30000,
			},
			setupData:      false,
			expectedStatus: http.StatusBadRequest,
			expectSuccess:  false,
		},
		{
			name:      "无效的Key索引",
			channelID: "1",
			keyIndex:  "invalid",
			requestBody: map[string]any{
				"duration_ms": 30000,
			},
			setupData:      false,
			expectedStatus: http.StatusBadRequest,
			expectSuccess:  false,
		},
		{
			name:      "负数Key索引",
			channelID: "1",
			keyIndex:  "-1",
			requestBody: map[string]any{
				"duration_ms": 30000,
			},
			setupData:      false,
			expectedStatus: http.StatusBadRequest,
			expectSuccess:  false,
		},
		{
			name:      "无效的请求体",
			channelID: "1",
			keyIndex:  "0",
			requestBody: map[string]any{
				"invalid_field": "value",
			},
			setupData:      true,
			expectedStatus: http.StatusBadRequest,
			expectSuccess:  false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// 创建测试服务器
			srv := newInMemoryServer(t)

			// 设置测试数据(如果需要)
			if tt.setupData {
				ctx := context.Background()

				// 创建渠道
				cfg := &model.Config{
					ID:           1,
					Name:         "test-channel",
					URL:          "http://test.example.com",
					Priority:     1,
					ModelEntries: []model.ModelEntry{{Model: "test-model", RedirectModel: ""}},
					Enabled:      true,
				}
				_, err := srv.store.CreateConfig(ctx, cfg)
				if err != nil {
					t.Fatalf("创建测试渠道失败: %v", err)
				}

				// 创建API Key
				key := &model.APIKey{
					ChannelID:   1,
					KeyIndex:    0,
					APIKey:      "test-key",
					KeyStrategy: model.KeyStrategySequential,
				}
				if err := srv.store.CreateAPIKeysBatch(ctx, []*model.APIKey{key}); err != nil {
					t.Fatalf("创建API Key失败: %v", err)
				}
			}

			c, w := newTestContext(t, newJSONRequest(t, http.MethodPost, "/admin/channels/"+tt.channelID+"/keys/"+tt.keyIndex+"/cooldown", tt.requestBody))
			c.Params = gin.Params{
				{Key: "id", Value: tt.channelID},
				{Key: "keyIndex", Value: tt.keyIndex},
			}

			// 调用处理函数
			srv.HandleSetKeyCooldown(c)

			// 验证响应状态码
			if w.Code != tt.expectedStatus {
				t.Errorf("期望状态码 %d, 实际 %d", tt.expectedStatus, w.Code)
			}

			resp := mustParseAPIResponse[json.RawMessage](t, w.Body.Bytes())
			if resp.Success != tt.expectSuccess {
				t.Errorf("期望 success=%v, 实际=%v, error=%q", tt.expectSuccess, resp.Success, resp.Error)
			}
		})
	}
}

// TestSetChannelCooldown_Integration 测试渠道冷却集成
func TestSetChannelCooldown_Integration(t *testing.T) {
	srv := newInMemoryServer(t)

	ctx := context.Background()

	// 创建测试渠道
	cfg := &model.Config{
		ID:           1,
		Name:         "test-channel",
		URL:          "http://test.example.com",
		Priority:     1,
		ModelEntries: []model.ModelEntry{{Model: "test-model", RedirectModel: ""}},
		Enabled:      true,
	}
	_, err := srv.store.CreateConfig(ctx, cfg)
	if err != nil {
		t.Fatalf("创建测试渠道失败: %v", err)
	}

	// 设置冷却
	requestBody := map[string]any{
		"duration_ms": 120000, // 2分钟
	}
	c, w := newTestContext(t, newJSONRequest(t, http.MethodPost, "/admin/channels/1/cooldown", requestBody))
	c.Params = gin.Params{{Key: "id", Value: "1"}}

	srv.HandleSetChannelCooldown(c)

	// 验证响应
	if w.Code != http.StatusOK {
		t.Errorf("期望状态码 200, 实际 %d", w.Code)
	}

	resp := mustParseAPIResponse[json.RawMessage](t, w.Body.Bytes())
	if !resp.Success {
		t.Error("期望 success=true")
	}

	// 验证数据库中的冷却状态
	updatedCfg, err := srv.store.GetConfig(ctx, 1)
	if err != nil {
		t.Fatalf("获取渠道失败: %v", err)
	}

	if updatedCfg.CooldownUntil == 0 {
		t.Error("期望渠道被冷却, 但 CooldownUntil=0")
	}
}

// TestSetKeyCooldown_Integration 测试Key冷却集成
func TestSetKeyCooldown_Integration(t *testing.T) {
	srv := newInMemoryServer(t)

	ctx := context.Background()

	// 创建测试渠道
	cfg := &model.Config{
		ID:           1,
		Name:         "test-channel",
		URL:          "http://test.example.com",
		Priority:     1,
		ModelEntries: []model.ModelEntry{{Model: "test-model", RedirectModel: ""}},
		Enabled:      true,
	}
	_, err := srv.store.CreateConfig(ctx, cfg)
	if err != nil {
		t.Fatalf("创建测试渠道失败: %v", err)
	}

	// 创建API Key
	key := &model.APIKey{
		ChannelID:   1,
		KeyIndex:    0,
		APIKey:      "test-key",
		KeyStrategy: model.KeyStrategySequential,
	}
	if err := srv.store.CreateAPIKeysBatch(ctx, []*model.APIKey{key}); err != nil {
		t.Fatalf("创建API Key失败: %v", err)
	}

	// 设置Key冷却
	requestBody := map[string]any{
		"duration_ms": 90000, // 90秒
	}
	c, w := newTestContext(t, newJSONRequest(t, http.MethodPost, "/admin/channels/1/keys/0/cooldown", requestBody))
	c.Params = gin.Params{
		{Key: "id", Value: "1"},
		{Key: "keyIndex", Value: "0"},
	}

	srv.HandleSetKeyCooldown(c)

	// 验证响应
	if w.Code != http.StatusOK {
		t.Errorf("期望状态码 200, 实际 %d", w.Code)
	}

	resp := mustParseAPIResponse[json.RawMessage](t, w.Body.Bytes())
	if !resp.Success {
		t.Error("期望 success=true")
	}

	// 验证数据库中的Key冷却状态
	updatedKey, err := srv.store.GetAPIKey(ctx, 1, 0)
	if err != nil {
		t.Fatalf("获取API Key失败: %v", err)
	}

	if updatedKey.CooldownUntil == 0 {
		t.Error("期望Key被冷却, 但 CooldownUntil=0")
	}
}