ModelRouter/internal/core/channels_more_test.go

package core

import (
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"

	"llmsproxy/internal/lua"
)

// channelsMoreTest：新引入的三条真实流程 + 两个边界流的行为对齐测试。
// 平台定义、决策表、端点逐字照抄参考（见 channels.lua 注释），测试照参考行为搬。

// ---- Cline（WorkOS 设备码 + register 两步） ----

// clineHarness 假上游：device code / authenticate / register 三端点。
type clineHarness struct {
	t    *testing.T
	ps   *lua.Plugins
	ts   *httptest.Server
	tok  atomic.Value // 当前 authenticate 响应脚本 []scriptStep
	step int32
}

func newClineHarness(t *testing.T) *clineHarness {
	t.Helper()
	h := &clineHarness{t: t}
	h.tok.Store([]scriptStep{})
	h.ts = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		switch r.URL.Path {
		case "/user_management/authorize/device":
			fmt.Fprint(w, `{"device_code":"CD1","user_code":"WXYZ-1234","verification_uri":"https://example/activate","interval":1,"expires_in":600}`)
		case "/user_management/authenticate":
			i := int(atomic.AddInt32(&h.step, 1)) - 1
			s := h.tok.Load().([]scriptStep)
			if i >= len(s) {
				w.WriteHeader(500)
				fmt.Fprint(w, `{"error":"unexpected"}`)
				return
			}
			if s[i].http != 0 && s[i].http != 200 {
				w.WriteHeader(s[i].http)
			}
			fmt.Fprint(w, s[i].body)
		case "/api/v1/auth/register":
			// 校验请求体是驼峰 accessToken/refreshToken（参考 registerClineTokens）
			b, _ := io.ReadAll(r.Body)
			for _, want := range []string{`"accessToken":"WO-AT-1"`, `"refreshToken":"WO-RT-1"`} {
				if !strings.Contains(string(b), want) {
					t.Errorf("register 请求体缺 %s: %s", want, b)
				}
			}
			fmt.Fprint(w, `{"success":true,"data":{"accessToken":"CLINE-AT-1","refreshToken":"CLINE-RT-1","userInfo":{"clineUserId":"u1","email":"dev@example.com"}}}`)
		default:
			w.WriteHeader(404)
		}
	}))
	t.Cleanup(h.ts.Close)

	c := kernelTestCore(t)
	h.ps = c.Plugins()
	src, _ := readBundledChannels(t)
	if err := h.ps.LoadSource("channels", src); err != nil {
		t.Fatal(err)
	}
	if _, err := h.ps.Call("channels", "set_endpoints", map[string]interface{}{
		"platform": "cline",
		"device":   map[string]interface{}{"code_url": h.ts.URL + "/user_management/authorize/device", "token_url": h.ts.URL + "/user_management/authenticate"},
		"register": map[string]interface{}{"url": h.ts.URL + "/api/v1/auth/register"},
	}); err != nil {
		t.Fatalf("set_endpoints: %v", err)
	}
	return h
}

func (h *clineHarness) begin() map[string]interface{} {
	out, err := h.ps.InvokeHandler("channels", "start_login", &lua.HandlerRequest{
		Method: "POST", Body: `{"platform":"cline"}`,
	})
	if err != nil {
		h.t.Fatal(err)
	}
	return out.JSON.(map[string]interface{})
}
func (h *clineHarness) poll(id string) map[string]interface{} {
	out, _ := h.ps.InvokeHandler("channels", "poll_login", &lua.HandlerRequest{
		Method: "GET", Query: map[string]string{"id": id},
	})
	return out.JSON.(map[string]interface{})
}

