ModelRouter/e2e/wire_golden_test.go

package e2e

import (
	"bufio"
	"encoding/json"
	"io"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"testing"
)

// mustAddr 选一个空闲端口并立即释放（与 e2e_test.go 的做法一致）。
func mustAddr(t *testing.T) string {
	t.Helper()
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("pick port: %v", err)
	}
	addr := l.Addr().String()
	l.Close()
	return addr
}

// 本文件建立**字节级**对外输出基线，供统一线格式 v2（对外序列化移交适配器）
// 做回归比对。
//
// 为什么必须有它：Phase 2 要把 gateway/chat.go 里硬编码的 OpenAI wire 结构与
// writeChatCompletion/pumpStream 的 JSON 拼装移交给适配器 serializer。现有 e2e
// 只用 strings.Contains 做宽松断言（"包含 Hel"），无法证明"输出字节未变"——
// 字段集、omitempty 裁剪、usage 零值处理都可能静默漂移，而客户端会因此看到
// 不同的 wire 形状。golden 把当前输出钉死，改动前后逐字节比对。
//
// 更新方式：GOLDEN_UPDATE=1 go test -tags luajit ./e2e/ -run TestWireGolden
// 更新后必须人工审阅 diff —— golden 是契约，不是快照垃圾桶。

// volatile 命中每次请求都变的字段，比对前替换为占位符。
// 允许冒号后的空白：golden 是 pretty-print 过的（`"created": 123`）。
var volatile = []*regexp.Regexp{
	regexp.MustCompile(`"id":\s*"[^"]*"`),
	regexp.MustCompile(`"created":\s*\d+`),
}

func scrub(s string) string {
	s = volatile[0].ReplaceAllString(s, `"id": "<ID>"`)
	return volatile[1].ReplaceAllString(s, `"created": <TS>`)
}

func compareGolden(t *testing.T, name, got string) {
	t.Helper()
	got = scrub(got)
	path := filepath.Join("..", ".snapshots", name)
	if os.Getenv("GOLDEN_UPDATE") == "1" {
		if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
			t.Fatal(err)
		}
		if err := os.WriteFile(path, []byte(got), 0o644); err != nil {
			t.Fatal(err)
		}
		t.Logf("golden updated: %s", path)
		return
	}
	want, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("golden missing (%s): %v\nrun GOLDEN_UPDATE=1 go test -tags luajit ./e2e/ -run TestWireGolden to create", path, err)
	}
	if got != string(want) {
		t.Errorf("对外输出与 golden 漂移 (%s)\n--- want ---\n%s\n--- got ---\n%s",
			name, string(want), got)
	}
}

// prettyJSON 缩进输出。走 map 会排序键，从而让「字段集」成为唯一被比较的东西，
// 忽略 Go struct 声明顺序（那不是对外契约的一部分）。
func prettyJSON(t *testing.T, raw string) string {
	t.Helper()
	var v interface{}
	if err := json.Unmarshal([]byte(raw), &v); err != nil {
		return raw
	}
	out, err := json.MarshalIndent(v, "", "  ")
	if err != nil {
		return raw
	}
	return string(out)
}

// TestWireGoldenNonStreamChat 钉死非流式 /v1/chat/completions 的字段集。
func TestWireGoldenNonStreamChat(t *testing.T) {
	up := newMockUpstream(t)
	dir := t.TempDir()
	addr := mustAddr(t)
	g := startGateway(t, buildBinary(t), addr, writeConfig(t, dir, addr, map[string]*mockUpstream{"good": up}))

	resp, body := g.do("POST", "/v1/chat/completions",
		`{"model":"good/good-m","messages":[{"role":"user","content":"hi"}]}`, true)

	// 语义：OpenAI 客户端依赖这些顶层字段
	var obj map[string]interface{}
	if err := json.Unmarshal([]byte(body), &obj); err != nil {
		t.Fatalf("not JSON: %v (%s)", err, body)
	}
	for _, k := range []string{"id", "object", "created", "model", "choices", "usage"} {
		if _, ok := obj[k]; !ok {
			t.Errorf("非流式响应缺少 OpenAI 必需字段 %q: %s", k, body)
		}
	}
	if obj["object"] != "chat.completion" {
		t.Errorf("object = %v, want chat.completion", obj["object"])
	}
	if resp.StatusCode != 200 {
		t.Errorf("status = %d, want 200", resp.StatusCode)
	}

	compareGolden(t, "chat_nonstream.json", prettyJSON(t, body))
}

// TestWireGoldenStreamChat 钉死流式 SSE 的帧形状与终止序列。
func TestWireGoldenStreamChat(t *testing.T) {
	up := newMockUpstream(t)
	dir := t.TempDir()
	addr := mustAddr(t)
	g := startGateway(t, buildBinary(t), addr, writeConfig(t, dir, addr, map[string]*mockUpstream{"good": up}))

	resp, _ := doStream(t, g, `{"model":"good/good-m","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
	defer resp.Body.Close()

	var frames []string
	sc := bufio.NewScanner(resp.Body)
	for sc.Scan() {
		line := sc.Text()
		if !strings.HasPrefix(line, "data: ") {
			continue
		}
		payload := strings.TrimPrefix(line, "data: ")
		if payload == "[DONE]" {
			frames = append(frames, "[DONE]")
			continue
		}
		frames = append(frames, prettyJSON(t, payload))
	}
	if len(frames) < 2 {
		t.Fatalf("期望多帧 SSE，实得 %d: %v", len(frames), frames)
	}
	if frames[len(frames)-1] != "[DONE]" {
		t.Errorf("末帧 = %q, want [DONE]", frames[len(frames)-1])
	}
	// 首帧必须是 role 前导：OpenAI 客户端靠它建立 assistant 消息
	if !strings.Contains(frames[0], `"role"`) {
		t.Errorf("首个 SSE 帧必须携带 role 前导，实得: %s", frames[0])
	}

	compareGolden(t, "chat_stream.txt", strings.Join(frames, "\n"))
}

// TestWireGoldenImage 钉死 /v1/images/generations 的响应形状。
func TestWireGoldenImage(t *testing.T) {
	up := newMockUpstream(t)
	dir := t.TempDir()
	addr := mustAddr(t)
	g := startGateway(t, buildBinary(t), addr, writeConfig(t, dir, addr, map[string]*mockUpstream{"image": up}))

	resp, body := g.do("POST", "/v1/images/generations", `{"model":"imagegen/flux-1","prompt":"a cat"}`, true)
	if resp.StatusCode != 200 {
		t.Fatalf("image status = %d: %s", resp.StatusCode, body)
	}
	if !strings.Contains(body, "QUJD") {
		t.Fatalf("生图响应丢失了载荷: %s", body)
	}
	compareGolden(t, "image_generation.json", prettyJSON(t, body))
}

// doStream 发流式请求并返回未读完的响应（调用方负责逐行读 SSE）。
func doStream(t *testing.T, g *gatewayUnderTest, body string) (*http.Response, error) {
	t.Helper()
	req, err := http.NewRequest("POST", "http://"+g.addr+"/v1/chat/completions", strings.NewReader(body))
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+g.key)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("stream request: %v", err)
	}
	return resp, nil
}

var _ = io.Discard