diff --git a/relay/helper/price_domain_test.go b/relay/helper/price_domain_test.go new file mode 100644 index 00000000..0252d3dd --- /dev/null +++ b/relay/helper/price_domain_test.go @@ -0,0 +1,68 @@ +package helper + +import ( + "net/http/httptest" + "testing" + + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/setting/ratio_setting" + + "github.com/gin-gonic/gin" +) + +// 验证「域名 × 命中分组 → 系数」在 HandleGroupRatio 里正确生效。 +// 模拟方式 B:通过 X-Forwarded-Host 头表示请求来自不同对外域名。 +func TestDomainRatioInHandleGroupRatio(t *testing.T) { + gin.SetMode(gin.TestMode) + + // 配置:api1 全部分组 ×1.2;api2 默认 ×1.5,但命中 vip 分组时 ×2.0 + cfg := `{"api1.test":{"*":1.2},"api2.test":{"*":1.5,"vip":2.0}}` + if err := ratio_setting.UpdateDomainRatioByJSONString(cfg); err != nil { + t.Fatalf("加载 DomainRatio 配置失败: %v", err) + } + if !ratio_setting.HasDomainRatioConfigured() { + t.Fatal("HasDomainRatioConfigured 应为 true") + } + + newCtx := func(host string) *gin.Context { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("POST", "/v1/chat/completions", nil) + if host != "" { + c.Request.Header.Set("X-Forwarded-Host", host) + } + return c + } + + // 基线:未命中任何域名配置 → 系数 1(GetGroupRatio("vip") 缺省也为 1) + base := HandleGroupRatio(newCtx("unknown.test"), &relaycommon.RelayInfo{UsingGroup: "vip"}).GroupRatio + if base != 1 { + t.Fatalf("基线倍率应为 1,实际 %v", base) + } + + cases := []struct { + name string + host string + group string + want float64 + }{ + {"api1 任意分组用默认*", "api1.test", "vip", 1.2}, + {"api1 端口应被忽略", "api1.test:3000", "default", 1.2}, + {"api2 命中vip用专属系数", "api2.test", "vip", 2.0}, + {"api2 其它分组回退*", "api2.test", "default", 1.5}, + {"未配置域名不影响", "other.test", "vip", 1.0}, + {"无Host头不影响", "", "vip", 1.0}, + } + + for _, tc := range cases { + got := HandleGroupRatio(newCtx(tc.host), &relaycommon.RelayInfo{UsingGroup: tc.group}).GroupRatio + if got != tc.want { + t.Errorf("[%s] host=%q group=%q 期望倍率 %v,实际 %v", tc.name, tc.host, tc.group, tc.want, got) + } else { + t.Logf("[OK] %s → %v", tc.name, got) + } + } + + // 清理,避免影响其它测试 + _ = ratio_setting.UpdateDomainRatioByJSONString(`{}`) +}