// runWorkosTable 是 WorkOS（Cline）分支的行为契约，供变异测试杀 workos 变异体。
// 决策表 5 行，跟 TestClineWorkOSBehavior 同源（这里只返回布尔，不报细节）。
func (h *clineHarness) runWorkosTable() bool {
	st := func(steps ...string) map[string]interface{} {
		h.tok.Store(scriptSteps(steps...))
		atomic.StoreInt32(&h.step, 0)
		j := h.begin()
		sid := j["session"].(map[string]interface{})["id"].(string)
		return h.poll(sid)
	}
	_ = h.ps.SetPluginKernelPermission("channels", []string{"add_source", "status", "list_sources"}, false)
	// 行1: 非200+error=authorization_pending → 等
	if j := st(`400 {"error":"authorization_pending"}`); j["state"] != "pending" {
		return false
	}
	// 行2: 非200+error=slow_down → 等 + interval+1（基线 1）
	h.tok.Store(scriptSteps(`400 {"error":"slow_down"}`))
	atomic.StoreInt32(&h.step, 0)
	j := h.begin()
	sid := j["session"].(map[string]interface{})["id"].(string)
	if j2 := h.poll(sid); j2["state"] != "pending" {
		return false
	} else if iv, _ := j2["interval"].(float64); iv != 2 {
		return false
	}
	// 行3: error=access_denied → failed
	if j := st(`400 {"error":"access_denied","error_description":"user declined"}`); j["state"] != "failed" {
		return false
	}
	// 行4: 成功（access+refresh）→ register 换 CLINE token → completed
	if j := st(`{"access_token":"WO-AT-1","refresh_token":"WO-RT-1"}`); j["state"] != "completed" || j["account"] != "dev@example.com" {
		return false
	}
	v, _, _ := h.ps.GetSecret("channels", "channels:cline")
	if v != "CLINE-AT-1" {
		return false
	}
	// 行5: 非200无error字段 → failed（不是死等）
	if j := st(`500 {"foo":"bar"}`); j["state"] != "failed" {
		return false
	}
	return true
}

// TestClineWorkOSBehavior：WorkOS 判据（error 字段）+ register 两步 + slow_down +1s。
func TestClineWorkOSBehavior(t *testing.T) {
	h := newClineHarness(t)
	_ = h.ps.SetPluginKernelPermission("channels", []string{"add_source", "status", "list_sources"}, false)

	st := func(steps ...string) map[string]interface{} {
		h.tok.Store(scriptSteps(steps...))
		atomic.StoreInt32(&h.step, 0)
		j := h.begin()
		sid := j["session"].(map[string]interface{})["id"].(string)
		return h.poll(sid)
	}

	t.Run("行1 error=authorization_pending_非200也继续等", func(t *testing.T) {
		if j := st(`400 {"error":"authorization_pending"}`); j["state"] != "pending" {
			t.Fatalf("应 pending: %v", j)
		}
	})
	t.Run("行2 error=slow_down_退避1秒", func(t *testing.T) {
		h.tok.Store(scriptSteps(`400 {"error":"slow_down"}`))
		atomic.StoreInt32(&h.step, 0)
		j := h.begin()
		sid := j["session"].(map[string]interface{})["id"].(string)
		j2 := h.poll(sid)
		if j2["state"] != "pending" {
			t.Fatalf("应 pending: %v", j2)
		}
		if iv, _ := j2["interval"].(float64); iv != 2 { // 基线 1 + slow_down_step 1
			t.Fatalf("WorkOS slow_down 应 interval+1s 得 2, 实得 %v", iv)
		}
	})
	t.Run("行3 error=access_denied_失败带拒绝", func(t *testing.T) {
		if j := st(`400 {"error":"access_denied","error_description":"user declined"}`); j["state"] != "failed" {
			t.Fatalf("应 failed: %v", j)
		}
	})
	t.Run("行4 成功_access+refresh_再打register_凭据是CLINE-AT-1", func(t *testing.T) {
		j := st(`{"access_token":"WO-AT-1","refresh_token":"WO-RT-1"}`)
		if j["state"] != "completed" || j["account"] != "dev@example.com" {
			t.Fatalf("应 completed+dev@example.com: %v", j)
		}
		v, _, _ := h.ps.GetSecret("channels", "channels:cline")
		if v != "CLINE-AT-1" {
			t.Fatalf("凭据应是 register 的 CLINE-AT-1, 实得 %q", v)
		}
	})
	t.Run("行5 非200无error字段_失败", func(t *testing.T) {
		if j := st(`500 {"foo":"bar"}`); j["state"] != "failed" {
			t.Fatalf("应 failed: %v", j)
		}
	})
}

// ---- Buddy（CodeBuddy state 轮询） ----

type buddyHarness struct {
	t       *testing.T
	ps      *lua.Plugins
	ts      *httptest.Server
	tokStep int32
	tokOk   int32
	acctOk  int32
}

