Add projects

backend/internal/api/handlers/management/config_weight_test.go

@@ -0,0 +1,104 @@
+package management
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/gin-gonic/gin"
+	"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
+)
+
+func TestPatchAPIKeyWeightForEveryFamily(t *testing.T) {
+	tests := []struct {
+		name  string
+		setup func(*config.Config)
+		patch func(*Handler, *gin.Context)
+		get   func(*config.Config) *int
+	}{
+		{name: "gemini", setup: func(cfg *config.Config) { cfg.GeminiKey = []config.GeminiKey{{APIKey: "key"}} }, patch: (*Handler).PatchGeminiKey, get: func(cfg *config.Config) *int { return cfg.GeminiKey[0].Weight }},
+		{name: "interactions", setup: func(cfg *config.Config) { cfg.InteractionsKey = []config.GeminiKey{{APIKey: "key"}} }, patch: (*Handler).PatchInteractionsKey, get: func(cfg *config.Config) *int { return cfg.InteractionsKey[0].Weight }},
+		{name: "claude", setup: func(cfg *config.Config) { cfg.ClaudeKey = []config.ClaudeKey{{APIKey: "key"}} }, patch: (*Handler).PatchClaudeKey, get: func(cfg *config.Config) *int { return cfg.ClaudeKey[0].Weight }},
+		{name: "vertex", setup: func(cfg *config.Config) {
+			cfg.VertexCompatAPIKey = []config.VertexCompatKey{{APIKey: "key", BaseURL: "https://example.com"}}
+		}, patch: (*Handler).PatchVertexCompatKey, get: func(cfg *config.Config) *int { return cfg.VertexCompatAPIKey[0].Weight }},
+		{name: "codex", setup: func(cfg *config.Config) {
+			cfg.CodexKey = []config.CodexKey{{APIKey: "key", BaseURL: "https://example.com"}}
+		}, patch: (*Handler).PatchCodexKey, get: func(cfg *config.Config) *int { return cfg.CodexKey[0].Weight }},
+		{name: "xai", setup: func(cfg *config.Config) {
+			cfg.XAIKey = []config.XAIKey{{APIKey: "key", BaseURL: "https://example.com"}}
+		}, patch: (*Handler).PatchXAIKey, get: func(cfg *config.Config) *int { return cfg.XAIKey[0].Weight }},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			cfg := &config.Config{}
+			test.setup(cfg)
+			h := &Handler{cfg: cfg, configFilePath: writeTestConfigFile(t)}
+
+			rec := httptest.NewRecorder()
+			ctx, _ := gin.CreateTestContext(rec)
+			ctx.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/key", strings.NewReader(`{"index":0,"value":{"weight":7}}`))
+			ctx.Request.Header.Set("Content-Type", "application/json")
+			test.patch(h, ctx)
+
+			if rec.Code != http.StatusOK {
+				t.Fatalf("status = %d, want 200; body=%s", rec.Code, rec.Body.String())
+			}
+			if weight := test.get(cfg); weight == nil || *weight != 7 {
+				t.Fatalf("weight = %v, want 7", weight)
+			}
+		})
+	}
+}
+
+func TestPatchAPIKeyWeightResetAndStrictValidation(t *testing.T) {
+	initial := 5
+	cfg := &config.Config{GeminiKey: []config.GeminiKey{{APIKey: "key", Weight: &initial}}}
+	h := &Handler{cfg: cfg, configFilePath: writeTestConfigFile(t)}
+
+	patch := func(raw string) *httptest.ResponseRecorder {
+		t.Helper()
+		rec := httptest.NewRecorder()
+		ctx, _ := gin.CreateTestContext(rec)
+		body := fmt.Sprintf(`{"index":0,"value":{"weight":%s}}`, raw)
+		ctx.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/gemini-api-key", strings.NewReader(body))
+		ctx.Request.Header.Set("Content-Type", "application/json")
+		h.PatchGeminiKey(ctx)
+		return rec
+	}
+
+	for _, invalid := range []string{"1.5", "1000001", "9223372036854775808", `"7"`} {
+		rec := patch(invalid)
+		if rec.Code != http.StatusBadRequest {
+			t.Fatalf("weight %s status = %d, want 400; body=%s", invalid, rec.Code, rec.Body.String())
+		}
+		if cfg.GeminiKey[0].Weight == nil || *cfg.GeminiKey[0].Weight != initial {
+			t.Fatalf("invalid weight %s changed config", invalid)
+		}
+	}
+
+	if rec := patch("null"); rec.Code != http.StatusOK {
+		t.Fatalf("reset status = %d, want 200; body=%s", rec.Code, rec.Body.String())
+	}
+	if cfg.GeminiKey[0].Weight != nil {
+		t.Fatalf("reset weight = %v, want nil default", cfg.GeminiKey[0].Weight)
+	}
+}
+
+func TestPutAPIKeyWeightRejectsAboveMaximum(t *testing.T) {
+	h := &Handler{cfg: &config.Config{}, configFilePath: writeTestConfigFile(t)}
+	rec := httptest.NewRecorder()
+	ctx, _ := gin.CreateTestContext(rec)
+	ctx.Request = httptest.NewRequest(http.MethodPut, "/v0/management/gemini-api-key", strings.NewReader(`[{"api-key":"key","weight":1000001}]`))
+	ctx.Request.Header.Set("Content-Type", "application/json")
+	h.PutGeminiKeys(ctx)
+	if rec.Code != http.StatusBadRequest {
+		t.Fatalf("status = %d, want 400; body=%s", rec.Code, rec.Body.String())
+	}
+	if len(h.cfg.GeminiKey) != 0 {
+		t.Fatal("invalid PUT changed config")
+	}
+}