package task import ( "context" "encoding/base64" "encoding/json" "errors" "fmt" "net" "net/http" "strings" "testing" "unicode/utf8" lcommon "github.com/qianfree/team-api/internal/logic/common" "github.com/qianfree/team-api/relay/dto" ) func TestMemResponseWriter(t *testing.T) { w := newMemResponseWriter() if w.StatusCode() != http.StatusOK { t.Fatalf("default status = %d, want 200", w.StatusCode()) } if _, err := w.Write([]byte("hello ")); err != nil { t.Fatalf("write: %v", err) } if _, err := w.Write([]byte("world")); err != nil { t.Fatalf("write: %v", err) } if got := string(w.Bytes()); got != "hello world" { t.Fatalf("body = %q, want %q", got, "hello world") } w.WriteHeader(http.StatusBadGateway) if w.StatusCode() != http.StatusBadGateway { t.Fatalf("status after WriteHeader = %d, want 502", w.StatusCode()) } w.Header().Set("Content-Type", "application/json") if w.Header().Get("Content-Type") != "application/json" { t.Fatalf("header not stored") } } func TestBuildImageResult_EmptyData(t *testing.T) { job := &SyncImageJob{PublicTaskID: "task_x"} body, _ := json.Marshal(map[string]any{"created": 1, "data": []any{}}) if _, _, err := buildImageResult(context.Background(), job, body); err == nil { t.Fatal("empty data should error") } } // TestExtractImageUsage_GptImage gpt-image 系列返回 usage,应被解析出 token 用量。 func TestExtractImageUsage_GptImage(t *testing.T) { body := []byte(`{"created":1,"data":[{"b64_json":"x"}],"usage":{"total_tokens":4160,"input_tokens":160,"output_tokens":4000}}`) prompt, completion, total := extractImageUsage(body) if prompt != 160 || completion != 4000 || total != 4160 { t.Fatalf("got prompt=%d completion=%d total=%d, want 160/4000/4160", prompt, completion, total) } } // TestExtractImageUsage_TotalFallback total 缺失时用 input+output 兜底。 func TestExtractImageUsage_TotalFallback(t *testing.T) { body := []byte(`{"data":[{"b64_json":"x"}],"usage":{"input_tokens":100,"output_tokens":250}}`) _, _, total := extractImageUsage(body) if total != 350 { t.Fatalf("total = %d, want 350 (input+output fallback)", total) } } // TestExtractImageUsage_NoUsage DALL·E 等无 usage 字段,返回全 0(保持按次计费预扣)。 func TestExtractImageUsage_NoUsage(t *testing.T) { body := []byte(`{"created":1,"data":[{"url":"https://example.com/a.png"}]}`) prompt, completion, total := extractImageUsage(body) if prompt != 0 || completion != 0 || total != 0 { t.Fatalf("got prompt=%d completion=%d total=%d, want all 0", prompt, completion, total) } } // TestExtractImageUsage_Malformed 响应体非法 JSON 时安全返回全 0,不 panic。 func TestExtractImageUsage_Malformed(t *testing.T) { if p, c, tot := extractImageUsage([]byte("not json")); p != 0 || c != 0 || tot != 0 { t.Fatalf("malformed body should yield 0/0/0, got %d/%d/%d", p, c, tot) } } func TestBuildImageResult_B64WithoutStorageErrors(t *testing.T) { // 对象存储不可用时,b64_json 无法 re-host → 报错。覆盖 accessor 直接返回错误, // 避免在无数据库的单测环境触达 sys_options 查询。 orig := acquireSyncImageFileSvc acquireSyncImageFileSvc = func(context.Context) (*lcommon.FileService, error) { return nil, lcommon.ErrStorageNotConfigured } defer func() { acquireSyncImageFileSvc = orig }() job := &SyncImageJob{PublicTaskID: "task_x", TenantID: 1, UserID: 2} raw, _ := json.Marshal(map[string]any{ "created": 1, "data": []any{map[string]any{"b64_json": base64.StdEncoding.EncodeToString([]byte("PNGDATA"))}}, }) if _, _, err := buildImageResult(context.Background(), job, raw); err == nil { t.Fatal("b64 without storage should error") } } // TestBuildImageResult_MultiURLPassthrough 多图 url 透传:全部图片进 normalized,首图作 ResultURL。 func TestBuildImageResult_MultiURLPassthrough(t *testing.T) { orig := rehostURLEnabled rehostURLEnabled = func(context.Context) bool { return false } // 透传,不下载 re-host defer func() { rehostURLEnabled = orig }() job := &SyncImageJob{PublicTaskID: "task_multi"} raw, _ := json.Marshal(map[string]any{ "created": 123, "data": []any{ map[string]any{"url": "https://cdn.example.com/a.png", "revised_prompt": "a"}, map[string]any{"url": "https://cdn.example.com/b.png"}, map[string]any{"url": "https://cdn.example.com/c.png"}, }, }) resultURL, normalized, err := buildImageResult(context.Background(), job, raw) if err != nil { t.Fatalf("unexpected error: %v", err) } if resultURL != "https://cdn.example.com/a.png" { t.Fatalf("resultURL = %q, want first image url", resultURL) } var out dto.ImageResponse if e := json.Unmarshal(normalized, &out); e != nil { t.Fatalf("normalized not a valid ImageResponse: %v", e) } if len(out.Data) != 3 { t.Fatalf("normalized data len = %d, want 3", len(out.Data)) } if out.Data[0].URL != "https://cdn.example.com/a.png" || out.Data[2].URL != "https://cdn.example.com/c.png" { t.Fatalf("normalized urls not preserved in order: %+v", out.Data) } if out.Data[0].RevisedPrompt != "a" { t.Fatalf("revised_prompt not preserved on first image") } } // TestBuildImageResult_AllOrNothingOnBadEntry 多图中任一条目无 url/b64 → 整体失败(all-or-nothing)。 func TestBuildImageResult_AllOrNothingOnBadEntry(t *testing.T) { orig := rehostURLEnabled rehostURLEnabled = func(context.Context) bool { return false } defer func() { rehostURLEnabled = orig }() job := &SyncImageJob{PublicTaskID: "task_bad"} raw, _ := json.Marshal(map[string]any{ "created": 1, "data": []any{ map[string]any{"url": "https://cdn.example.com/a.png"}, map[string]any{}, // 无 url 无 b64 → 整单失败 }, }) if _, _, err := buildImageResult(context.Background(), job, raw); err == nil { t.Fatal("entry without url/b64 should fail the whole task (all-or-nothing)") } } // TestBuildImageResult_MultiB64WithoutStorageErrors 多图 b64 无对象存储 → 整体失败。 func TestBuildImageResult_MultiB64WithoutStorageErrors(t *testing.T) { orig := acquireSyncImageFileSvc acquireSyncImageFileSvc = func(context.Context) (*lcommon.FileService, error) { return nil, lcommon.ErrStorageNotConfigured } defer func() { acquireSyncImageFileSvc = orig }() job := &SyncImageJob{PublicTaskID: "task_b64multi", TenantID: 1, UserID: 2} raw, _ := json.Marshal(map[string]any{ "created": 1, "data": []any{ map[string]any{"b64_json": base64.StdEncoding.EncodeToString([]byte("PNG1"))}, map[string]any{"b64_json": base64.StdEncoding.EncodeToString([]byte("PNG2"))}, }, }) if _, _, err := buildImageResult(context.Background(), job, raw); err == nil { t.Fatal("multi b64 without storage should error") } } // TestImageFailReason_StorageNotConfiguredFriendly 存储未配置的错误(穿过多层 %w 包裹) // 应映射为中文友好提示;其它错误保留技术细节。 func TestImageFailReason_StorageNotConfiguredFriendly(t *testing.T) { // 模拟 buildImageResult → rehostImage 的多层包裹,验证 errors.Is 仍能识别哨兵。 wrapped := fmt.Errorf("rehost b64_json[0]: %w", fmt.Errorf("object storage not configured, cannot re-host image: %w", lcommon.ErrStorageNotConfigured)) msg := imageFailReason(wrapped) if !strings.Contains(msg, "对象存储") { t.Fatalf("storage-not-configured should map to friendly message, got %q", msg) } other := imageFailReason(errors.New("upstream returned garbage")) if strings.Contains(other, "对象存储") { t.Fatalf("non-storage error must not map to storage message, got %q", other) } if !strings.Contains(other, "upstream returned garbage") { t.Fatalf("non-storage error should retain technical detail, got %q", other) } } func TestExtFromContentType(t *testing.T) { cases := map[string]string{ "image/png": ".png", "image/jpeg": ".jpg", "image/webp": ".webp", "": ".png", "text/plain": ".png", } for ct, want := range cases { if got := extFromContentType(ct); got != want { t.Errorf("extFromContentType(%q) = %q, want %q", ct, got, want) } } } func TestTruncateStr(t *testing.T) { if got := truncateStr("abcdef", 3); got != "abc" { t.Fatalf("truncate = %q, want abc", got) } if got := truncateStr("ab", 5); got != "ab" { t.Fatalf("truncate = %q, want ab", got) } // UTF-8 边界:不得把多字节字符(中文每字 3 字节)从中间切断。 utf8Cases := []struct { s string max int want string }{ {"中文", 3, "中"}, // 恰好一个完整字 {"中文", 4, "中"}, // 第 4 字节落在「文」中间 → 回退到「中」 {"中文", 5, "中"}, // 第 5 字节仍在「文」中间 → 回退到「中」 {"中文", 6, "中文"}, // 恰好容纳两字 {"中", 2, ""}, // 连一个字都放不下 → 空串(而非非法字节) } for _, c := range utf8Cases { got := truncateStr(c.s, c.max) if got != c.want { t.Fatalf("truncateStr(%q, %d) = %q, want %q", c.s, c.max, got, c.want) } if !utf8.ValidString(got) { t.Fatalf("truncateStr(%q, %d) = %q is not valid UTF-8", c.s, c.max, got) } } } func TestIsForbiddenIP(t *testing.T) { cases := []struct { ip string want bool }{ // 禁止:云元数据 / 私网 / loopback / link-local / multicast / unspecified {"169.254.169.254", true}, // AWS/云元数据 {"169.254.0.1", true}, // link-local {"10.0.0.1", true}, // 私网 10/8 {"172.16.0.1", true}, // 私网 172.16/12 {"172.31.255.255", true}, // 私网 172.16/12 上界 {"172.32.0.1", false}, // 公网(紧邻私网上界) {"192.168.1.1", true}, // 私网 192.168/16 {"127.0.0.1", true}, // loopback {"0.0.0.0", true}, // unspecified {"224.0.0.1", true}, // multicast {"::1", true}, // IPv6 loopback {"fe80::1", true}, // IPv6 link-local {"fc00::1", true}, // IPv6 ULA(私网) // 允许:公网 {"8.8.8.8", false}, {"1.1.1.1", false}, {"104.16.0.1", false}, } for _, c := range cases { ip := net.ParseIP(c.ip) if ip == nil { t.Fatalf("ParseIP(%q) failed", c.ip) } if got := isForbiddenIP(ip); got != c.want { t.Errorf("isForbiddenIP(%q) = %v, want %v", c.ip, got, c.want) } } } func TestValidateSSRFURL(t *testing.T) { blocked := []string{ // 字面量内网 IP(预检阶段直接拒绝) "http://169.254.169.254/latest/meta-data/", "http://10.0.0.1/x.png", "http://127.0.0.1/x.png", "http://192.168.1.1/x.png", "http://172.16.0.1/x.png", "https://[::1]/x.png", // 非 http(s) scheme "file:///etc/passwd", "gopher://127.0.0.1:6379/_FLUSHALL", "ftp://example.com/x.png", // 空 host "http:///x.png", } for _, u := range blocked { if err := validateSSRFURL(u); err == nil { t.Errorf("validateSSRFURL(%q) = nil, want error", u) } } allowed := []string{ // 公网字面量 IP 与域名均应通过预检(域名的网段校验由 dial 时兜底) "https://8.8.8.8/x.png", "http://example.com/image.png", "https://cdn.example.com/a/b/c.jpeg", } for _, u := range allowed { if err := validateSSRFURL(u); err != nil { t.Errorf("validateSSRFURL(%q) = %v, want nil", u, err) } } }