mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-11-07 09:47:23 +00:00
6fe6014d92
产品代码统一为纯文本风格（照参考实现），去掉 emoji：billing/quota/usage 插件、 channels 的提示文案、WebUI 文案，以及 go 测试、docs、packaging 脚本中的 注意/OK/FAIL 等符号化写法。第三方目录（vendor、node_modules、.probe）不动。 另：packaging/upload-release.sh 的仓库名改为真实全名 JianFeeeee/ModelRouter → HeNanMU-EEA/ModelRouter（前者是重定向，POST/PUT 不跟随 会 404）。
265 lines
8.0 KiB
Go
265 lines
8.0 KiB
Go
package core
|
||
|
||
import (
|
||
"fmt"
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"os"
|
||
"path/filepath"
|
||
"strings"
|
||
"sync"
|
||
"sync/atomic"
|
||
"testing"
|
||
|
||
"llmsproxy/internal/config"
|
||
"llmsproxy/internal/lua"
|
||
)
|
||
|
||
// mustBox 构造一个真 SecretBox（与运行期同一实现，密钥落 runtime dir）。
|
||
func mustBox(t *testing.T, runtimeFile string) *config.SecretBox {
|
||
t.Helper()
|
||
box, err := config.NewSecretBox(runtimeFile)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return box
|
||
}
|
||
|
||
// channelsCore 起一个带 channels 插件、可指定源码与 box 的核心（压力测试用）。
|
||
func channelsCore(t *testing.T, src string) (*Core, *lua.Plugins) {
|
||
t.Helper()
|
||
c := kernelTestCore(t)
|
||
ps := c.Plugins()
|
||
if err := ps.LoadSource("channels", src); err != nil {
|
||
t.Fatalf("加载 channels: %v", err)
|
||
}
|
||
ps.SetSecretBox(mustBox(t, filepath.Join(t.TempDir(), "runtime.json")))
|
||
_ = ps.SetPluginKernelPermission("channels", []string{"add_source", "status", "list_sources"}, false)
|
||
return c, ps
|
||
}
|
||
|
||
// cookieishPlatformFixture 是压力测试用的占位平台（对应此前内联在生产源码里的
|
||
// "cookieish"）。夹具经运行时注入，生产目录只含真实厂商。
|
||
const cookieishPlatformFixture = `[{
|
||
"id": "cookieish",
|
||
"name": "Cookie 平台（测试夹具）",
|
||
"flow": "paste",
|
||
"base_url": "https://api.cookieish.example/v1",
|
||
"adapter": "openai",
|
||
"models": ["c-model-1"],
|
||
"checkin": {"url": "https://api.cookieish.example/v1/api/credit/claim", "method": "POST", "body": "{}"}
|
||
}]`
|
||
|
||
// TestChannelsConcurrentLogins 压力：并发走 device-poll 登录。
|
||
//
|
||
// 插件是单 VM 单锁 —— 本测试验证的是**正确性下的并发**（会话互不串台、
|
||
// 凭据互不可见、最终都 completed），真吞吐压力在 e2e 层（真二进制+HTTP 并发）。
|
||
func TestChannelsConcurrentLogins(t *testing.T) {
|
||
var hits int32
|
||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
switch r.URL.Path {
|
||
case "/oauth2/device/code":
|
||
n := atomic.AddInt32(&hits, 1)
|
||
w.Header().Set("Content-Type", "application/json")
|
||
fmt.Fprintf(w, `{"device_code":"DEV-%d","user_code":"UC-%d","verification_uri":"https://example/activate","interval":1,"expires_in":600}`, n, n)
|
||
case "/oauth2/token":
|
||
w.Header().Set("Content-Type", "application/json")
|
||
fmt.Fprint(w, `{"access_token":"TOK","refresh_token":"RF","nickname":"u"}`)
|
||
default:
|
||
w.WriteHeader(404)
|
||
}
|
||
}))
|
||
defer ts.Close()
|
||
|
||
_, ps := channelsCore(t, mustChannelsSrc(t))
|
||
// 把 minimax 端点指到假上游
|
||
if _, err := ps.Call("channels", "set_endpoints", map[string]interface{}{
|
||
"platform": "minimax",
|
||
"device": map[string]interface{}{
|
||
"code_url": ts.URL + "/oauth2/device/code",
|
||
"token_url": ts.URL + "/oauth2/token",
|
||
},
|
||
}); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
|
||
const rounds = 12
|
||
var wg sync.WaitGroup
|
||
var failures int32
|
||
for i := 0; i < rounds; i++ {
|
||
wg.Add(1)
|
||
go func(n int) {
|
||
defer wg.Done()
|
||
out, err := ps.InvokeHandler("channels", "start_login", &lua.HandlerRequest{
|
||
Method: "POST", Body: `{"platform":"minimax"}`,
|
||
})
|
||
if err != nil {
|
||
t.Errorf("start_login #%d: %v", n, err)
|
||
atomic.AddInt32(&failures, 1)
|
||
return
|
||
}
|
||
j := out.JSON.(map[string]interface{})
|
||
sess, _ := j["session"].(map[string]interface{})
|
||
if sess == nil {
|
||
t.Errorf("start_login #%d 无会话: %v", n, j)
|
||
atomic.AddInt32(&failures, 1)
|
||
return
|
||
}
|
||
sid := sess["id"].(string)
|
||
for k := 0; k < 10; k++ {
|
||
out2, _ := ps.InvokeHandler("channels", "poll_login", &lua.HandlerRequest{
|
||
Method: "GET", Query: map[string]string{"id": sid},
|
||
})
|
||
m2 := out2.JSON.(map[string]interface{})
|
||
if m2["state"] == "completed" {
|
||
return
|
||
}
|
||
if m2["state"] == "failed" {
|
||
t.Errorf("会话 #%d failed: %v", n, m2)
|
||
atomic.AddInt32(&failures, 1)
|
||
return
|
||
}
|
||
}
|
||
t.Errorf("会话 #%d 10 次轮询未终态", n)
|
||
atomic.AddInt32(&failures, 1)
|
||
}(i)
|
||
}
|
||
wg.Wait()
|
||
if failures != 0 {
|
||
t.Fatalf("并发登录失败 %d 轮", failures)
|
||
}
|
||
}
|
||
|
||
// TestChannelsClaimStress 压力：同一平台并发领签到（幂等 + 无登录态丢失）。
|
||
func TestChannelsClaimStress(t *testing.T) {
|
||
var claims int32
|
||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
if r.URL.Path == "/v1/checkin" {
|
||
atomic.AddInt32(&claims, 1)
|
||
w.Header().Set("Content-Type", "application/json")
|
||
fmt.Fprint(w, `{"ok":true,"credits":10}`)
|
||
return
|
||
}
|
||
w.WriteHeader(404)
|
||
fmt.Fprint(w, `{}`)
|
||
}))
|
||
defer ts.Close()
|
||
|
||
_, ps := channelsCore(t, mustChannelsSrc(t))
|
||
injectTestPlatforms(t, ps, cookieishPlatformFixture)
|
||
if _, err := ps.Call("channels", "set_endpoints", map[string]interface{}{
|
||
"platform": "cookieish",
|
||
"checkin": map[string]interface{}{"url": ts.URL + "/v1/checkin"},
|
||
}); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if err := ps.SetSecret("channels", "channels:cookieish", "PASTE-TOKEN"); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
|
||
const n = 30
|
||
var wg sync.WaitGroup
|
||
for i := 0; i < n; i++ {
|
||
wg.Add(1)
|
||
go func(k int) {
|
||
defer wg.Done()
|
||
out, err := ps.InvokeHandler("channels", "claim_daily", &lua.HandlerRequest{
|
||
Method: "POST", Body: `{"platform":"cookieish"}`,
|
||
})
|
||
if err != nil {
|
||
t.Errorf("claim #%d: %v", k, err)
|
||
return
|
||
}
|
||
j := out.JSON.(map[string]interface{})
|
||
if j["ok"] != true {
|
||
t.Errorf("claim #%d 未成功: %v", k, j)
|
||
}
|
||
}(i)
|
||
}
|
||
wg.Wait()
|
||
if got := atomic.LoadInt32(&claims); got != n {
|
||
t.Errorf("上游应收到 %d 次领取, 实得 %d", n, got)
|
||
}
|
||
}
|
||
|
||
// TestLegacyVaultPlaintextLoad 旧通路：老版本（无加密）写的明文库在新版本
|
||
// （带 box）下必须能读 —— 凭据升级不丢。落盘则会在下次写时整体重封为密文。
|
||
func TestLegacyVaultPlaintextLoad(t *testing.T) {
|
||
dir := t.TempDir()
|
||
ps := lua.NewPlugins(nil, dir)
|
||
if err := ps.SetSecret("legacy", "channels:old", "PLAIN-OLD-TOKEN"); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
ps.SetSecretBox(mustBox(t, filepath.Join(dir, "runtime.json")))
|
||
v, ok, err := ps.GetSecret("legacy", "channels:old")
|
||
if err != nil || !ok || v != "PLAIN-OLD-TOKEN" {
|
||
t.Fatalf("旧明文库读回失败: v=%q ok=%v err=%v", v, ok, err)
|
||
}
|
||
// 任意一次写触发整体重封
|
||
if err := ps.SetSecret("legacy", "k2", "NEW"); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
raw, err := os.ReadFile(filepath.Join(dir, ".secrets", "legacy.json"))
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if strings.Contains(string(raw), "PLAIN-OLD-TOKEN") {
|
||
t.Fatalf("重封后旧明文仍在盘上:\n%s", raw)
|
||
}
|
||
if !strings.Contains(string(raw), "enc:v1:") {
|
||
t.Fatalf("重封后应含密文:\n%s", raw)
|
||
}
|
||
}
|
||
|
||
// TestLegacySourceSealedConfig 旧通路：config.yaml 里明文 api_key 的旧配置，
|
||
// 新版本加载后运行期拿到解密值、落盘被 seal（SealIfNeeded 语义）。
|
||
func TestLegacySourceSealedConfig(t *testing.T) {
|
||
dir := t.TempDir()
|
||
cfgPath := filepath.Join(dir, "config.yaml")
|
||
yaml := "listen: 127.0.0.1:0\n" +
|
||
"gateway_keys:\n - sk-legacy-0001\n" +
|
||
"runtime_file: " + filepath.Join(dir, "runtime.json") + "\n" +
|
||
"adapter_dir: " + filepath.Join(dir, "adapters") + "\n" +
|
||
"sources:\n" +
|
||
" - name: legacy-src\n" +
|
||
" base_url: https://legacy.example.com/v1\n" +
|
||
" adapter: openai\n" +
|
||
" api_key: PLAIN-LEGACY-KEY\n" +
|
||
" max_concurrent: 2\n" +
|
||
" models:\n" +
|
||
" - id: m1\n" +
|
||
" kind: chat\n" +
|
||
" priority: 10\n"
|
||
if err := os.WriteFile(cfgPath, []byte(yaml), 0o644); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
cfg, err := config.Load(cfgPath)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
c, err := NewFromConfig(cfg)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer c.Close()
|
||
|
||
got := c.Sources()
|
||
if len(got) != 1 || got[0].APIKey != "PLAIN-LEGACY-KEY" {
|
||
t.Fatalf("运行期 api_key 应解密: %#v", got)
|
||
}
|
||
raw, err := os.ReadFile(cfgPath)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if strings.Contains(string(raw), "PLAIN-LEGACY-KEY") {
|
||
t.Fatalf("旧明文 api_key 升级后仍明文落盘:\n%s", raw)
|
||
}
|
||
if !strings.Contains(string(raw), "enc:v1:") {
|
||
t.Fatalf("落盘应含 enc:v1 密文:\n%s", raw)
|
||
}
|
||
mk, err := os.ReadFile(filepath.Join(dir, "master.key"))
|
||
if err != nil || len(mk) != 64 {
|
||
t.Fatalf("master.key 应 64 hex, len=%d err=%v", len(mk), err)
|
||
}
|
||
}
|