func newBuddyHarness(t *testing.T) *buddyHarness {
	t.Helper()
	h := &buddyHarness{t: t, tokOk: 1, acctOk: 1}
	h.ts = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		switch {
		case r.Method == "POST" && r.URL.Path == "/v2/plugin/auth/state":
			fmt.Fprint(w, `{"data":{"state":"ST-1","authUrl":"https://www.codebuddy.cn/login/?platform=ide&state=ST-1"}}`)
		case r.Method == "GET" && r.URL.Path == "/v2/plugin/auth/token":
			i := atomic.AddInt32(&h.tokStep, 1)
			if atomic.LoadInt32(&h.tokOk) == 0 {
				w.WriteHeader(200)
				fmt.Fprint(w, `{"code":11217,"message":"token not ready"}`)
				return
			}
			if i == 1 {
				fmt.Fprint(w, `{"code":11217,"message":"token not ready"}`)
			} else {
				fmt.Fprint(w, `{"data":{"accessToken":"BD-AT-1","refreshToken":"BD-RT-1"}}`)
			}
		case r.Method == "GET" && r.URL.Path == "/v2/plugin/login/account":
			if atomic.LoadInt32(&h.acctOk) == 0 {
				w.WriteHeader(500)
				fmt.Fprint(w, `{"code":1}`)
				return
			}
			fmt.Fprint(w, `{"data":{"nickName":"buddy-dev"}}`)
		default:
			w.WriteHeader(404)
		}
	}))
	t.Cleanup(h.ts.Close)

	c := kernelTestCore(t)
	h.ps = c.Plugins()
	src, _ := readBundledChannels(t)
	if err := h.ps.LoadSource("channels", src); err != nil {
		t.Fatal(err)
	}
	if _, err := h.ps.Call("channels", "set_endpoints", map[string]interface{}{
		"platform": "buddy",
		"state":    map[string]interface{}{"state_url": h.ts.URL + "/v2/plugin/auth/state", "token_url": h.ts.URL + "/v2/plugin/auth/token", "account_url": h.ts.URL + "/v2/plugin/login/account"},
	}); err != nil {
		t.Fatalf("set_endpoints: %v", err)
	}
	return h
}

func (h *buddyHarness) begin() (map[string]interface{}, string) {
	out, err := h.ps.InvokeHandler("channels", "start_login", &lua.HandlerRequest{
		Method: "POST", Body: `{"platform":"buddy"}`,
	})
	if err != nil {
		h.t.Fatal(err)
	}
	j := out.JSON.(map[string]interface{})
	return j, j["session"].(map[string]interface{})["id"].(string)
}
func (h *buddyHarness) poll(id string) map[string]interface{} {
	out, _ := h.ps.InvokeHandler("channels", "poll_login", &lua.HandlerRequest{
		Method: "GET", Query: map[string]string{"id": id},
	})
	return out.JSON.(map[string]interface{})
}

// TestBuddyStatePollBehavior：state-poll 两拍（先 11217 后 token）+ 账号查询 + account 失败不阻塞。
func TestBuddyStatePollBehavior(t *testing.T) {
	h := newBuddyHarness(t)
	_ = h.ps.SetPluginKernelPermission("channels", []string{"add_source", "status", "list_sources"}, false)

	t.Run("行1 state申请给出authUrl", func(t *testing.T) {
		j, _ := h.begin()
		if j["url"] == "" {
			t.Fatalf("应有 authUrl: %v", j)
		}
	})
	t.Run("行2 首拍11217仍pending_次拍token_completed", func(t *testing.T) {
		atomic.StoreInt32(&h.tokStep, 0)
		_, sid := h.begin()
		if j := h.poll(sid); j["state"] != "pending" {
			t.Fatalf("首拍应 pending: %v", j)
		}
		if j := h.poll(sid); j["state"] != "completed" || j["account"] != "buddy-dev" {
			t.Fatalf("次拍应 completed+buddy-dev: %v", j)
		}
		v, _, _ := h.ps.GetSecret("channels", "channels:buddy")
		if v != "BD-AT-1" {
			t.Fatalf("凭据应 BD-AT-1, 实得 %q", v)
		}
	})
	t.Run("行3 account查询失败不阻塞_回退默认名", func(t *testing.T) {
		atomic.StoreInt32(&h.tokStep, 0)
		atomic.StoreInt32(&h.acctOk, 0)
		defer atomic.StoreInt32(&h.acctOk, 1)
		_, sid := h.begin()
		h.poll(sid)
		j := h.poll(sid)
		if j["state"] != "completed" {
			t.Fatalf("account 失败不应阻塞: %v", j)
		}
	})
}

// ---- 通用 ----

func scriptSteps(steps ...string) []scriptStep {
	var out []scriptStep
	for _, b := range steps {
		code := 200
		switch {
		case strings.HasPrefix(b, "400 "):
			code, b = 400, b[4:]
		case strings.HasPrefix(b, "403 "):
			code, b = 403, b[4:]
		case strings.HasPrefix(b, "500 "):
			code, b = 500, b[4:]
		}
		out = append(out, scriptStep{http: code, body: b})
	}
	return out
}