hunjing/shenshi/internal/server/byok_quota_test.go

// BYOK / 配额三段式 / 用量归属的服务级场景测试。
package server

import (
	"bufio"
	"encoding/hex"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"hunjing/shenshi/internal/byok"
	"hunjing/shenshi/internal/quota"
	"hunjing/shenshi/internal/route"
	"hunjing/shenshi/internal/store"
)

func testMasterKey() string {
	b := make([]byte, 32)
	for i := range b {
		b[i] = byte(i + 1)
	}
	return hex.EncodeToString(b)
}

// fullServer 装配 store/byok/quota 的完整服务；eps 为注册表候选（按优先级升序尝试）。
func fullServer(t *testing.T, budget int64, eps ...route.Endpoint) (*httptest.Server, store.Store, *byok.Vault) {
	t.Helper()
	st := store.NewMem()
	vault, err := byok.NewVault(testMasterKey(), st)
	if err != nil {
		t.Fatal(err)
	}
	registry := registryWith(t, time.Second, eps...)
	srv := New(Config{Store: st, Vault: vault, Quota: quota.New(st, budget)}, registry)
	ts := httptest.NewServer(srv)
	t.Cleanup(ts.Close)
	return ts, st, vault
}

func postChat(ts *httptest.Server, body, userID string) (*http.Response, string) {
	req, _ := http.NewRequest(http.MethodPost, ts.URL+"/v1/chat/completions", strings.NewReader(body))
	req.Header.Set("content-type", "application/json")
	if userID != "" {
		req.Header.Set("X-Hunjing-User-Id", userID)
	}
	resp, err := ts.Client().Do(req)
	if err != nil {
		return nil, ""
	}
	var sb strings.Builder
	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		sb.WriteString(scanner.Text())
		sb.WriteString("\n")
	}
	return resp, sb.String()
}

func byokChatBody(run, extra string) string {
	return `{"ir_version":1,"slot":"chat","messages":[{"role":"user","content":"hi"}],"run_id":"` + run + `"` + extra + `,"stream":true}`
}

// 场景：BYOK 独占路由——请求携带 byok_key_id 时只走用户端点，不回落平台候选。
func TestByokExclusiveRouting(t *testing.T) {
	platform := fakeOpenAI(t, 0, 0, "平台端点")
	defer platform.Close()
	ts, _, vault := fullServer(t, 0, oaEndpoint("platform", platform.URL, 10))

	byokUpstream := fakeOpenAI(t, 0, 0, "用户端点")
	defer byokUpstream.Close()
	key, err := vault.Register("u1", "自带", "openai_compatible", byokUpstream.URL, "sk-mine", "my-model")
	if err != nil {
		t.Fatal(err)
	}

	body := byokChatBody("run_byok", `,"user_id":"u1","byok_key_id":"`+key.ID+`"`)
	resp, out := postChat(ts, body, "u1")
	if resp == nil || resp.StatusCode != 200 {
		t.Fatalf("BYOK 请求失败: %s", out)
	}
	if !strings.Contains(out, "用户端点") || strings.Contains(out, "平台端点") {
		t.Fatalf("应独占 BYOK 端点: %s", out)
	}
}

// 场景：BYOK 无效引用 → 整体失败，不回落平台端点。
func TestByokInvalidReferenceRejected(t *testing.T) {
	platform := fakeOpenAI(t, 0, 0, "平台端点")
	defer platform.Close()
	ts, _, _ := fullServer(t, 0, oaEndpoint("platform", platform.URL, 10))

	resp, out := postChat(ts, byokChatBody("run_bad", `,"user_id":"u1","byok_key_id":"byok_missing"`), "u1")
	if resp == nil || resp.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("应 503: %s", out)
	}
	if !strings.Contains(out, "HJ-2001") || strings.Contains(out, "平台端点") {
		t.Fatalf("应整体失败而非回落: %s", out)
	}
}

