hunjing/shenshi/internal/server/failover_test.go

package server

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

	"hunjing/shenshi/internal/route"
	"hunjing/shenshi/internal/transform"
)

// fakeOpenAI 构造一个 OpenAI 兼容 SSE 假上游：可控延迟/状态码/内容。
func fakeOpenAI(t *testing.T, delay time.Duration, status int, content string) *httptest.Server {
	t.Helper()
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		if status != 0 {
			w.WriteHeader(status)
			_, _ = w.Write([]byte(`{"error":{"message":"boom"}}`))
			return
		}
		if delay > 0 {
			time.Sleep(delay)
		}
		w.Header().Set("content-type", "text/event-stream")
		flusher := w.(http.Flusher)
		_, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"" + content + "\"}}]}\n\n"))
		flusher.Flush()
		_, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":4}}\n\n"))
		flusher.Flush()
		_, _ = w.Write([]byte("data: [DONE]\n\n"))
	}))
}

func registryWith(t *testing.T, firstByte time.Duration, eps ...route.Endpoint) *route.Registry {
	t.Helper()
	r := route.New()
	r.FirstByteTimeout = firstByte
	// 清掉默认 stub，仅保留测试端点
	r.Reset("chat")
	for _, ep := range eps {
		r.Register("chat", ep)
	}
	return r
}

func oaEndpoint(name, baseURL string, priority int) route.Endpoint {
	return route.Endpoint{
		Name: name, Type: "openai_compatible", Priority: priority,
		Adaptor: &transform.OpenAICompatAdaptor{EndpointName: name, BaseURL: baseURL, UpstreamModel: "test-model"},
	}
}

func foPostChat(ts *httptest.Server, body string) (*http.Response, string) {
	resp, err := ts.Client().Post(ts.URL+"/v1/chat/completions", "application/json", strings.NewReader(body))
	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 foChatBody(run string) string {
	return `{"ir_version":1,"slot":"chat","messages":[{"role":"user","content":"你好"}],"run_id":"` + run + `","stream":true}`
}

// 场景：首端点首字超时 → 切道到次端点成功（不向客户端暴露失败）。
func TestFailoverOnFirstByteTimeout(t *testing.T) {
	slow := fakeOpenAI(t, 500*time.Millisecond, 0, "慢端点")
	fast := fakeOpenAI(t, 0, 0, "快端点")
	defer slow.Close()
	defer fast.Close()

	srv := New(Config{}, registryWith(t, 50*time.Millisecond, oaEndpoint("slow", slow.URL, 10), oaEndpoint("fast", fast.URL, 20)))
	ts := httptest.NewServer(srv)
	defer ts.Close()

	resp, body := foPostChat(ts, foChatBody("run_fo"))
	if resp == nil || resp.StatusCode != http.StatusOK {
		t.Fatalf("状态码异常：%v", resp)
	}
	if !strings.Contains(body, "快端点") || strings.Contains(body, "慢端点") {
		t.Fatalf("应切道到快端点：%s", body)
	}
	if !strings.Contains(body, `"finish_reason":"stop"`) || !strings.Contains(body, "data: [DONE]") {
		t.Fatalf("缺少终帧/DONE：%s", body)
	}
}

// 场景：首端点 401 → 凭证级冷却；下一请求直接跳过它走次端点。
func TestCredentialCooldownSkipsEndpoint(t *testing.T) {
	bad := fakeOpenAI(t, 0, http.StatusUnauthorized, "")
	good := fakeOpenAI(t, 0, 0, "好端点")
	defer bad.Close()
	defer good.Close()

	srv := New(Config{}, registryWith(t, time.Second, oaEndpoint("bad", bad.URL, 10), oaEndpoint("good", good.URL, 20)))
	ts := httptest.NewServer(srv)
	defer ts.Close()

	// 第一次：bad 401 → 冷却 → 切 good
	resp, body := foPostChat(ts, foChatBody("run_a"))
	if resp == nil || resp.StatusCode != http.StatusOK || !strings.Contains(body, "好端点") {
		t.Fatalf("首次应切到好端点：%v %s", resp, body)
	}
	// 第二次：bad 仍在冷却，直接命中 good（若未跳过会再花一次 401 往返——这里断言成功即可）
	resp2, body2 := foPostChat(ts, foChatBody("run_b"))
	if resp2 == nil || resp2.StatusCode != http.StatusOK || !strings.Contains(body2, "好端点") {
		t.Fatalf("冷却后应直达好端点：%v", resp2)
	}
}

// 场景：全部候选失败且未写字节 → HJ-2001 JSON（不是半截 SSE）。
func TestAllCandidatesFailedReturnsJsonError(t *testing.T) {
	bad := fakeOpenAI(t, 0, http.StatusBadGateway, "")
	defer bad.Close()
	srv := New(Config{}, registryWith(t, time.Second, oaEndpoint("bad", bad.URL, 10)))
	ts := httptest.NewServer(srv)
	defer ts.Close()

	resp, body := foPostChat(ts, foChatBody("run_x"))
	if resp.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("状态码 = %d", resp.StatusCode)
	}
	if !strings.Contains(body, "HJ-2001") || strings.Contains(body, "text/event-stream") {
		t.Fatalf("应为 JSON 错误而非 SSE：%s", body)
	}
}

// 场景：用量聚合端点返回 ok 计数与 token。
func TestUsageAggregate(t *testing.T) {
	good := fakeOpenAI(t, 0, 0, "内容")
	defer good.Close()
	srv := New(Config{}, registryWith(t, time.Second, oaEndpoint("good", good.URL, 10)))
	ts := httptest.NewServer(srv)
	defer ts.Close()

	if resp, _ := foPostChat(ts, foChatBody("run_u")); resp == nil || resp.StatusCode != 200 {
		t.Fatal("请求失败")
	}
	usageResp, err := ts.Client().Get(ts.URL + "/me/usage")
	if err != nil {
		t.Fatal(err)
	}
	var agg []map[string]any
	defer func() { _ = usageResp.Body.Close() }()
	if err := json.NewDecoder(usageResp.Body).Decode(&agg); err != nil {
		t.Fatal(err)
	}
	if len(agg) != 1 || agg[0]["ok"].(float64) != 1 || agg[0]["prompt_tokens"].(float64) != 3 {
		t.Fatalf("聚合不符：%v", agg)
	}
}