54 lines
1.5 KiB
Go
54 lines
1.5 KiB
Go
package matting
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"os"
|
|
"path/filepath"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
// TestMatting 使用仓库根目录的 image.png 调用真实抠图服务进行端到端验证。
|
|
func TestMatting(t *testing.T) {
|
|
client := NewClientWithAddr("127.0.0.1", 8001)
|
|
ctx, cancel := context.WithTimeout(context.Background(), 180*time.Second)
|
|
defer cancel()
|
|
|
|
// 前置：确认抠图服务已就绪（模型加载完成后 /healthz 才会响应）
|
|
if err := client.HealthCheck(ctx); err != nil {
|
|
t.Skipf("抠图服务未就绪，跳过真实抠图测试: %v", err)
|
|
}
|
|
|
|
imgPath := filepath.Join("..", "..", "image.png")
|
|
raw, err := os.ReadFile(imgPath)
|
|
if err != nil {
|
|
t.Fatalf("读取 image.png 失败: %v", err)
|
|
}
|
|
|
|
out, err := client.Matting(ctx, raw)
|
|
if err != nil {
|
|
t.Fatalf("Matting 失败: %v", err)
|
|
}
|
|
|
|
if len(out) == 0 {
|
|
t.Fatal("抠图返回空结果")
|
|
}
|
|
|
|
pngMagic := []byte{0x89, 'P', 'N', 'G', 0x0D, 0x0A, 0x1A, 0x0A}
|
|
if len(out) < len(pngMagic) || !bytes.Equal(out[:len(pngMagic)], pngMagic) {
|
|
preview := out
|
|
if len(preview) > len(pngMagic) {
|
|
preview = preview[:len(pngMagic)]
|
|
}
|
|
t.Fatalf("返回结果不是 PNG，前 %d 字节: % x", len(preview), preview)
|
|
}
|
|
|
|
// 保存结果便于人工核验
|
|
resultPath := filepath.Join(t.TempDir(), "matting_result.png")
|
|
if err := os.WriteFile(resultPath, out, 0o644); err != nil {
|
|
t.Fatalf("保存抠图结果失败: %v", err)
|
|
}
|
|
t.Logf("抠图成功，输入 %d 字节 -> 输出 %d 字节，结果已保存到 %s", len(raw), len(out), resultPath)
|
|
}
|