package proxy import ( "encoding/json" "strings" "testing" ) func TestRewriteModelFieldTopLevelOnly(t *testing.T) { in := `{"model":"a","messages":[{"role":"user","content":"{\"model\":\"nested\"}"}],"max_tokens":5}` out, ok := rewriteModelField([]byte(in), "b") if !ok { t.Fatal("应改写顶层 model") } var m map[string]any if err := json.Unmarshal(out, &m); err != nil { t.Fatalf("输出必须是合法 JSON: %v\n%s", err, out) } if m["model"] != "b" { t.Fatalf("顶层 model 未改写:%v", m["model"]) } // 嵌套的 model 引用必须原样保留 if !strings.Contains(string(out), `\"model\":\"nested\"`) { t.Fatalf("嵌套内容被改动了:%s", out) } if m["max_tokens"].(float64) != 5 { t.Fatalf("其它字段丢失:%v", out) } } func TestRewriteModelFieldKeepsByteShape(t *testing.T) { in := "{\n \"model\" : \"old\",\n \"stream\": true\n}" out, ok := rewriteModelField([]byte(in), "new") if !ok { t.Fatal("应改写") } want := "{\n \"model\" : \"new\",\n \"stream\": true\n}" if string(out) != want { t.Fatalf("应只替换值本身,实际:%q", out) } } func TestRewriteModelFieldMissing(t *testing.T) { if _, ok := rewriteModelField([]byte(`{"stream":true}`), "x"); ok { t.Fatal("没有 model 时不应改写") } if _, ok := rewriteModelField([]byte(`not json`), "x"); ok { t.Fatal("非 JSON 不应改写") } } func TestInjectStreamOptions(t *testing.T) { cases := []struct{ in, want string }{ {`{"model":"m"}`, `{"model":"m","stream_options":{"include_usage":true}}`}, {"{\n \"model\": \"m\"\n}", "{\n \"model\": \"m\"\n,\"stream_options\":{\"include_usage\":true}}"}, {`{}`, `{"stream_options":{"include_usage":true}}`}, // 空对象也能安全追加 } for _, c := range cases { out, ok := injectStreamOptions([]byte(c.in), `{"include_usage":true}`) if !ok { t.Fatalf("应注入:%s", c.in) } if string(out) != c.want { t.Fatalf("注入结果 = %q,期望 %q", out, c.want) } var m map[string]any if err := json.Unmarshal(out, &m); err != nil { t.Fatalf("注入后不是合法 JSON: %v (%s)", err, out) } if _, ok := m["stream_options"]; !ok { t.Fatalf("缺少 stream_options: %s", out) } } // 已有 stream_options 时不重复注入 if _, ok := injectStreamOptions([]byte(`{"stream_options":{"include_usage":true}}`), `{}`); ok { t.Fatal("已有该键时不应再注入") } } // 快路径必须与 JSON 解析结果一致(含转义、非正文事件、三协议形态)。 func TestFastContentTextMatchesJSON(t *testing.T) { lines := []string{ `data: {"choices":[{"delta":{"content":"hello"}}]}`, `data: {"choices":[{"delta":{"content":"中文内容"}}]}`, `data: {"choices":[{"delta":{"content":"带\"引号\"和\\反斜杠"}}]}`, `data: {"choices":[{"delta":{"role":"assistant"}}]}`, `data: {"choices":[{"delta":{},"finish_reason":"stop"}]}`, `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"块文本"}}`, `data: {"type":"response.output_text.delta","delta":"响应增量"}`, `data: {"type":"message_start","message":{"usage":{}}}`, `data: [DONE]`, } for _, line := range lines { got := fastContentText([]byte(line)) want := sseContentText([]byte(line)) if got != want { t.Fatalf("不一致\n行: %s\n快路径: %q\nJSON: %q", line, got, want) } } }