m1saka/token_thief
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: restore reviewable migration evidence

Browse Source
This commit is contained in:
MiMoCode
2026-07-10 18:26:48 +08:00
commit d1bbb5370c
42 changed files with 6419 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/README.md
+61
View File
@@ -0,0 +1,61 @@
# tests/
集中放置项目的测试代码。每个子目录对应被测包，使用 Go 的黑盒测试包模式（`package xxx_test`）通过被测包的导出 API 进行验证。
## 目录结构
```
tests/
├── config/ # 配置与 glob 过滤器测试
├── logger/ # 异步日志队列与重试逻辑测试
├── proxy/ # 反向代理测试（WebSocket、SSE、上游错误）
└── scripts/ # 手工端到端测试脚本与 DB 校验工具
```
## 运行
```bash
go test ./tests/...
```
或者运行某一个子包：
```bash
go test ./tests/proxy/...
```
## 端到端冒烟测试
先构建本地二进制：
```powershell
go build -o TokenThief.exe .
```
加载 `.env` 并设置 newapi Key：
```powershell
. .\tests\scripts\load-env.ps1
$env:NEWAPI_KEY = "sk-..."
```
`.env` 需要包含可写的 `CLICKHOUSE_URL`，例如 `clickhouse://tokenthief:tokenthief@localhost:9000/tokenthief`。
运行聊天端点冒烟测试：
```powershell
.\tests\scripts\smoke.ps1 -Model "gpt-5.4-mini"
```
按 request_id 打印数据库里的流式 `response_body`：
```powershell
.\tests\scripts\dump-stream-body.ps1 -RequestID "<request_id>"
```
## 构建排除
- 测试目录中的 `_test.go` 文件不会参与 `go build`。
- `tests/scripts/` 只放手工测试脚本，不被主程序 import。
- `Dockerfile` 仅构建 main 包（`./`），不会触及 `tests/`。
- 仓库根目录的 `.dockerignore` 把 `tests/` 整体排除在 build context 之外，镜像中不会包含测试代码或测试脚本。
tests/config/config_test.go
+283
View File
@@ -0,0 +1,283 @@
package config_test
import (
"net/netip"
"strings"
"testing"
"time"
"git.misaka.ren/M1saka/token_thief/config"
)
func TestLoadTimeoutDefaults(t *testing.T) {
t.Setenv("UPSTREAM_URL", "https://example.com")
t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db")
cfg, err := config.Load()
if err != nil {
t.Fatalf("Load: %v", err)
}
if cfg.ReadTimeout != 30*time.Second {
t.Fatalf("ReadTimeout=%s want 30s", cfg.ReadTimeout)
}
if cfg.WriteTimeout != 10*time.Minute {
t.Fatalf("WriteTimeout=%s want 10m", cfg.WriteTimeout)
}
if cfg.IdleTimeout != 5*time.Minute {
t.Fatalf("IdleTimeout=%s want 5m", cfg.IdleTimeout)
}
if cfg.UpstreamTimeout != 30*time.Second {
t.Fatalf("UpstreamTimeout=%s want 30s", cfg.UpstreamTimeout)
}
if cfg.UpstreamResponseTimeout != 30*time.Second {
t.Fatalf("UpstreamResponseTimeout=%s want 30s", cfg.UpstreamResponseTimeout)
}
if cfg.UpstreamStreamIdleTimeout != 2*time.Minute {
t.Fatalf("UpstreamStreamIdleTimeout=%s want 2m", cfg.UpstreamStreamIdleTimeout)
}
}
func TestLoadQueueDefaults(t *testing.T) {
t.Setenv("UPSTREAM_URL", "https://example.com")
t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db")
cfg, err := config.Load()
if err != nil {
t.Fatalf("Load: %v", err)
}
if cfg.MaxBodyBytes != 1<<20 {
t.Fatalf("MaxBodyBytes=%d want %d", cfg.MaxBodyBytes, 1<<20)
}
if cfg.LogQueueSize != 256 {
t.Fatalf("LogQueueSize=%d want 256", cfg.LogQueueSize)
}
if cfg.LogQueueBytes != 64<<20 {
t.Fatalf("LogQueueBytes=%d want %d", cfg.LogQueueBytes, 64<<20)
}
}
func TestLoadTimeoutOverrides(t *testing.T) {
t.Setenv("UPSTREAM_URL", "https://example.com")
t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db")
t.Setenv("READ_TIMEOUT", "10s")
t.Setenv("WRITE_TIMEOUT", "20s")
t.Setenv("IDLE_TIMEOUT", "30s")
t.Setenv("UPSTREAM_TIMEOUT", "40s")
t.Setenv("UPSTREAM_RESPONSE_TIMEOUT", "50s")
t.Setenv("UPSTREAM_STREAM_IDLE_TIMEOUT", "60s")
cfg, err := config.Load()
if err != nil {
t.Fatalf("Load: %v", err)
}
if cfg.ReadTimeout != 10*time.Second {
t.Fatalf("ReadTimeout=%s want 10s", cfg.ReadTimeout)
}
if cfg.WriteTimeout != 20*time.Second {
t.Fatalf("WriteTimeout=%s want 20s", cfg.WriteTimeout)
}
if cfg.IdleTimeout != 30*time.Second {
t.Fatalf("IdleTimeout=%s want 30s", cfg.IdleTimeout)
}
if cfg.UpstreamTimeout != 40*time.Second {
t.Fatalf("UpstreamTimeout=%s want 40s", cfg.UpstreamTimeout)
}
if cfg.UpstreamResponseTimeout != 50*time.Second {
t.Fatalf("UpstreamResponseTimeout=%s want 50s", cfg.UpstreamResponseTimeout)
}
if cfg.UpstreamStreamIdleTimeout != 60*time.Second {
t.Fatalf("UpstreamStreamIdleTimeout=%s want 60s", cfg.UpstreamStreamIdleTimeout)
}
}
func TestLoadRejectsInvalidTypedValues(t *testing.T) {
tests := []struct {
name string
key string
value string
}{
{name: "int", key: "LOG_QUEUE_SIZE", value: "not-an-int"},
{name: "int64", key: "MAX_BODY_BYTES", value: "1.5"},
{name: "duration", key: "READ_TIMEOUT", value: "30"},
{name: "bool", key: "UPSTREAM_TLS_INSECURE_SKIP_VERIFY", value: "yes"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Setenv("UPSTREAM_URL", "https://example.com")
t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db")
t.Setenv(tt.key, tt.value)
_, err := config.Load()
if err == nil {
t.Fatalf("Load succeeded with %s=%q", tt.key, tt.value)
}
if !strings.Contains(err.Error(), tt.key) {
t.Fatalf("err=%q does not identify %s", err, tt.key)
}
})
}
}
func TestLoadRejectsNonPositiveNumericAndDurationValues(t *testing.T) {
tests := []struct {
key string
value string
}{
{key: "MAX_BODY_BYTES", value: "0"},
{key: "LOG_QUEUE_SIZE", value: "-1"},
{key: "LOG_QUEUE_BYTES", value: "0"},
{key: "LOG_BATCH_SIZE", value: "0"},
{key: "LOG_BATCH_INTERVAL", value: "-1s"},
{key: "LOG_WORKERS", value: "0"},
{key: "DB_RECONNECT_INTERVAL", value: "0s"},
{key: "READ_TIMEOUT", value: "0s"},
{key: "WRITE_TIMEOUT", value: "0s"},
{key: "IDLE_TIMEOUT", value: "0s"},
{key: "UPSTREAM_TIMEOUT", value: "0s"},
{key: "UPSTREAM_RESPONSE_TIMEOUT", value: "0s"},
{key: "UPSTREAM_STREAM_IDLE_TIMEOUT", value: "0s"},
}
for _, tt := range tests {
t.Run(tt.key, func(t *testing.T) {
t.Setenv("UPSTREAM_URL", "https://example.com")
t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db")
t.Setenv(tt.key, tt.value)
_, err := config.Load()
if err == nil {
t.Fatalf("Load succeeded with %s=%q", tt.key, tt.value)
}
if !strings.Contains(err.Error(), tt.key) {
t.Fatalf("err=%q does not identify %s", err, tt.key)
}
})
}
}
func TestLoadTrustedProxies(t *testing.T) {
t.Setenv("UPSTREAM_URL", "https://example.com")
t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db")
t.Setenv("TRUSTED_PROXIES", "10.0.0.0/8, 192.0.2.1,2001:db8::/32, 2001:db8::1")
cfg, err := config.Load()
if err != nil {
t.Fatalf("Load: %v", err)
}
want := []netip.Prefix{
netip.MustParsePrefix("10.0.0.0/8"),
netip.MustParsePrefix("192.0.2.1/32"),
netip.MustParsePrefix("2001:db8::/32"),
netip.MustParsePrefix("2001:db8::1/128"),
}
if len(cfg.TrustedProxies) != len(want) {
t.Fatalf("TrustedProxies=%v want %v", cfg.TrustedProxies, want)
}
for i := range want {
if cfg.TrustedProxies[i] != want[i] {
t.Fatalf("TrustedProxies[%d]=%v want %v", i, cfg.TrustedProxies[i], want[i])
}
}
}
func TestLoadRejectsInvalidTrustedProxies(t *testing.T) {
for _, value := range []string{"not-an-ip", "10.0.0.0/33", "10.0.0.1/8", "10.0.0.1,,192.0.2.1"} {
t.Run(value, func(t *testing.T) {
t.Setenv("UPSTREAM_URL", "https://example.com")
t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db")
t.Setenv("TRUSTED_PROXIES", value)
_, err := config.Load()
if err == nil {
t.Fatalf("Load succeeded with TRUSTED_PROXIES=%q", value)
}
})
}
}
func TestLoadUpstreamTLSInsecureSkipVerifyDefault(t *testing.T) {
t.Setenv("UPSTREAM_URL", "https://example.com")
t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db")
cfg, err := config.Load()
if err != nil {
t.Fatalf("Load: %v", err)
}
if cfg.UpstreamTLSInsecureSkipVerify {
t.Fatal("UpstreamTLSInsecureSkipVerify=true want false")
}
}
func TestLoadUpstreamTLSInsecureSkipVerifyOverride(t *testing.T) {
t.Setenv("UPSTREAM_URL", "https://example.com")
t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db")
t.Setenv("UPSTREAM_TLS_INSECURE_SKIP_VERIFY", "true")
cfg, err := config.Load()
if err != nil {
t.Fatalf("Load: %v", err)
}
if !cfg.UpstreamTLSInsecureSkipVerify {
t.Fatal("UpstreamTLSInsecureSkipVerify=false want true")
}
}
func TestLoadRequiresClickHouseURL(t *testing.T) {
t.Setenv("UPSTREAM_URL", "https://example.com")
t.Setenv("DATABASE_URL", "postgres://user:pass@localhost/db")
_, err := config.Load()
if err == nil {
t.Fatal("Load succeeded without CLICKHOUSE_URL")
}
if err.Error() != "CLICKHOUSE_URL is required" {
t.Fatalf("err=%q want CLICKHOUSE_URL is required", err.Error())
}
}
func TestLoadValidatesClickHouseURL(t *testing.T) {
tests := []struct {
name string
dsn string
}{
{name: "missing port", dsn: "clickhouse://user:pass@localhost/db"},
{name: "unsupported scheme", dsn: "https://localhost:9440/db"},
{name: "plaintext skip verify", dsn: "clickhouse://localhost:9000/db?skip_verify=true"},
{name: "conflicting secure parameter", dsn: "clickhouses://localhost:9440/db?secure=false"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Setenv("UPSTREAM_URL", "https://example.com")
t.Setenv("CLICKHOUSE_URL", tt.dsn)
_, err := config.Load()
if err == nil {
t.Fatalf("Load succeeded with CLICKHOUSE_URL=%q", tt.dsn)
}
if !strings.Contains(err.Error(), "CLICKHOUSE_URL") {
t.Fatalf("err=%q does not identify CLICKHOUSE_URL", err)
}
})
}
}
func TestLoadAcceptsSecureClickHouseURL(t *testing.T) {
t.Setenv("UPSTREAM_URL", "https://example.com")
t.Setenv("CLICKHOUSE_URL", "clickhouses://user:pass@localhost:9440/db?skip_verify=false")
cfg, err := config.Load()
if err != nil {
t.Fatalf("Load: %v", err)
}
if cfg.ClickHouseURL != "clickhouses://user:pass@localhost:9440/db?skip_verify=false" {
t.Fatalf("ClickHouseURL=%q", cfg.ClickHouseURL)
}
}
tests/config/filter_test.go
+164
View File
@@ -0,0 +1,164 @@
package config_test
import (
"os"
"path/filepath"
"testing"
"git.misaka.ren/M1saka/token_thief/config"
)
func TestCompileGlobMatch(t *testing.T) {
cases := []struct {
pattern string
target string
want bool
}{
// 单段 *
{"/v1/audio/*", "/v1/audio/speech", true},
{"/v1/audio/*", "/v1/audio/transcriptions", true},
{"/v1/audio/*", "/v1/audio/sub/x", false},
// 字面匹配
{"/v1/chat/completions", "/v1/chat/completions", true},
{"/v1/chat/completions", "/v1/chat/completions/extra", false},
// 跨段 **
{"/v1/videos/**", "/v1/videos/", true},
{"/v1/videos/**", "/v1/videos/abc", true},
{"/v1/videos/**", "/v1/videos/abc/content", true},
// Gemini 风格 :generateContent
{"/v1beta/models/*:generateContent", "/v1beta/models/gemini-pro:generateContent", true},
{"/v1beta/models/*:generateContent", "/v1beta/models/gemini-1.5-flash:generateContent", true},
{"/v1beta/models/*:generateContent", "/v1beta/models/x/y:generateContent", false},
// engines 嵌套
{"/v1/engines/*/embeddings", "/v1/engines/text-embedding-ada-002/embeddings", true},
{"/v1/engines/*/embeddings", "/v1/engines/a/b/embeddings", false},
// 含正则元字符
{"/v1/models/*", "/v1/models/gpt-4.1", true},
}
for _, c := range cases {
re, err := config.CompileGlob(c.pattern)
if err != nil {
t.Fatalf("compile %q: %v", c.pattern, err)
}
got := re.MatchString(c.target)
if got != c.want {
t.Errorf("%q vs %q: got %v want %v (regex=%s)", c.pattern, c.target, got, c.want, re.String())
}
}
}
func TestFilterShouldLog(t *testing.T) {
f, err := config.NewFilter(config.FilterWhitelist, []string{"/v1/chat/completions", "/v1/videos/*"})
if err != nil {
t.Fatalf("NewFilter: %v", err)
}
if !f.ShouldLog("/v1/chat/completions") {
t.Error("whitelist must allow /v1/chat/completions")
}
if !f.ShouldLog("/v1/videos/abc") {
t.Error("whitelist must allow /v1/videos/abc")
}
if f.ShouldLog("/v1/embeddings") {
t.Error("whitelist must reject /v1/embeddings")
}
fb, err := config.NewFilter(config.FilterBlacklist, []string{"/v1/chat/completions", "/v1/videos/*"})
if err != nil {
t.Fatalf("NewFilter blacklist: %v", err)
}
if fb.ShouldLog("/v1/chat/completions") {
t.Error("blacklist must reject /v1/chat/completions")
}
if !fb.ShouldLog("/v1/embeddings") {
t.Error("blacklist must allow /v1/embeddings")
}
fd, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatalf("NewFilter disabled: %v", err)
}
if !fd.ShouldLog("/anything") {
t.Error("disabled must allow everything")
}
}
// TestRepoFilterOnlyAllowsChatAndCompletions 加载仓库根目录的 filter.yaml，
// 验证默认过滤器只记录聊天与补全端点。
func TestRepoFilterOnlyAllowsChatAndCompletions(t *testing.T) {
cwd, err := os.Getwd()
if err != nil {
t.Fatalf("getwd: %v", err)
}
// 测试目录位于 <repo>/tests/config，filter.yaml 在 <repo>/filter.yaml
path := filepath.Join(cwd, "..", "..", "filter.yaml")
if _, err := os.Stat(path); err != nil {
t.Fatalf("filter.yaml not found: %v", err)
}
f, err := config.LoadFilter(path)
if err != nil {
t.Fatalf("load filter: %v", err)
}
allowed := []string{
// Chat
"/v1/chat/completions",
"/v1/responses",
"/v1/messages",
"/v1beta/models/gemini-1.5-pro:generateContent",
"/v1beta/models/gemini-1.5-pro:generateContent/",
"/v1beta/models/gemini-1.5-pro:streamGenerateContent",
// Completions
"/v1/completions",
}
for _, p := range allowed {
if !f.ShouldLog(p) {
t.Errorf("filter.yaml should match chat/completion endpoint %q", p)
}
}
blocked := []string{
// Models
"/v1/models",
"/v1beta/models",
// Embeddings
"/v1/embeddings",
"/v1/engines/text-embedding-ada-002/embeddings",
// Moderations / Rerank / Realtime
"/v1/moderations",
"/v1/rerank",
"/v1/realtime",
// Audio
"/v1/audio/speech",
"/v1/audio/transcriptions",
"/v1/audio/translations",
// Images
"/v1/images/generations",
"/v1/images/generations/",
"/v1/images/edits",
"/v1/images/edits/",
// Videos - 通用
"/v1/video/generations",
"/v1/video/generations/task_abc",
// Videos - Sora
"/v1/videos",
"/v1/videos/task_abc",
"/v1/videos/task_abc/content",
// Videos - 即梦
"/jimeng/",
// Videos - Kling
"/kling/v1/videos/text2video",
"/kling/v1/videos/text2video/task_abc",
"/kling/v1/videos/image2video",
"/kling/v1/videos/image2video/task_abc",
}
for _, p := range blocked {
if f.ShouldLog(p) {
t.Errorf("filter.yaml should not match non-chat/completion endpoint %q", p)
}
}
}
tests/deployment/compose_test.go
+66
View File
@@ -0,0 +1,66 @@
package deployment_test
import (
"os"
"path/filepath"
"regexp"
"strings"
"testing"
)
func TestComposeUsesSafeDeploymentDefaults(t *testing.T) {
root := filepath.Join("..", "..")
compose := readFile(t, filepath.Join(root, "compose.yml"))
envExample := readFile(t, filepath.Join(root, ".env.example"))
for _, required := range []string{
`${UPSTREAM_URL:?set UPSTREAM_URL}`,
`${CLICKHOUSE_URL:?set CLICKHOUSE_URL with URL-encoded credentials}`,
`${CLICKHOUSE_PASSWORD:?set a strong CLICKHOUSE_PASSWORD}`,
} {
if !strings.Contains(compose, required) {
t.Errorf("compose.yml must reject an empty required setting with %q", required)
}
}
if !regexp.MustCompile(`(?m)^\s+image: clickhouse/clickhouse-server:\d+\.\d+\.\d+\.\d+-alpine\s*$`).MatchString(compose) {
t.Error("ClickHouse image must use an exact version tag")
}
serviceStart := strings.LastIndex(compose, " thief_clickhouse:")
if serviceStart < 0 {
t.Fatal("compose.yml is missing the thief_clickhouse service")
}
clickhouseService := compose[serviceStart:]
if strings.Contains(clickhouseService, "\n ports:") {
t.Error("ClickHouse must not publish host ports by default")
}
for _, emptySecret := range []string{"CLICKHOUSE_URL=\n", "CLICKHOUSE_PASSWORD=\n"} {
if !strings.Contains(envExample, emptySecret) {
t.Errorf(".env.example must leave %q empty", strings.TrimSpace(emptySecret))
}
}
}
func TestReviewDocumentsUseStablePaths(t *testing.T) {
root := filepath.Join("..", "..")
for _, path := range []string{
filepath.Join(root, "docs", "compose", "specs", "reliability-security-fixes.md"),
filepath.Join(root, "docs", "compose", "plans", "reliability-security-fixes.md"),
} {
content := readFile(t, path)
if !strings.Contains(content, "2026-07-09-clickhouse-migration.md") {
t.Errorf("%s must link to the dated source document", path)
}
}
}
func readFile(t *testing.T, path string) string {
t.Helper()
b, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
return string(b)
}
tests/logger/queue_test.go
+323
View File
@@ -0,0 +1,323 @@
package logger_test
import (
"context"
"errors"
"sync"
"testing"
"time"
"git.misaka.ren/M1saka/token_thief/logger"
)
type batchScript struct {
prepareErr error
appendAt int
appendErr error
sendErr error
}
type fakeBatch struct {
script batchScript
rows [][]any
appendCall int
abortCalls int
}
func (f *fakeBatch) Append(v ...any) error {
f.appendCall++
if f.script.appendErr != nil && f.appendCall == f.script.appendAt {
return f.script.appendErr
}
f.rows = append(f.rows, append([]any(nil), v...))
return nil
}
func (f *fakeBatch) Send() error { return f.script.sendErr }
func (f *fakeBatch) Abort() error { f.abortCalls++; return nil }
type fakePool struct {
calls int
queries []string
batches []*fakeBatch
scripts []batchScript
}
func (p *fakePool) PrepareBatch(_ context.Context, query string) (logger.Batch, error) {
idx := p.calls
p.calls++
p.queries = append(p.queries, query)
var script batchScript
if idx < len(p.scripts) {
script = p.scripts[idx]
}
if script.prepareErr != nil {
return nil, script.prepareErr
}
b := &fakeBatch{script: script}
p.batches = append(p.batches, b)
return b, nil
}
func TestFlushSuccessUsesClickHouseTypes(t *testing.T) {
p := &fakePool{}
started := time.Date(2026, 7, 9, 1, 2, 3, 4_000_000, time.UTC)
entry := &logger.LogEntry{
RequestID: "rid", Method: "POST", Path: "/v1", Query: "a=b", ClientIP: "127.0.0.1",
RequestHeaders: []byte(`{"x":"y"}`), RequestBody: []byte("request"), RequestTruncated: true,
StatusCode: 201, ResponseHeaders: []byte(`{"h":"v"}`), ResponseBody: []byte("response"),
ResponseTruncated: true, IsStream: true, LatencyMS: 150, StartedAt: started,
FinishedAt: started.Add(150 * time.Millisecond),
}
result := logger.Flush(context.Background(), p, []*logger.LogEntry{entry})
if result.Err != nil || result.Failed != 0 || len(result.Retry) != 0 {
t.Fatalf("unexpected result: %+v", result)
}
if p.calls != 1 || len(p.batches) != 1 || len(p.batches[0].rows) != 1 {
t.Fatalf("calls=%d batches=%d rows=%d", p.calls, len(p.batches), len(p.batches[0].rows))
}
if contains(p.queries[0], "VALUES") {
t.Fatalf("query contains VALUES: %s", p.queries[0])
}
row := p.batches[0].rows[0]
if len(row) != 17 {
t.Fatalf("columns=%d want 17", len(row))
}
if row[5] != string(entry.RequestHeaders) || row[6] != string(entry.RequestBody) || row[9] != string(entry.ResponseHeaders) || row[10] != string(entry.ResponseBody) {
t.Fatalf("byte fields were not converted to strings: %#v", row)
}
if _, ok := row[8].(int32); !ok {
t.Fatalf("status code type=%T want int32", row[8])
}
if p.batches[0].abortCalls != 0 {
t.Fatalf("abort calls=%d want 0", p.batches[0].abortCalls)
}
}
func TestFlushEmptyBatchDoesNotPrepare(t *testing.T) {
p := &fakePool{}
result := logger.Flush(context.Background(), p, nil)
if result.Err != nil || result.Failed != 0 || result.Ambiguous != 0 || len(result.Retry) != 0 {
t.Fatalf("unexpected result: %+v", result)
}
if p.calls != 0 {
t.Fatalf("PrepareBatch calls=%d want 0", p.calls)
}
}
func TestFlushPrepareFailureIsRetryable(t *testing.T) {
p := &fakePool{scripts: []batchScript{{prepareErr: errors.New("prepare")}}}
entries := []*logger.LogEntry{{RequestID: "a"}, {RequestID: "b"}}
result := logger.Flush(context.Background(), p, entries)
if result.Err == nil || result.Failed != 0 || len(result.Retry) != 2 {
t.Fatalf("unexpected result: %+v", result)
}
if len(p.batches) != 0 {
t.Fatalf("prepare failure created %d batches", len(p.batches))
}
}
func TestFlushAppendFailureIsolatesOnlyBadRow(t *testing.T) {
appendErr := errors.New("bad row")
p := &fakePool{scripts: []batchScript{
{appendAt: 2, appendErr: appendErr},
{},
{appendAt: 1, appendErr: appendErr},
{},
}}
entries := []*logger.LogEntry{{RequestID: "a"}, {RequestID: "bad"}, {RequestID: "c"}}
result := logger.Flush(context.Background(), p, entries)
if result.Err == nil || result.Failed != 1 || len(result.Retry) != 0 {
t.Fatalf("unexpected result: %+v", result)
}
if p.calls != 4 {
t.Fatalf("PrepareBatch calls=%d want 4", p.calls)
}
if p.batches[0].abortCalls != 1 || p.batches[2].abortCalls != 1 {
t.Fatalf("abort calls initial=%d bad-row=%d want 1 each", p.batches[0].abortCalls, p.batches[2].abortCalls)
}
if p.batches[1].abortCalls != 0 || p.batches[3].abortCalls != 0 {
t.Fatalf("successful batches were aborted")
}
}
func TestFlushAppendFailureCanFullyRecover(t *testing.T) {
p := &fakePool{scripts: []batchScript{
{appendAt: 2, appendErr: errors.New("batch append")},
{},
{},
}}
entries := []*logger.LogEntry{{RequestID: "a"}, {RequestID: "b"}}
result := logger.Flush(context.Background(), p, entries)
if result.Err != nil || result.Failed != 0 || len(result.Retry) != 0 {
t.Fatalf("unexpected result: %+v", result)
}
if p.calls != 3 || p.batches[0].abortCalls != 1 {
t.Fatalf("calls=%d abort=%d", p.calls, p.batches[0].abortCalls)
}
}
func TestFlushSendFailureIsAmbiguousAndNotRetried(t *testing.T) {
p := &fakePool{scripts: []batchScript{{sendErr: errors.New("connection lost")}}}
entries := []*logger.LogEntry{{RequestID: "a"}, {RequestID: "b"}}
result := logger.Flush(context.Background(), p, entries)
if result.Err == nil || result.Failed != 2 || len(result.Retry) != 0 {
t.Fatalf("unexpected result: %+v", result)
}
if result.Ambiguous != 2 {
t.Fatalf("Ambiguous=%d want 2", result.Ambiguous)
}
if p.calls != 1 {
t.Fatalf("PrepareBatch calls=%d want 1", p.calls)
}
if p.batches[0].abortCalls != 1 {
t.Fatalf("abort calls=%d want 1", p.batches[0].abortCalls)
}
}
func TestFlushAppendRecoveryCountsAmbiguousSingleSend(t *testing.T) {
p := &fakePool{scripts: []batchScript{
{appendAt: 1, appendErr: errors.New("bad batch")},
{},
{sendErr: errors.New("ack lost")},
}}
entries := []*logger.LogEntry{{RequestID: "a"}, {RequestID: "b"}}
result := logger.Flush(context.Background(), p, entries)
if result.Failed != 1 || result.Ambiguous != 1 || len(result.Retry) != 0 {
t.Fatalf("unexpected result: %+v", result)
}
}
func TestEstimatedBytesCoversStringsAndByteSlices(t *testing.T) {
entry := &logger.LogEntry{
RequestID: "1", Method: "22", Path: "333", Query: "4444", ClientIP: "55555", Error: "666666",
RequestHeaders: []byte("7777777"), RequestBody: []byte("88888888"),
ResponseHeaders: []byte("999999999"), ResponseBody: []byte("0000000000"),
}
if got, want := logger.EstimatedBytes(entry), int64(55); got != want {
t.Fatalf("EstimatedBytes=%d want %d", got, want)
}
}
func TestQueueByteBudgetReleasedAfterFinalDiscard(t *testing.T) {
entry := &logger.LogEntry{
RequestID: "request", Method: "POST", Path: "/path", Query: "q=1", ClientIP: "ip", Error: "error",
RequestHeaders: []byte("rh"), RequestBody: []byte("rb"), ResponseHeaders: []byte("sh"), ResponseBody: []byte("sb"),
}
budget := logger.EstimatedBytes(entry)
q := logger.NewQueue(nil, 4, 2, 1, time.Hour, budget)
q.Submit(entry)
q.Submit(entry)
stats := q.Stats()
if stats.Enqueued != 1 || stats.Dropped != 1 || stats.Bytes != budget {
t.Fatalf("before drain: %+v budget=%d", stats, budget)
}
q.Start(context.Background())
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
if err := q.Stop(ctx); err != nil {
t.Fatalf("Stop: %v", err)
}
stats = q.Stats()
if stats.Bytes != 0 || stats.Failed != 1 {
t.Fatalf("after drain: %+v", stats)
}
}
func TestQueueStopBeforeStartReleasesBudget(t *testing.T) {
entry := &logger.LogEntry{RequestID: "queued"}
q := logger.NewQueue(nil, 2, 2, 1, time.Hour, logger.EstimatedBytes(entry))
q.Submit(entry)
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
if err := q.Stop(ctx); err != nil {
t.Fatalf("Stop: %v", err)
}
if stats := q.Stats(); stats.Bytes != 0 || stats.Failed != 1 {
t.Fatalf("stats=%+v", stats)
}
}
func TestQueueCanceledStopEventuallyReleasesAllBudget(t *testing.T) {
entry := &logger.LogEntry{RequestID: "queued"}
q := logger.NewQueue(nil, 32, 32, 1, time.Hour, 32*logger.EstimatedBytes(entry))
q.Start(context.Background())
for i := 0; i < 32; i++ {
q.Submit(entry)
}
ctx, cancel := context.WithCancel(context.Background())
cancel()
_ = q.Stop(ctx)
deadline := time.Now().Add(time.Second)
for q.Stats().Bytes != 0 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if stats := q.Stats(); stats.Bytes != 0 || stats.Failed != 32 {
t.Fatalf("stats=%+v", stats)
}
}
func TestQueueStopAndSubmitAreConcurrentAndRepeatSafe(t *testing.T) {
q := logger.NewQueue(nil, 32, 8, 2, time.Millisecond, 1<<20)
q.Start(context.Background())
var submitters sync.WaitGroup
for i := 0; i < 8; i++ {
submitters.Add(1)
go func() {
defer submitters.Done()
for j := 0; j < 2_000; j++ {
q.Submit(&logger.LogEntry{RequestID: "x"})
}
}()
}
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
if err := q.Stop(ctx); err != nil {
t.Fatalf("first Stop: %v", err)
}
submitters.Wait()
if err := q.Stop(ctx); err != nil {
t.Fatalf("second Stop: %v", err)
}
q.Submit(&logger.LogEntry{RequestID: "after-stop"})
if stats := q.Stats(); stats.Bytes != 0 {
t.Fatalf("bytes after Stop=%d want 0", stats.Bytes)
}
}
func TestStopIsIndependentFromStartContext(t *testing.T) {
root, cancelRoot := context.WithCancel(context.Background())
q := logger.NewQueue(nil, 4, 4, 1, time.Hour, 1024)
q.Start(root)
q.Submit(&logger.LogEntry{RequestID: "x"})
cancelRoot()
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
if err := q.Stop(ctx); err != nil {
t.Fatalf("Stop after root cancellation: %v", err)
}
if stats := q.Stats(); stats.Failed != 1 || stats.Bytes != 0 {
t.Fatalf("stats=%+v", stats)
}
}
func contains(s, substr string) bool {
for i := 0; i+len(substr) <= len(s); i++ {
if s[i:i+len(substr)] == substr {
return true
}
}
return false
}
tests/proxy/error_test.go
+139
View File
@@ -0,0 +1,139 @@
package proxy_test
import (
"io"
"net/http"
"net/http/httptest"
"net/url"
"testing"
"time"
"git.misaka.ren/M1saka/token_thief/config"
"git.misaka.ren/M1saka/token_thief/logger"
"git.misaka.ren/M1saka/token_thief/proxy"
)
type sliceSubmitter struct{ entries []*logger.LogEntry }
func (s *sliceSubmitter) Submit(e *logger.LogEntry) { s.entries = append(s.entries, e) }
// TestUpstreamErrorRecorded 验证上游不可达时 502 响应被记录、错误信息进入 LogEntry.Error。
func TestUpstreamErrorRecorded(t *testing.T) {
// 指向一个一定不可用的端口
badURL, _ := url.Parse("http://127.0.0.1:1") // port 1 几乎肯定 connection refused
sub := &sliceSubmitter{}
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.New(badURL, filter, sub, 1024)
srv := httptest.NewServer(h)
defer srv.Close()
resp, err := http.Get(srv.URL + "/v1/chat/completions")
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
body, _ := io.ReadAll(resp.Body)
if resp.StatusCode != http.StatusBadGateway {
t.Errorf("status=%d want 502, body=%q", resp.StatusCode, body)
}
if resp.Header.Get("X-Request-Id") == "" {
t.Errorf("missing X-Request-Id header")
}
deadline := time.Now().Add(2 * time.Second)
for len(sub.entries) == 0 && time.Now().Before(deadline) {
time.Sleep(10 * time.Millisecond)
}
if len(sub.entries) == 0 {
t.Fatal("no log entry captured")
}
e := sub.entries[0]
if e.StatusCode != http.StatusBadGateway {
t.Errorf("LogEntry.StatusCode=%d want 502", e.StatusCode)
}
if e.Error == "" {
t.Errorf("LogEntry.Error should be set, got empty")
}
if string(e.ResponseBody) != "bad gateway\n" {
t.Errorf("response_body should be fixed bad gateway, got %q", e.ResponseBody)
}
if e.RequestID == "" {
t.Errorf("LogEntry.RequestID empty")
}
}
// TestModifyResponseSetsRequestID 验证正常上游响应也会带上 X-Request-Id。
func TestModifyResponseSetsRequestID(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(200)
_, _ = w.Write([]byte(`{"ok":true}`))
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &sliceSubmitter{}
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.New(u, filter, sub, 1024)
srv := httptest.NewServer(h)
defer srv.Close()
resp, err := http.Get(srv.URL + "/v1/anything")
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
rid := resp.Header.Get("X-Request-Id")
if len(rid) != 32 {
t.Errorf("X-Request-Id length=%d want 32, value=%q", len(rid), rid)
}
deadline := time.Now().Add(2 * time.Second)
for len(sub.entries) == 0 && time.Now().Before(deadline) {
time.Sleep(10 * time.Millisecond)
}
if len(sub.entries) == 0 {
t.Fatal("no log entry")
}
if sub.entries[0].RequestID != rid {
t.Errorf("LogEntry.RequestID=%q response header=%q (should match)", sub.entries[0].RequestID, rid)
}
}
func TestUpstreamTLSInsecureSkipVerifyAllowsSelfSigned(t *testing.T) {
upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`{"ok":true}`))
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &sliceSubmitter{}
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.NewWithOptions(u, filter, sub, 1024, proxy.Options{UpstreamTLSInsecureSkipVerify: true})
srv := httptest.NewServer(h)
defer srv.Close()
resp, err := http.Get(srv.URL + "/v1/anything")
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
t.Fatalf("status=%d want 200, body=%q", resp.StatusCode, body)
}
}
tests/proxy/robustness_test.go
+336
View File
@@ -0,0 +1,336 @@
package proxy_test
import (
"context"
"errors"
"io"
"net/http"
"net/http/httptest"
"net/netip"
"net/url"
"strings"
"sync/atomic"
"testing"
"time"
"git.misaka.ren/M1saka/token_thief/config"
"git.misaka.ren/M1saka/token_thief/proxy"
)
type failingBody struct{ err error }
func (b failingBody) Read([]byte) (int, error) { return 0, b.err }
func (failingBody) Close() error { return nil }
type lateFailingBody struct {
remaining int
err error
}
func (b *lateFailingBody) Read(p []byte) (int, error) {
if b.remaining == 0 {
return 0, b.err
}
n := min(len(p), b.remaining)
for i := range p[:n] {
p[i] = 'x'
}
b.remaining -= n
return n, nil
}
func (*lateFailingBody) Close() error { return nil }
type failingResponseWriter struct {
header http.Header
short bool
}
func (w *failingResponseWriter) Header() http.Header { return w.header }
func (*failingResponseWriter) WriteHeader(int) {}
func (w *failingResponseWriter) Write(p []byte) (int, error) {
if w.short {
return len(p) - 1, nil
}
return 0, errors.New("write failed")
}
func newTestHandler(t *testing.T, upstream *url.URL, sub *captureSubmitter, opts proxy.Options) *proxy.Handler {
t.Helper()
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
return proxy.NewWithOptions(upstream, filter, sub, 1024*1024, opts)
}
func TestRequestBodyReadFailureReturnsFixed400WithoutUpstream(t *testing.T) {
var upstreamCalls atomic.Int32
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
upstreamCalls.Add(1)
w.WriteHeader(http.StatusOK)
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
h := newTestHandler(t, u, &captureSubmitter{}, proxy.Options{})
req := httptest.NewRequest(http.MethodPost, "http://proxy.test/v1/chat", nil)
req.Body = failingBody{err: errors.New("secret read failure")}
req.ContentLength = 1
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusBadRequest || rec.Body.String() != "bad request\n" {
t.Fatalf("response=(%d, %q), want fixed 400 bad request", rec.Code, rec.Body.String())
}
if upstreamCalls.Load() != 0 {
t.Fatalf("upstream called %d times", upstreamCalls.Load())
}
}
func TestRequestBodyReadFailureAfterCaptureLimitDoesNotReachUpstream(t *testing.T) {
var upstreamCalls atomic.Int32
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
upstreamCalls.Add(1)
_, _ = io.Copy(io.Discard, r.Body)
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
h := newTestHandler(t, u, &captureSubmitter{}, proxy.Options{})
req := httptest.NewRequest(http.MethodPost, "http://proxy.test/v1/chat", nil)
req.Body = &lateFailingBody{remaining: 1024*1024 + 1, err: errors.New("late read failure")}
req.ContentLength = 1024*1024 + 2
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusBadRequest || rec.Body.String() != "bad request\n" {
t.Fatalf("response=(%d, %q), want fixed 400 bad request", rec.Code, rec.Body.String())
}
if upstreamCalls.Load() != 0 {
t.Fatalf("upstream called %d times", upstreamCalls.Load())
}
}
func TestFilteredRequestBodyReadFailureDoesNotReachUpstream(t *testing.T) {
var upstreamCalls atomic.Int32
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
upstreamCalls.Add(1)
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
filter, err := config.NewFilter(config.FilterBlacklist, []string{"/ignored"})
if err != nil {
t.Fatal(err)
}
h := proxy.NewWithOptions(u, filter, &captureSubmitter{}, 1024, proxy.Options{})
req := httptest.NewRequest(http.MethodPost, "http://proxy.test/ignored", nil)
req.Body = failingBody{err: errors.New("read failure")}
req.ContentLength = 1
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusBadRequest || rec.Body.String() != "bad request\n" {
t.Fatalf("response=(%d, %q), want fixed 400 bad request", rec.Code, rec.Body.String())
}
if upstreamCalls.Load() != 0 {
t.Fatalf("upstream called %d times", upstreamCalls.Load())
}
}
func TestBadGatewayResponseDoesNotLeakUpstreamError(t *testing.T) {
badURL, _ := url.Parse("http://127.0.0.1:1")
h := newTestHandler(t, badURL, &captureSubmitter{}, proxy.Options{})
rec := httptest.NewRecorder()
h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/test", nil))
if rec.Code != http.StatusBadGateway || rec.Body.String() != "bad gateway\n" {
t.Fatalf("response=(%d, %q), want fixed 502 bad gateway", rec.Code, rec.Body.String())
}
}
func TestResponseWriteFailurePreventsCommit(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = io.WriteString(w, "response")
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
for _, short := range []bool{false, true} {
sub := &captureSubmitter{}
h := newTestHandler(t, u, sub, proxy.Options{})
w := &failingResponseWriter{header: make(http.Header), short: short}
h.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/test", nil))
if sub.Len() != 0 {
t.Fatalf("short=%v: failed response write was committed", short)
}
}
}
func TestTrustedProxyClientIP(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = io.WriteString(w, "ok")
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
trusted := netip.MustParsePrefix("10.0.0.0/8")
tests := []struct {
name string
peer string
xff string
wantIP string
}{
{name: "untrusted peer ignores xff", peer: "203.0.113.9:1234", xff: "198.51.100.1", wantIP: "203.0.113.9"},
{name: "strip trusted from right", peer: "10.0.0.2:1234", xff: "198.51.100.7, 10.0.0.3", wantIP: "198.51.100.7"},
{name: "ipv6", peer: "[2001:db8::2]:1234", xff: "198.51.100.7", wantIP: "2001:db8::2"},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
sub := &captureSubmitter{}
h := newTestHandler(t, u, sub, proxy.Options{TrustedProxies: []netip.Prefix{trusted}})
req := httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/test", nil)
req.RemoteAddr = tc.peer
req.Header.Set("X-Forwarded-For", tc.xff)
h.ServeHTTP(httptest.NewRecorder(), req)
if sub.Len() != 1 || sub.Entry(0).ClientIP != tc.wantIP {
t.Fatalf("ClientIP=%q, want %q", sub.Entry(0).ClientIP, tc.wantIP)
}
})
}
}
func TestTrustedProxyUsesValidXRealIPWithoutXFF(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
trusted := netip.MustParsePrefix("192.0.2.0/24")
h := newTestHandler(t, u, sub, proxy.Options{TrustedProxies: []netip.Prefix{trusted}})
req := httptest.NewRequest(http.MethodGet, "http://proxy.test/x", nil)
req.RemoteAddr = "192.0.2.10:1234"
req.Header.Set("X-Real-IP", "198.51.100.20")
h.ServeHTTP(httptest.NewRecorder(), req)
if sub.Len() != 1 || sub.Entry(0).ClientIP != "198.51.100.20" {
t.Fatalf("entries=%d", sub.Len())
}
}
func TestResponseBodyTimeoutPreventsCommit(t *testing.T) {
release := make(chan struct{})
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
w.(http.Flusher).Flush()
<-release
}))
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
h := newTestHandler(t, u, sub, proxy.Options{ResponseTimeout: 50 * time.Millisecond})
srv := httptest.NewServer(h)
resp, err := http.Get(srv.URL + "/v1/test")
if err == nil {
_, _ = io.ReadAll(resp.Body)
_ = resp.Body.Close()
}
time.Sleep(50 * time.Millisecond)
if sub.Len() != 0 {
t.Fatal("timed out response must not be committed")
}
close(release)
srv.Close()
upstream.Close()
}
func TestSSEIdleTimeoutResetsAfterReads(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
f := w.(http.Flusher)
for _, event := range []string{"data: one\n\n", "data: two\n\n", "data: [DONE]\n\n"} {
_, _ = io.WriteString(w, event)
f.Flush()
time.Sleep(30 * time.Millisecond)
}
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
h := newTestHandler(t, u, sub, proxy.Options{SSEIdleTimeout: 60 * time.Millisecond})
srv := httptest.NewServer(h)
defer srv.Close()
resp, err := http.Get(srv.URL + "/v1/test")
if err != nil {
t.Fatal(err)
}
_, _ = io.ReadAll(resp.Body)
_ = resp.Body.Close()
if sub.Len() != 1 {
t.Fatalf("entries=%d, want completed SSE commit", sub.Len())
}
}
func TestIncompleteSSEEventPreventsCommit(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
_, _ = io.WriteString(w, "data: partial")
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
h := newTestHandler(t, u, sub, proxy.Options{})
h.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/test", nil))
if sub.Len() != 0 {
t.Fatal("SSE ending mid-event must not be committed as complete")
}
}
func TestIncompleteSSEEventAfterTerminalEventPreventsCommit(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
_, _ = io.WriteString(w, "data: [DONE]\n\n")
w.(http.Flusher).Flush()
_, _ = io.WriteString(w, "data: partial")
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
h := newTestHandler(t, u, sub, proxy.Options{})
h.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/test", nil))
if sub.Len() != 0 {
t.Fatal("SSE ending mid-event after a terminal event must not be committed")
}
}
func TestShutdownReturnsWithoutHijackedConnections(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) }))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
h := newTestHandler(t, u, &captureSubmitter{}, proxy.Options{})
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
if err := h.Shutdown(ctx); err != nil {
t.Fatal(err)
}
}
func TestTerminalTextInNonSSEBodyDoesNotAffectCommit(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
_, _ = io.WriteString(w, `{"text":"data: [DONE]"}`)
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
h := newTestHandler(t, u, sub, proxy.Options{})
h.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/test", strings.NewReader("")))
if sub.Len() != 1 {
t.Fatalf("entries=%d, want 1", sub.Len())
}
}
tests/proxy/sse_capture_test.go
+824
View File
@@ -0,0 +1,824 @@
package proxy_test
import (
"bufio"
"encoding/json"
"fmt"
"io"
"net"
"net/http"
"net/http/httptest"
"net/url"
"strings"
"sync"
"testing"
"time"
"git.misaka.ren/M1saka/token_thief/config"
"git.misaka.ren/M1saka/token_thief/logger"
"git.misaka.ren/M1saka/token_thief/proxy"
)
type captureSubmitter struct {
mu sync.Mutex
entries []*logger.LogEntry
}
func (c *captureSubmitter) Submit(e *logger.LogEntry) {
c.mu.Lock()
defer c.mu.Unlock()
c.entries = append(c.entries, e)
}
func (c *captureSubmitter) Len() int {
c.mu.Lock()
defer c.mu.Unlock()
return len(c.entries)
}
func (c *captureSubmitter) Entry(i int) *logger.LogEntry {
c.mu.Lock()
defer c.mu.Unlock()
return c.entries[i]
}
// fakeOpenAIStreamUpstream 模拟一个 OpenAI 兼容的 SSE 上游：
// 分 5 次往响应里 write 一行 SSE 数据，每次都 Flush。
func fakeOpenAIStreamUpstream() *httptest.Server {
chunks := []string{
`data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}` + "\n\n",
`data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"你"},"finish_reason":null}]}` + "\n\n",
`data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"好"},"finish_reason":null}]}` + "\n\n",
`data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}` + "\n\n",
"data: [DONE]\n\n",
}
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Cache-Control", "no-cache")
w.WriteHeader(http.StatusOK)
flusher := w.(http.Flusher)
for _, ch := range chunks {
_, _ = io.WriteString(w, ch)
flusher.Flush()
time.Sleep(5 * time.Millisecond)
}
}))
}
func fakeAnthropicStreamUpstream() *httptest.Server {
chunks := []string{
`event: message_start` + "\n" + `data: {"type":"message_start","message":{"id":"msg-1","type":"message","role":"assistant","content":[],"model":"claude-3-5-sonnet","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}}` + "\n\n",
`event: content_block_start` + "\n" + `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}` + "\n\n",
`event: content_block_delta` + "\n" + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"你"}}` + "\n\n",
`event: content_block_delta` + "\n" + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"好"}}` + "\n\n",
`event: message_delta` + "\n" + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":3}}` + "\n\n",
`event: message_stop` + "\n" + `data: {"type":"message_stop"}` + "\n\n",
}
return newSSEUpstream(chunks)
}
func fakeGeminiStreamUpstream() *httptest.Server {
chunks := []string{
`data: {"candidates":[{"content":{"parts":[{"text":"你"}],"role":"model"},"index":0}]}` + "\n\n",
`data: {"candidates":[{"content":{"parts":[{"text":"好"}],"role":"model"},"finishReason":"STOP","index":0}],"usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":2,"totalTokenCount":4}}` + "\n\n",
}
return newSSEUpstream(chunks)
}
func fakeUnknownStreamUpstream() *httptest.Server {
return newSSEUpstream([]string{
`event: custom` + "\n" + `data: not-json` + "\n\n",
})
}
func fakeOpenAIStreamUpstreamThatStaysOpen(release <-chan struct{}) *httptest.Server {
chunks := []string{
`data: {"id":"chatcmpl-hang","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}` + "\n\n",
`data: {"id":"chatcmpl-hang","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":null}]}` + "\n\n",
`data: {"id":"chatcmpl-hang","object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}` + "\n\n",
"data: [DONE]\n\n",
}
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
flusher := w.(http.Flusher)
for _, ch := range chunks {
_, _ = io.WriteString(w, ch)
flusher.Flush()
}
<-release
}))
}
func newSSEUpstream(chunks []string) *httptest.Server {
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Cache-Control", "no-cache")
w.WriteHeader(http.StatusOK)
flusher := w.(http.Flusher)
for _, ch := range chunks {
_, _ = io.WriteString(w, ch)
flusher.Flush()
time.Sleep(5 * time.Millisecond)
}
}))
}
func TestSSEChunkAssembly(t *testing.T) {
upstream := fakeOpenAIStreamUpstream()
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.New(u, filter, sub, 1024*1024)
proxySrv := httptest.NewServer(h)
defer proxySrv.Close()
// 客户端走原始 TCP，逐字节读 + 打印，验证流式实时到达
pu, _ := url.Parse(proxySrv.URL)
conn, err := net.DialTimeout("tcp", pu.Host, 3*time.Second)
if err != nil {
t.Fatal(err)
}
defer conn.Close()
fmt.Fprintf(conn, "POST /v1/chat/completions HTTP/1.1\r\nHost: %s\r\nContent-Length: 0\r\n\r\n", pu.Host)
_ = conn.SetReadDeadline(time.Now().Add(3 * time.Second))
br := bufio.NewReader(conn)
resp, err := http.ReadResponse(br, nil)
if err != nil {
t.Fatal(err)
}
clientBody, _ := io.ReadAll(resp.Body)
// 等待 proxy.ServeHTTP 返回并把 entry 提交
deadline := time.Now().Add(2 * time.Second)
for sub.Len() == 0 && time.Now().Before(deadline) {
time.Sleep(10 * time.Millisecond)
}
if sub.Len() == 0 {
t.Fatal("no log entry captured")
}
e := sub.Entry(0)
t.Logf("\n========== 客户端收到的字节 ==========\n%s", clientBody)
t.Logf("\n========== 数据库 response_body 字段（按字节原样存储）==========\n%s", e.ResponseBody)
t.Logf("\n========== 元数据 ==========")
t.Logf("is_stream = %v", e.IsStream)
t.Logf("status_code = %d", e.StatusCode)
t.Logf("len(body) = %d bytes", len(e.ResponseBody))
t.Logf("response_truncated = %v", e.ResponseTruncated)
if !strings.Contains(string(clientBody), `"content":"你"`) ||
!strings.Contains(string(clientBody), `"content":"好"`) ||
!strings.Contains(string(clientBody), "[DONE]") {
t.Errorf("client response body 缺少预期 chunk 内容")
}
var captured struct {
Choices []struct {
Message struct {
Role string `json:"role"`
Content string `json:"content"`
} `json:"message"`
FinishReason string `json:"finish_reason"`
} `json:"choices"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("stream response_body should be assembled JSON: %v; body=%q", err, e.ResponseBody)
}
if len(captured.Choices) != 1 {
t.Fatalf("assembled JSON choices length=%d, want 1", len(captured.Choices))
}
if captured.Choices[0].Message.Role != "assistant" {
t.Errorf("assembled role=%q, want assistant", captured.Choices[0].Message.Role)
}
if captured.Choices[0].Message.Content != "你好" {
t.Errorf("assembled content=%q, want 你好", captured.Choices[0].Message.Content)
}
if captured.Choices[0].FinishReason != "stop" {
t.Errorf("assembled finish_reason=%q, want stop", captured.Choices[0].FinishReason)
}
if !e.IsStream {
t.Errorf("is_stream 应为 true")
}
}
func TestOpenAIStreamPreservesReasoningAndMetadata(t *testing.T) {
upstream := newSSEUpstream([]string{
`data: {"id":"resp-1","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini-2026-03-17","choices":[{"index":0,"delta":{"role":"assistant","reasoning_content":"think "},"finish_reason":null,"native_finish_reason":null}]}` + "\n\n",
`data: {"id":"resp-1","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini-2026-03-17","choices":[{"index":0,"delta":{"reasoning_content":"hard"},"finish_reason":null,"native_finish_reason":null}]}` + "\n\n",
`data: {"id":"resp-1","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini-2026-03-17","choices":[{"index":0,"delta":{"content":"final"},"finish_reason":null,"native_finish_reason":null}]}` + "\n\n",
`data: {"id":"resp-1","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini-2026-03-17","choices":[{"index":0,"delta":{},"finish_reason":"stop","native_finish_reason":"stop"}]}` + "\n\n",
"data: [DONE]\n\n",
})
e := requestStreamEntry(t, upstream, "/v1/chat/completions")
var captured struct {
ID string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
Model string `json:"model"`
Choices []struct {
Index int `json:"index"`
Message struct {
Role string `json:"role"`
Content string `json:"content"`
ReasoningContent string `json:"reasoning_content"`
} `json:"message"`
FinishReason string `json:"finish_reason"`
NativeFinishReason string `json:"native_finish_reason"`
} `json:"choices"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("openai stream response_body should be JSON: %v; body=%q", err, e.ResponseBody)
}
if captured.ID != "resp-1" || captured.Object != "chat.completion" || captured.Created != 1779335544 || captured.Model != "gpt-5.4-mini-2026-03-17" {
t.Fatalf("unexpected metadata: %+v", captured)
}
if len(captured.Choices) != 1 {
t.Fatalf("choices length=%d, want 1", len(captured.Choices))
}
choice := captured.Choices[0]
if choice.Message.Role != "assistant" || choice.Message.Content != "final" || choice.Message.ReasoningContent != "think hard" {
t.Fatalf("unexpected message: %+v", choice.Message)
}
if choice.FinishReason != "stop" || choice.NativeFinishReason != "stop" {
t.Fatalf("unexpected finish reasons: %+v", choice)
}
}
func TestOpenAIStreamDoneDoesNotSubmitBeforeUpstreamCloses(t *testing.T) {
release := make(chan struct{})
upstream := fakeOpenAIStreamUpstreamThatStaysOpen(release)
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.New(u, filter, sub, 1024*1024)
proxySrv := httptest.NewServer(h)
released := false
defer func() {
if !released {
close(release)
}
proxySrv.Close()
upstream.Close()
}()
clientDone := make(chan error, 1)
go func() {
resp, err := http.Post(proxySrv.URL+"/v1/chat/completions", "application/json", strings.NewReader(`{"stream":true}`))
if err != nil {
clientDone <- err
return
}
_, _ = io.ReadAll(resp.Body)
_ = resp.Body.Close()
clientDone <- nil
}()
time.Sleep(100 * time.Millisecond)
if sub.Len() != 0 {
t.Fatal("terminal SSE event must not submit before ReverseProxy returns")
}
close(release)
released = true
select {
case err := <-clientDone:
if err != nil {
t.Fatal(err)
}
if sub.Len() != 1 {
t.Fatalf("stream should be submitted once, got %d entries", sub.Len())
}
case <-time.After(2 * time.Second):
t.Fatal("client did not finish after upstream closed")
}
}
func TestOpenAIStreamPreservesUsageChunk(t *testing.T) {
upstream := newSSEUpstream([]string{
`data: {"id":"resp-usage","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","system_fingerprint":"fp_123","choices":[{"index":0,"delta":{"role":"assistant","content":"ok"},"finish_reason":null}]}` + "\n\n",
`data: {"id":"resp-usage","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","system_fingerprint":"fp_123","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":null}` + "\n\n",
`data: {"id":"resp-usage","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","system_fingerprint":"fp_123","choices":[],"usage":{"prompt_tokens":5,"completion_tokens":2,"total_tokens":7}}` + "\n\n",
"data: [DONE]\n\n",
})
e := requestStreamEntry(t, upstream, "/v1/chat/completions")
var captured struct {
ID string `json:"id"`
SystemFingerprint string `json:"system_fingerprint"`
Choices []struct {
Message struct {
Content string `json:"content"`
} `json:"message"`
} `json:"choices"`
Usage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
} `json:"usage"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("openai stream with usage should be assembled JSON: %v; body=%q", err, e.ResponseBody)
}
if captured.ID != "resp-usage" || captured.SystemFingerprint != "fp_123" {
t.Fatalf("metadata not preserved: %+v", captured)
}
if len(captured.Choices) != 1 || captured.Choices[0].Message.Content != "ok" {
t.Fatalf("choices not assembled: %+v", captured.Choices)
}
if captured.Usage.PromptTokens != 5 || captured.Usage.CompletionTokens != 2 || captured.Usage.TotalTokens != 7 {
t.Fatalf("usage not preserved: %+v", captured.Usage)
}
}
func TestOpenAIStreamAssemblesToolCalls(t *testing.T) {
upstream := newSSEUpstream([]string{
`data: {"id":"resp-tools","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"delta":{"role":"assistant","tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"lookup","arguments":"{\"q\":"}}]},"finish_reason":null}]}` + "\n\n",
`data: {"id":"resp-tools","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"weather\"}"}}]},"finish_reason":null}]}` + "\n\n",
`data: {"id":"resp-tools","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}` + "\n\n",
"data: [DONE]\n\n",
})
e := requestStreamEntry(t, upstream, "/v1/chat/completions")
var captured struct {
Choices []struct {
Message struct {
Role string `json:"role"`
Content string `json:"content"`
ToolCalls []struct {
ID string `json:"id"`
Type string `json:"type"`
Function struct {
Name string `json:"name"`
Arguments string `json:"arguments"`
} `json:"function"`
} `json:"tool_calls"`
} `json:"message"`
FinishReason string `json:"finish_reason"`
} `json:"choices"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("openai tool stream should be assembled JSON: %v; body=%q", err, e.ResponseBody)
}
if len(captured.Choices) != 1 || captured.Choices[0].FinishReason != "tool_calls" {
t.Fatalf("unexpected choices: %+v", captured.Choices)
}
message := captured.Choices[0].Message
if message.Role != "assistant" || message.Content != "" {
t.Fatalf("unexpected message basics: %+v", message)
}
if len(message.ToolCalls) != 1 {
t.Fatalf("tool_calls length=%d, want 1; body=%s", len(message.ToolCalls), e.ResponseBody)
}
tool := message.ToolCalls[0]
if tool.ID != "call_1" || tool.Type != "function" || tool.Function.Name != "lookup" || tool.Function.Arguments != `{"q":"weather"}` {
t.Fatalf("unexpected tool call: %+v", tool)
}
}
func TestOpenAIResponsesStreamAssemblesCompletedResponse(t *testing.T) {
upstream := newSSEUpstream([]string{
`event: response.created` + "\n" + `data: {"type":"response.created","response":{"id":"resp-1","object":"response","status":"in_progress","model":"gpt-5.4-mini","output":[]}}` + "\n\n",
`event: response.output_text.delta` + "\n" + `data: {"type":"response.output_text.delta","item_id":"msg-1","output_index":0,"content_index":0,"delta":"hello"}` + "\n\n",
`event: response.completed` + "\n" + `data: {"type":"response.completed","response":{"id":"resp-1","object":"response","status":"completed","model":"gpt-5.4-mini","output":[{"id":"msg-1","type":"message","role":"assistant","content":[{"type":"output_text","text":"hello"}]}],"usage":{"input_tokens":3,"output_tokens":1,"total_tokens":4}}}` + "\n\n",
})
e := requestStreamEntry(t, upstream, "/v1/responses")
var captured struct {
ID string `json:"id"`
Object string `json:"object"`
Status string `json:"status"`
Model string `json:"model"`
Output []struct {
Type string `json:"type"`
Role string `json:"role"`
Content []struct {
Type string `json:"type"`
Text string `json:"text"`
} `json:"content"`
} `json:"output"`
Usage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
TotalTokens int `json:"total_tokens"`
} `json:"usage"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("responses stream should be completed response JSON: %v; body=%q", err, e.ResponseBody)
}
if captured.ID != "resp-1" || captured.Object != "response" || captured.Status != "completed" || captured.Model != "gpt-5.4-mini" {
t.Fatalf("unexpected response metadata: %+v", captured)
}
if len(captured.Output) != 1 || len(captured.Output[0].Content) != 1 || captured.Output[0].Content[0].Text != "hello" {
t.Fatalf("unexpected response output: %+v", captured.Output)
}
if captured.Usage.TotalTokens != 4 {
t.Fatalf("usage not preserved: %+v", captured.Usage)
}
}
func TestOpenAICompletionsStreamAssemblesText(t *testing.T) {
upstream := newSSEUpstream([]string{
`data: {"id":"cmpl-1","object":"text_completion","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"text":"hello","finish_reason":null}]}` + "\n\n",
`data: {"id":"cmpl-1","object":"text_completion","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"text":" world","finish_reason":null}]}` + "\n\n",
`data: {"id":"cmpl-1","object":"text_completion","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"text":"","finish_reason":"stop"}],"usage":{"prompt_tokens":2,"completion_tokens":2,"total_tokens":4}}` + "\n\n",
"data: [DONE]\n\n",
})
e := requestStreamEntry(t, upstream, "/v1/completions")
var captured struct {
ID string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
Model string `json:"model"`
Choices []struct {
Index int `json:"index"`
Text string `json:"text"`
FinishReason string `json:"finish_reason"`
} `json:"choices"`
Usage struct {
TotalTokens int `json:"total_tokens"`
} `json:"usage"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("completions stream should be assembled JSON: %v; body=%q", err, e.ResponseBody)
}
if captured.ID != "cmpl-1" || captured.Object != "text_completion" || captured.Model != "gpt-5.4-mini" || captured.Created != 1779335544 {
t.Fatalf("unexpected metadata: %+v", captured)
}
if len(captured.Choices) != 1 || captured.Choices[0].Text != "hello world" || captured.Choices[0].FinishReason != "stop" {
t.Fatalf("unexpected choices: %+v", captured.Choices)
}
if captured.Usage.TotalTokens != 4 {
t.Fatalf("usage not preserved: %+v", captured.Usage)
}
}
func TestAnthropicSSEAssemblesNativeMessageJSON(t *testing.T) {
e := requestStreamEntry(t, fakeAnthropicStreamUpstream(), "/v1/messages")
var captured struct {
ID string `json:"id"`
Type string `json:"type"`
Role string `json:"role"`
Content []struct {
Type string `json:"type"`
Text string `json:"text"`
} `json:"content"`
StopReason string `json:"stop_reason"`
Usage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
} `json:"usage"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("anthropic stream response_body should be native JSON: %v; body=%q", err, e.ResponseBody)
}
if captured.ID != "msg-1" || captured.Type != "message" || captured.Role != "assistant" {
t.Fatalf("unexpected anthropic message metadata: %+v", captured)
}
if len(captured.Content) != 1 || captured.Content[0].Type != "text" || captured.Content[0].Text != "你好" {
t.Fatalf("unexpected anthropic content: %+v", captured.Content)
}
if captured.StopReason != "end_turn" {
t.Errorf("stop_reason=%q, want end_turn", captured.StopReason)
}
if captured.Usage.InputTokens != 10 || captured.Usage.OutputTokens != 3 {
t.Errorf("usage=%+v, want input=10 output=3", captured.Usage)
}
}
func TestAnthropicSSEAssemblesToolUseContent(t *testing.T) {
upstream := newSSEUpstream([]string{
`event: message_start` + "\n" + `data: {"type":"message_start","message":{"id":"msg-tool","type":"message","role":"assistant","content":[],"model":"claude-3-5-sonnet","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}}` + "\n\n",
`event: content_block_start` + "\n" + `data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu_1","name":"lookup","input":{}}}` + "\n\n",
`event: content_block_delta` + "\n" + `data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"q\":"}}` + "\n\n",
`event: content_block_delta` + "\n" + `data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"\"weather\"}"}}` + "\n\n",
`event: message_delta` + "\n" + `data: {"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":8}}` + "\n\n",
`event: message_stop` + "\n" + `data: {"type":"message_stop"}` + "\n\n",
})
e := requestStreamEntry(t, upstream, "/v1/messages")
var captured struct {
Content []struct {
Type string `json:"type"`
ID string `json:"id"`
Name string `json:"name"`
Input map[string]any `json:"input"`
} `json:"content"`
StopReason string `json:"stop_reason"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("anthropic tool stream should be JSON: %v; body=%q", err, e.ResponseBody)
}
if len(captured.Content) != 1 {
t.Fatalf("content length=%d, want 1; body=%s", len(captured.Content), e.ResponseBody)
}
tool := captured.Content[0]
if tool.Type != "tool_use" || tool.ID != "toolu_1" || tool.Name != "lookup" || tool.Input["q"] != "weather" {
t.Fatalf("unexpected tool content: %+v", tool)
}
if captured.StopReason != "tool_use" {
t.Fatalf("stop_reason=%q, want tool_use", captured.StopReason)
}
}
func TestGeminiSSEAssemblesNativeGenerateContentJSON(t *testing.T) {
e := requestStreamEntry(t, fakeGeminiStreamUpstream(), "/v1beta/models/gemini-1.5-pro:generateContent")
var captured struct {
Candidates []struct {
Content struct {
Role string `json:"role"`
Parts []struct {
Text string `json:"text"`
} `json:"parts"`
} `json:"content"`
FinishReason string `json:"finishReason"`
Index int `json:"index"`
} `json:"candidates"`
UsageMetadata struct {
PromptTokenCount int `json:"promptTokenCount"`
CandidatesTokenCount int `json:"candidatesTokenCount"`
TotalTokenCount int `json:"totalTokenCount"`
} `json:"usageMetadata"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("gemini stream response_body should be native JSON: %v; body=%q", err, e.ResponseBody)
}
if len(captured.Candidates) != 1 {
t.Fatalf("candidates length=%d, want 1", len(captured.Candidates))
}
candidate := captured.Candidates[0]
if candidate.Content.Role != "model" || len(candidate.Content.Parts) != 1 || candidate.Content.Parts[0].Text != "你好" {
t.Fatalf("unexpected gemini content: %+v", candidate.Content)
}
if candidate.FinishReason != "STOP" {
t.Errorf("finishReason=%q, want STOP", candidate.FinishReason)
}
if captured.UsageMetadata.TotalTokenCount != 4 {
t.Errorf("usageMetadata=%+v, want totalTokenCount=4", captured.UsageMetadata)
}
}
func TestGeminiStreamPreservesSafetyRatingsAndUsageMetadata(t *testing.T) {
upstream := newSSEUpstream([]string{
`data: {"candidates":[{"content":{"parts":[{"text":"你"}],"role":"model"},"finishReason":null,"index":0,"safetyRatings":[{"category":"HARM_CATEGORY_HARASSMENT","probability":"NEGLIGIBLE"}]}],"usageMetadata":{"promptTokenCount":8,"toolUsePromptTokenCount":0,"candidatesTokenCount":0,"totalTokenCount":8,"thoughtsTokenCount":10}}` + "\n\n",
`data: {"candidates":[{"content":{"parts":[{"text":"好"}],"role":"model"},"finishReason":"STOP","index":0,"safetyRatings":[{"category":"HARM_CATEGORY_HARASSMENT","probability":"NEGLIGIBLE"}]}],"usageMetadata":{"promptTokenCount":8,"toolUsePromptTokenCount":0,"candidatesTokenCount":2,"totalTokenCount":20,"thoughtsTokenCount":10}}` + "\n\n",
})
e := requestStreamEntry(t, upstream, "/v1beta/models/gemini-1.5-pro:streamGenerateContent")
var captured struct {
Candidates []struct {
Content struct {
Parts []struct {
Text string `json:"text"`
} `json:"parts"`
} `json:"content"`
SafetyRatings []struct {
Category string `json:"category"`
Probability string `json:"probability"`
} `json:"safetyRatings"`
} `json:"candidates"`
UsageMetadata struct {
PromptTokenCount int `json:"promptTokenCount"`
ToolUsePromptTokenCount int `json:"toolUsePromptTokenCount"`
CandidatesTokenCount int `json:"candidatesTokenCount"`
TotalTokenCount int `json:"totalTokenCount"`
ThoughtsTokenCount int `json:"thoughtsTokenCount"`
} `json:"usageMetadata"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("gemini stream response_body should be JSON: %v; body=%q", err, e.ResponseBody)
}
if len(captured.Candidates) != 1 || len(captured.Candidates[0].SafetyRatings) != 1 {
t.Fatalf("expected safetyRatings to be preserved, got %+v", captured.Candidates)
}
if captured.Candidates[0].SafetyRatings[0].Category != "HARM_CATEGORY_HARASSMENT" {
t.Fatalf("unexpected safetyRatings: %+v", captured.Candidates[0].SafetyRatings)
}
if captured.UsageMetadata.ToolUsePromptTokenCount != 0 || captured.UsageMetadata.ThoughtsTokenCount != 10 || captured.UsageMetadata.TotalTokenCount != 20 {
t.Fatalf("usageMetadata fields not preserved: %+v", captured.UsageMetadata)
}
}
func TestGeminiStreamPreservesFunctionCallParts(t *testing.T) {
upstream := newSSEUpstream([]string{
`data: {"candidates":[{"content":{"parts":[{"functionCall":{"name":"lookup","args":{"q":"weather"}}}],"role":"model"},"finishReason":"STOP","index":0}],"usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":2,"totalTokenCount":4}}` + "\n\n",
})
e := requestStreamEntry(t, upstream, "/v1beta/models/gemini-1.5-pro:streamGenerateContent")
var captured struct {
Candidates []struct {
Content struct {
Parts []struct {
FunctionCall struct {
Name string `json:"name"`
Args map[string]any `json:"args"`
} `json:"functionCall"`
} `json:"parts"`
} `json:"content"`
} `json:"candidates"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("gemini functionCall stream should be JSON: %v; body=%q", err, e.ResponseBody)
}
call := captured.Candidates[0].Content.Parts[0].FunctionCall
if call.Name != "lookup" || call.Args["q"] != "weather" {
t.Fatalf("functionCall not preserved: %+v; body=%s", call, e.ResponseBody)
}
}
func TestUnknownSSEKeepsRawBody(t *testing.T) {
e := requestStreamEntry(t, fakeUnknownStreamUpstream(), "/v1/chat/completions")
if string(e.ResponseBody) != "event: custom\ndata: not-json\n\n" {
t.Fatalf("unknown stream should keep raw body, got %q", e.ResponseBody)
}
}
func TestTruncatedSSEKeepsCapturedRawBody(t *testing.T) {
upstream := fakeOpenAIStreamUpstream()
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.New(u, filter, sub, 520)
proxySrv := httptest.NewServer(h)
defer proxySrv.Close()
resp, err := http.Post(proxySrv.URL+"/v1/chat/completions", "application/json", nil)
if err != nil {
t.Fatal(err)
}
_, _ = io.ReadAll(resp.Body)
_ = resp.Body.Close()
deadline := time.Now().Add(2 * time.Second)
for sub.Len() == 0 && time.Now().Before(deadline) {
time.Sleep(10 * time.Millisecond)
}
if sub.Len() == 0 {
t.Fatal("no log entry captured")
}
e := sub.Entry(0)
if !e.ResponseTruncated {
t.Fatal("response should be marked truncated")
}
if !strings.HasPrefix(string(e.ResponseBody), "data: ") {
t.Fatalf("truncated stream should keep captured raw body, got %q", e.ResponseBody)
}
if json.Valid(e.ResponseBody) {
t.Fatalf("truncated stream should not be assembled as JSON, got %q", e.ResponseBody)
}
}
func TestMultimodalStreamRequestBodyIsCaptured(t *testing.T) {
upstream := newSSEUpstream([]string{
`data: {"id":"chatcmpl-image","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}` + "\n\n",
`data: {"id":"chatcmpl-image","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":null}]}` + "\n\n",
`data: {"id":"chatcmpl-image","object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}` + "\n\n",
"data: [DONE]\n\n",
})
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.New(u, filter, sub, 1024*1024)
proxySrv := httptest.NewServer(h)
defer proxySrv.Close()
reqBody := `{"model":"gpt-5.4-mini","messages":[{"role":"user","content":[{"type":"text","text":"describe"},{"type":"image_url","image_url":{"url":"data:image/png;base64,AAAA"}}]}],"stream":true}`
resp, err := http.Post(proxySrv.URL+"/v1/chat/completions", "application/json", strings.NewReader(reqBody))
if err != nil {
t.Fatal(err)
}
_, _ = io.ReadAll(resp.Body)
_ = resp.Body.Close()
deadline := time.Now().Add(2 * time.Second)
for sub.Len() == 0 && time.Now().Before(deadline) {
time.Sleep(10 * time.Millisecond)
}
if sub.Len() == 0 {
t.Fatal("no log entry captured")
}
e := sub.Entry(0)
if e.RequestTruncated {
t.Fatal("multimodal request should not be truncated")
}
requestText := string(e.RequestBody)
if !strings.Contains(requestText, `"image_url"`) || !strings.Contains(requestText, `data:image/png;base64,AAAA`) {
t.Fatalf("request_body should contain image input, got %q", requestText)
}
if !e.IsStream {
t.Fatal("response should still be marked stream")
}
}
func TestChunkedMultimodalRequestBodyIsCaptured(t *testing.T) {
upstream := newSSEUpstream([]string{
`data: {"id":"chatcmpl-image","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":"stop"}]}` + "\n\n",
"data: [DONE]\n\n",
})
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.New(u, filter, sub, 1024*1024)
proxySrv := httptest.NewServer(h)
defer proxySrv.Close()
reqBody := `{"model":"gpt-5.4-mini","messages":[{"role":"user","content":[{"type":"text","text":"describe"},{"type":"image_url","image_url":{"url":"data:image/png;base64,AAAA"}}]}],"stream":true}`
req, err := http.NewRequest(http.MethodPost, proxySrv.URL+"/v1/chat/completions", strings.NewReader(reqBody))
if err != nil {
t.Fatal(err)
}
req.ContentLength = -1
req.Header.Set("Content-Type", "application/json")
req.TransferEncoding = []string{"chunked"}
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
_, _ = io.ReadAll(resp.Body)
_ = resp.Body.Close()
deadline := time.Now().Add(2 * time.Second)
for sub.Len() == 0 && time.Now().Before(deadline) {
time.Sleep(10 * time.Millisecond)
}
if sub.Len() == 0 {
t.Fatal("no log entry captured")
}
e := sub.Entry(0)
requestText := string(e.RequestBody)
if !strings.Contains(requestText, `"image_url"`) || !strings.Contains(requestText, `data:image/png;base64,AAAA`) {
t.Fatalf("chunked request_body should contain image input, got %q", requestText)
}
}
func requestStreamEntry(t *testing.T, upstream *httptest.Server, path string) *logger.LogEntry {
t.Helper()
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.New(u, filter, sub, 1024*1024)
proxySrv := httptest.NewServer(h)
defer proxySrv.Close()
resp, err := http.Post(proxySrv.URL+path, "application/json", nil)
if err != nil {
t.Fatal(err)
}
_, _ = io.ReadAll(resp.Body)
_ = resp.Body.Close()
deadline := time.Now().Add(2 * time.Second)
for sub.Len() == 0 && time.Now().Before(deadline) {
time.Sleep(10 * time.Millisecond)
}
if sub.Len() == 0 {
t.Fatal("no log entry captured")
}
return sub.Entry(0)
}
tests/proxy/websocket_test.go
+299
View File
@@ -0,0 +1,299 @@
package proxy_test
import (
"bufio"
"context"
"io"
"net"
"net/http"
"net/http/httptest"
"net/url"
"strings"
"testing"
"time"
"git.misaka.ren/M1saka/token_thief/config"
"git.misaka.ren/M1saka/token_thief/logger"
"git.misaka.ren/M1saka/token_thief/proxy"
)
type noopSubmitter struct{}
func (noopSubmitter) Submit(*logger.LogEntry) {}
type chanSubmitter chan *logger.LogEntry
func (c chanSubmitter) Submit(e *logger.LogEntry) { c <- e }
// fakeUpstream 模拟一个最简 WebSocket 升级：
// 收到 GET + Upgrade: websocket 后回 101，然后做字节回声直到对端关闭。
func fakeUpstream(t *testing.T) *httptest.Server {
t.Helper()
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if strings.ToLower(r.Header.Get("Upgrade")) != "websocket" {
http.Error(w, "expected websocket upgrade", http.StatusBadRequest)
return
}
hj, ok := w.(http.Hijacker)
if !ok {
http.Error(w, "no hijack", http.StatusInternalServerError)
return
}
conn, brw, err := hj.Hijack()
if err != nil {
t.Errorf("upstream hijack: %v", err)
return
}
defer conn.Close()
// 直接回 101 握手响应（简化版，不做真正 Sec-WebSocket-Accept 计算）
_, _ = brw.WriteString("HTTP/1.1 101 Switching Protocols\r\n" +
"Upgrade: websocket\r\n" +
"Connection: Upgrade\r\n" +
"\r\n")
_ = brw.Flush()
// echo
buf := make([]byte, 1024)
for {
n, err := conn.Read(buf)
if err != nil {
return
}
if _, err := conn.Write(buf[:n]); err != nil {
return
}
}
}))
return srv
}
func TestWebSocketHandshakeCapturedAndShutdownClosesConnection(t *testing.T) {
upstream := fakeUpstream(t)
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
entries := make(chanSubmitter, 1)
h := proxy.New(u, filter, entries, 1024)
proxySrv := httptest.NewServer(h)
defer proxySrv.Close()
pu, _ := url.Parse(proxySrv.URL)
conn, err := net.DialTimeout("tcp", pu.Host, time.Second)
if err != nil {
t.Fatal(err)
}
defer conn.Close()
_, _ = io.WriteString(conn, "GET /v1/realtime HTTP/1.1\r\nHost: "+pu.Host+"\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n")
br := bufio.NewReader(conn)
resp, err := http.ReadResponse(br, nil)
if err != nil {
t.Fatal(err)
}
if resp.StatusCode != http.StatusSwitchingProtocols {
t.Fatalf("status=%d", resp.StatusCode)
}
select {
case entry := <-entries:
if entry.StatusCode != http.StatusSwitchingProtocols || len(entry.ResponseBody) != 0 {
t.Fatalf("handshake entry status=%d body=%q", entry.StatusCode, entry.ResponseBody)
}
case <-time.After(time.Second):
t.Fatal("websocket handshake was not captured")
}
payload := []byte("frame-data-must-not-be-logged")
if _, err := conn.Write(payload); err != nil {
t.Fatal(err)
}
echo := make([]byte, len(payload))
if _, err := io.ReadFull(br, echo); err != nil {
t.Fatalf("read websocket payload: %v", err)
}
if string(echo) != string(payload) {
t.Fatalf("echo mismatch: got %q want %q", echo, payload)
}
select {
case entry := <-entries:
t.Fatalf("websocket frame produced an extra log entry: %+v", entry)
case <-time.After(50 * time.Millisecond):
}
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
if err := h.Shutdown(ctx); err != nil {
t.Fatal(err)
}
_ = conn.SetReadDeadline(time.Now().Add(time.Second))
if _, err := conn.Read(make([]byte, 1)); err == nil {
t.Fatal("connection remains open after Shutdown")
}
if err := h.Shutdown(ctx); err != nil {
t.Fatalf("second Shutdown: %v", err)
}
}
func TestWebSocketNaturalCloseUnregistersConnection(t *testing.T) {
upstream := fakeUpstream(t)
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.New(u, filter, noopSubmitter{}, 1024)
proxySrv := httptest.NewServer(h)
defer proxySrv.Close()
pu, _ := url.Parse(proxySrv.URL)
conn, err := net.DialTimeout("tcp", pu.Host, time.Second)
if err != nil {
t.Fatal(err)
}
_, _ = io.WriteString(conn, "GET /v1/realtime HTTP/1.1\r\nHost: "+pu.Host+"\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n")
resp, err := http.ReadResponse(bufio.NewReader(conn), nil)
if err != nil {
t.Fatal(err)
}
if resp.StatusCode != http.StatusSwitchingProtocols {
t.Fatalf("status=%d", resp.StatusCode)
}
if err := conn.Close(); err != nil {
t.Fatal(err)
}
deadline := time.Now().Add(time.Second)
for {
ctx, cancel := context.WithCancel(context.Background())
cancel()
err := h.Shutdown(ctx)
if err == nil {
break
}
if time.Now().After(deadline) {
t.Fatalf("connection was not unregistered after natural close: %v", err)
}
time.Sleep(10 * time.Millisecond)
}
}
func TestNonWebSocketUpgradeIsManagedButNotLogged(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
hj := w.(http.Hijacker)
conn, brw, err := hj.Hijack()
if err != nil {
t.Errorf("upstream hijack: %v", err)
return
}
defer conn.Close()
_, _ = brw.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: test-protocol\r\nConnection: Upgrade\r\n\r\n")
_ = brw.Flush()
_, _ = io.Copy(io.Discard, conn)
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
entries := make(chanSubmitter, 1)
h := proxy.New(u, filter, entries, 1024)
proxySrv := httptest.NewServer(h)
defer proxySrv.Close()
pu, _ := url.Parse(proxySrv.URL)
conn, err := net.DialTimeout("tcp", pu.Host, time.Second)
if err != nil {
t.Fatal(err)
}
defer conn.Close()
_, _ = io.WriteString(conn, "GET /upgrade HTTP/1.1\r\nHost: "+pu.Host+"\r\nUpgrade: test-protocol\r\nConnection: Upgrade\r\n\r\n")
resp, err := http.ReadResponse(bufio.NewReader(conn), nil)
if err != nil {
t.Fatal(err)
}
if resp.StatusCode != http.StatusSwitchingProtocols {
t.Fatalf("status=%d", resp.StatusCode)
}
select {
case entry := <-entries:
t.Fatalf("non-WebSocket upgrade was logged: %+v", entry)
case <-time.After(50 * time.Millisecond):
}
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
if err := h.Shutdown(ctx); err != nil {
t.Fatal(err)
}
}
// TestWebSocketProxyPassthrough 验证 Upgrade 请求能正确透传，
// 确认 captureWriter 的 Hijacker 实现没破坏 ReverseProxy 的 WS 行为。
func TestWebSocketProxyPassthrough(t *testing.T) {
upstream := fakeUpstream(t)
defer upstream.Close()
u, err := url.Parse(upstream.URL)
if err != nil {
t.Fatal(err)
}
// disabled 模式下 ShouldLog 返回 true（全量记录），会进入捕获分支。
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.New(u, filter, noopSubmitter{}, 1024)
proxySrv := httptest.NewServer(h)
defer proxySrv.Close()
// 建立到代理的 TCP 连接，手写 Upgrade 请求
pu, _ := url.Parse(proxySrv.URL)
d := net.Dialer{Timeout: 3 * time.Second}
conn, err := d.DialContext(context.Background(), "tcp", pu.Host)
if err != nil {
t.Fatal(err)
}
defer conn.Close()
req := "GET /v1/realtime HTTP/1.1\r\n" +
"Host: " + pu.Host + "\r\n" +
"Upgrade: websocket\r\n" +
"Connection: Upgrade\r\n" +
"Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n" +
"Sec-WebSocket-Version: 13\r\n" +
"\r\n"
if _, err := io.WriteString(conn, req); err != nil {
t.Fatal(err)
}
_ = conn.SetReadDeadline(time.Now().Add(3 * time.Second))
br := bufio.NewReader(conn)
resp, err := http.ReadResponse(br, nil)
if err != nil {
t.Fatalf("read upgrade response: %v", err)
}
if resp.StatusCode != http.StatusSwitchingProtocols {
t.Fatalf("expected 101, got %d", resp.StatusCode)
}
if !strings.EqualFold(resp.Header.Get("Upgrade"), "websocket") {
t.Fatalf("expected Upgrade: websocket, got %q", resp.Header.Get("Upgrade"))
}
// echo 测试
payload := "hello-websocket"
if _, err := io.WriteString(conn, payload); err != nil {
t.Fatal(err)
}
got := make([]byte, len(payload))
if _, err := io.ReadFull(br, got); err != nil {
t.Fatalf("read echo: %v", err)
}
if string(got) != payload {
t.Fatalf("echo mismatch: got %q want %q", got, payload)
}
}
tests/scripts/dbcheck/main.go
+109
View File
@@ -0,0 +1,109 @@
// Standalone helper for smoke.ps1. It reads CHECK_RIDS (JSON array) and verifies rows in proxy_logs.
package main
import (
"context"
"encoding/json"
"fmt"
"log"
"os"
"time"
"github.com/ClickHouse/clickhouse-go/v2"
"git.misaka.ren/M1saka/token_thief/db"
)
type row struct {
RequestID string
Method string
Path string
StatusCode int
RequestTruncated bool
ResponseTruncated bool
IsStream bool
LatencyMS int64
ReqBodyLen int64
RespBodyLen int64
ReqHeaders string
ErrorMsg string
}
func main() {
log.SetFlags(0)
dsn := os.Getenv("CLICKHOUSE_URL")
if dsn == "" {
log.Fatalf("CLICKHOUSE_URL not set")
}
ridsRaw := os.Getenv("CHECK_RIDS")
if ridsRaw == "" {
log.Fatalf("CHECK_RIDS not set")
}
var rids []string
if err := json.Unmarshal([]byte(ridsRaw), &rids); err != nil {
log.Fatalf("parse CHECK_RIDS: %v", err)
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
opts, err := db.ClickHouseOptions(dsn)
if err != nil {
log.Fatalf("connect: %v", err)
}
conn, err := clickhouse.Open(opts)
if err != nil {
log.Fatalf("connect: %v", err)
}
defer conn.Close()
const q = `
SELECT request_id, method, path, status_code,
request_truncated, response_truncated, is_stream, latency_ms,
toInt64(length(request_body)),
toInt64(length(response_body)),
request_headers,
error
FROM proxy_logs WHERE request_id = ?
ORDER BY started_at DESC LIMIT 1`
ok := 0
for _, rid := range rids {
var r row
err := conn.QueryRow(ctx, q, rid).Scan(
&r.RequestID, &r.Method, &r.Path, &r.StatusCode,
&r.RequestTruncated, &r.ResponseTruncated, &r.IsStream, &r.LatencyMS,
&r.ReqBodyLen, &r.RespBodyLen, &r.ReqHeaders, &r.ErrorMsg,
)
if err != nil {
fmt.Printf(" FAIL rid=%s: not found in db (%v)\n", rid, err)
continue
}
ok++
fmt.Printf(" OK rid=%s\n", rid)
fmt.Printf(" method=%s path=%s status=%d latency_ms=%d\n", r.Method, r.Path, r.StatusCode, r.LatencyMS)
fmt.Printf(" req_body=%d B (truncated=%v) resp_body=%d B (truncated=%v) is_stream=%v\n",
r.ReqBodyLen, r.RequestTruncated, r.RespBodyLen, r.ResponseTruncated, r.IsStream)
hasAuth := false
var hm map[string][]string
if err := json.Unmarshal([]byte(r.ReqHeaders), &hm); err == nil {
_, hasAuth = hm["Authorization"]
}
fmt.Printf(" headers has Authorization=%v\n", hasAuth)
if r.ErrorMsg != "" {
fmt.Printf(" error=%s\n", trunc(r.ErrorMsg, 200))
}
}
fmt.Printf("\n found %d / %d\n", ok, len(rids))
if ok != len(rids) {
os.Exit(1)
}
}
func trunc(s string, n int) string {
if len(s) <= n {
return s
}
return s[:n] + "..."
}
tests/scripts/dump-stream-body.ps1
+54
View File
@@ -0,0 +1,54 @@
param([string]$RequestID)
if (-not $RequestID) { throw "RequestID is required" }
. .\tests\scripts\load-env.ps1 | Out-Null
$env:DUMP_RID = $RequestID
@'
package main
import (
"context"
"fmt"
"log"
"os"
"time"
"github.com/ClickHouse/clickhouse-go/v2"
"git.misaka.ren/M1saka/token_thief/db"
)
func main() {
log.SetFlags(0)
dsn := os.Getenv("CLICKHOUSE_URL")
if dsn == "" {
log.Fatal("CLICKHOUSE_URL not set")
}
rid := os.Getenv("DUMP_RID")
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
opts, err := db.ClickHouseOptions(dsn)
if err != nil {
log.Fatal(err)
}
conn, err := clickhouse.Open(opts)
if err != nil {
log.Fatal(err)
}
defer conn.Close()
var body string
if err := conn.QueryRow(ctx,
`SELECT response_body FROM proxy_logs WHERE request_id = ? ORDER BY started_at DESC LIMIT 1`,
rid,
).Scan(&body); err != nil {
log.Fatal(err)
}
fmt.Println(body)
}
'@ | Set-Content -Path tmp_dump.go -Encoding UTF8 -NoNewline
go run tmp_dump.go
Remove-Item tmp_dump.go -Force
Remove-Item Env:DUMP_RID -ErrorAction SilentlyContinue
tests/scripts/load-env.ps1
+25
View File
@@ -0,0 +1,25 @@
# Load KEY=VALUE pairs from .env into the current PowerShell process.
# Usage: . .\tests\scripts\load-env.ps1
param(
[string]$Path = ".env"
)
if (-not (Test-Path $Path)) {
Write-Error "env file not found: $Path"
return
}
Get-Content $Path | ForEach-Object {
$line = $_.Trim()
if ($line -eq "" -or $line.StartsWith("#")) { return }
$idx = $line.IndexOf("=")
if ($idx -lt 1) { return }
$key = $line.Substring(0, $idx).Trim()
$val = $line.Substring($idx + 1).Trim()
if (($val.StartsWith('"') -and $val.EndsWith('"')) -or
($val.StartsWith("'") -and $val.EndsWith("'"))) {
$val = $val.Substring(1, $val.Length - 2)
}
[Environment]::SetEnvironmentVariable($key, $val, "Process")
Write-Host " loaded $key"
}
tests/scripts/smoke.ps1
+135
View File
@@ -0,0 +1,135 @@
# End-to-end smoke test:
# start proxy -> run chat cases -> wait batch flush -> verify ClickHouse -> stop proxy.
param(
[string]$Model = "gpt-5.4-mini",
[string]$ApiKey = $env:NEWAPI_KEY,
[string]$ProxyBase = "http://127.0.0.1:8080"
)
if (-not $ApiKey) { throw "NEWAPI_KEY not set" }
$ErrorActionPreference = "Stop"
function Section($name) {
Write-Host ""
Write-Host ("=" * 70) -ForegroundColor Cyan
Write-Host $name -ForegroundColor Cyan
Write-Host ("=" * 70) -ForegroundColor Cyan
}
Section "Start proxy"
if (Test-Path proxy.log) { Remove-Item proxy.log -Force }
if (Test-Path proxy.err.log) { Remove-Item proxy.err.log -Force }
$proxy = Start-Process -FilePath .\TokenThief.exe -PassThru -RedirectStandardOutput proxy.log -RedirectStandardError proxy.err.log -WindowStyle Hidden
Write-Host "proxy pid=$($proxy.Id)"
Start-Sleep -Seconds 2
$health = & curl.exe -s -o NUL -w "%{http_code}" "$ProxyBase/healthz"
Write-Host "healthz: $health"
if ($health -ne "200") {
Get-Content proxy.log -Tail 30 | Write-Host
Get-Content proxy.err.log -Tail 30 | Write-Host
throw "proxy did not start"
}
$results = @{}
function CallChat {
param([string]$Url, [string]$JsonBody)
$bodyTmp = New-TemporaryFile
$headTmp = New-TemporaryFile
$outTmp = New-TemporaryFile
# PowerShell 5.1 Set-Content -Encoding UTF8 writes a BOM; newapi rejects BOM-prefixed JSON.
$utf8NoBom = New-Object System.Text.UTF8Encoding($false)
[System.IO.File]::WriteAllText($bodyTmp.FullName, $JsonBody, $utf8NoBom)
$code = & curl.exe -s -X POST $Url `
-H "Content-Type: application/json" `
-H "Authorization: Bearer $ApiKey" `
-D $headTmp.FullName `
-o $outTmp.FullName `
--data-binary "@$($bodyTmp.FullName)" `
-w "%{http_code}"
$rid = ""
foreach ($line in Get-Content $headTmp.FullName) {
if ($line -match '^X-Request-Id:\s*(.+)$') {
$rid = $matches[1].Trim()
break
}
}
$body = Get-Content $outTmp.FullName -Raw -ErrorAction SilentlyContinue
if (-not $body) { $body = "" }
Remove-Item $bodyTmp.FullName, $headTmp.FullName, $outTmp.FullName -Force -ErrorAction SilentlyContinue
return @{ StatusCode = $code; RequestID = $rid; Body = $body }
}
try {
Section "T1: non-stream chat completions"
$t1Body = '{"model":"' + $Model + '","messages":[{"role":"user","content":"Say hello in one short sentence."}],"stream":false}'
$r = CallChat "$ProxyBase/v1/chat/completions" $t1Body
Write-Host " status=$($r.StatusCode) rid=$($r.RequestID)"
$preview = if ($r.Body.Length -gt 200) { $r.Body.Substring(0, 200) } else { $r.Body }
Write-Host " body(first 200): $preview"
$results.T1 = $r
Section "T2: stream chat completions"
$t2Body = '{"model":"' + $Model + '","messages":[{"role":"user","content":"Say exactly: one two three four five"}],"stream":true}'
$r = CallChat "$ProxyBase/v1/chat/completions" $t2Body
Write-Host " status=$($r.StatusCode) rid=$($r.RequestID)"
$chunkCount = ([regex]::Matches($r.Body, "^data:", "Multiline")).Count
$hasDone = $r.Body.Contains("[DONE]")
Write-Host " SSE chunk lines=$chunkCount contains [DONE]=$hasDone body_len=$($r.Body.Length)"
$results.T2 = $r
Section "T3: large body (request_truncated should be true)"
$bigContent = "x" * (12 * 1024 * 1024)
$t3Body = '{"model":"' + $Model + '","messages":[{"role":"user","content":"' + $bigContent + '"}],"stream":false,"max_tokens":5}'
Write-Host " request body size: $($t3Body.Length) bytes"
$r = CallChat "$ProxyBase/v1/chat/completions" $t3Body
Write-Host " status=$($r.StatusCode) rid=$($r.RequestID)"
$preview = if ($r.Body.Length -gt 200) { $r.Body.Substring(0, 200) } else { $r.Body }
Write-Host " body(first 200): $preview"
$results.T3 = $r
Section "T4: nonexistent model (upstream error response should be logged)"
$t4Body = '{"model":"definitely-not-a-real-model-xyz","messages":[{"role":"user","content":"hi"}]}'
$r = CallChat "$ProxyBase/v1/chat/completions" $t4Body
Write-Host " status=$($r.StatusCode) rid=$($r.RequestID)"
$preview = if ($r.Body.Length -gt 200) { $r.Body.Substring(0, 200) } else { $r.Body }
Write-Host " body(first 200): $preview"
$results.T4 = $r
Section "T5: healthz (should not be logged)"
$code = & curl.exe -s -o NUL -w "%{http_code}" "$ProxyBase/healthz"
Write-Host " /healthz status=$code"
Section "Wait batch flush (4s)"
Start-Sleep -Seconds 4
Section "T6: ClickHouse verification"
$rids = @()
foreach ($k in @("T1","T2","T3","T4")) {
if ($results[$k].RequestID) { $rids += $results[$k].RequestID }
}
Write-Host " request_ids: $($rids -join ', ')"
$env:CHECK_RIDS = ($rids | ConvertTo-Json -Compress)
if ($rids.Count -eq 1) { $env:CHECK_RIDS = "[`"$($rids[0])`"]" }
& go run .\tests\scripts\dbcheck\main.go
Remove-Item Env:CHECK_RIDS -ErrorAction SilentlyContinue
} finally {
Section "Stop proxy"
Stop-Process -Id $proxy.Id -Force
Write-Host "tail proxy.log:"
Get-Content proxy.log -Tail 30 | Write-Host
if (Test-Path proxy.err.log) {
$err = Get-Content proxy.err.log -ErrorAction SilentlyContinue
if ($err) {
Write-Host "proxy.err.log:"
$err | Write-Host
}
}
}