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/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)
}
}
}