// 场景：配额不足 → HJ-2004 预拒，且不产生用量记录。
func TestQuotaRejectsWhenExhausted(t *testing.T) {
	upstream := fakeOpenAI(t, 0, 0, "内容")
	defer upstream.Close()
	ts, st, _ := fullServer(t, 10, oaEndpoint("up", upstream.URL, 10))

	resp, out := postChat(ts, byokChatBody("run_q", `,"user_id":"u1","max_tokens":100`), "u1")
	if resp == nil || resp.StatusCode != http.StatusPaymentRequired {
		t.Fatalf("配额不足应 402: %s", out)
	}
	if !strings.Contains(out, "HJ-2004") {
		t.Fatalf("应携带 HJ-2004: %s", out)
	}
	sample, _ := st.SampleUsage(10)
	if len(sample) != 0 {
		t.Fatalf("预拒请求不应产生用量记录: %+v", sample)
	}
}

// 场景：成功结算释放差值；全败全额退款——余额不残留。
func TestQuotaSettleAndRefundBalance(t *testing.T) {
	upstream := fakeOpenAI(t, 0, 0, "内容")
	bad := fakeOpenAI(t, 0, http.StatusForbidden, "")
	defer upstream.Close()
	defer bad.Close()
	ts, st, _ := fullServer(t, 1000, oaEndpoint("up", upstream.URL, 10))

	// 成功：预扣 max_tokens=200，实际 3+4=7
	if resp, out := postChat(ts, byokChatBody("run_ok", `,"user_id":"u1","max_tokens":200`), "u1"); resp == nil || resp.StatusCode != 200 {
		t.Fatalf("成功请求失败: %s", out)
	}
	b, _ := st.Balance("u1")
	if b.Used != 7 || b.Reserved != 0 {
		t.Fatalf("结算后账目错误: %+v", b)
	}

	// 已用 7，预算 1000：请求 max_tokens 994 应被拒（7+994 > 1000）
	if resp, out := postChat(ts, byokChatBody("run_over", `,"user_id":"u1","max_tokens":994`), "u1"); resp == nil || resp.StatusCode != http.StatusPaymentRequired {
		t.Fatalf("超预算应拒: %s", out)
	}

	// 全候选失败 → 全额退款
	tsBad, stBad, _ := fullServer(t, 1000, oaEndpoint("bad", bad.URL, 10))
	if resp, out := postChat(tsBad, byokChatBody("run_af", `,"user_id":"u2","max_tokens":500`), "u2"); resp == nil || resp.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("全败应 503: %s", out)
	}
	b2, _ := stBad.Balance("u2")
	if b2.Reserved != 0 || b2.Used != 0 {
		t.Fatalf("全败后应全额退款: %+v", b2)
	}
}

// 场景：用量归属计量——/me/usage 按用户过滤并附配额总览。
func TestUsageAttributionByUser(t *testing.T) {
	upstream := fakeOpenAI(t, 0, 0, "内容")
	defer upstream.Close()
	ts, _, _ := fullServer(t, 0, oaEndpoint("up", upstream.URL, 10))

	for _, b := range []string{
		byokChatBody("run_u1", `,"user_id":"u1"`),
		byokChatBody("run_u2", `,"user_id":"u2"`),
		byokChatBody("run_u1b", `,"user_id":"u1"`),
	} {
		if resp, out := postChat(ts, b, ""); resp == nil || resp.StatusCode != 200 {
			t.Fatalf("请求失败: %s", out)
		}
	}

	req, _ := http.NewRequest(http.MethodGet, ts.URL+"/me/usage", nil)
	req.Header.Set("X-Hunjing-User-Id", "u1")
	resp, err := ts.Client().Do(req)
	if err != nil {
		t.Fatal(err)
	}
	var payload struct {
		Rows []struct {
			UserID   string `json:"user_id"`
			Requests int    `json:"requests"`
		} `json:"rows"`
		Quota *struct {
			Budget int64 `json:"budget"`
			Used   int64 `json:"used"`
		} `json:"quota"`
	}
	_ = json.NewDecoder(resp.Body).Decode(&payload)
	if len(payload.Rows) != 1 || payload.Rows[0].UserID != "u1" || payload.Rows[0].Requests != 2 {
		t.Fatalf("u1 聚合错误: %+v", payload.Rows)
	}
	if payload.Quota == nil || payload.Quota.Used != 14 {
		t.Fatalf("配额总览错误: %+v", payload.Quota)
	}
}