66 lines
1.9 KiB
Go
66 lines
1.9 KiB
Go
package classification
|
|
|
|
import (
|
|
"context"
|
|
"os"
|
|
"path/filepath"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
// TestClassify 使用仓库根目录的 image.png 调用真实分类服务进行端到端验证。
|
|
func TestClassify(t *testing.T) {
|
|
client := NewClientWithAddr("127.0.0.1", 8000)
|
|
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)
|
|
}
|
|
|
|
result, err := client.Classify(ctx, raw)
|
|
if err != nil {
|
|
t.Fatalf("Classify 失败: %v", err)
|
|
}
|
|
|
|
if result == nil {
|
|
t.Fatal("分类返回空结果")
|
|
}
|
|
if result.CategoryMain == "" {
|
|
t.Fatal("分类大类 category_main 为空")
|
|
}
|
|
if !validLevel(result.Formality) {
|
|
t.Fatalf("正式程度 formality 应为 1~5，实际 %q", result.Formality)
|
|
}
|
|
if !validLevel(result.VisualImpact) {
|
|
t.Fatalf("视觉冲击力 visual_impact 应为 1~5，实际 %q", result.VisualImpact)
|
|
}
|
|
|
|
t.Logf("分类成功: category_main=%s(%.3f), category_sub=%s(%.3f), style=%s(%.3f), color_main=%s(%.3f)/%s(%.3f), formality=%s(%s), visual_impact=%s(%s), occasions=%s",
|
|
result.CategoryMain, result.CategoryMainConfidence,
|
|
result.CategorySub, result.CategorySubConfidence,
|
|
result.Style, result.StyleConfidence,
|
|
result.ColorMain, result.ColorMainConfidence,
|
|
result.ColorSub, result.ColorSubConfidence,
|
|
result.Formality, result.FormalityLabel,
|
|
result.VisualImpact, result.VisualImpactLabel,
|
|
result.Occasions)
|
|
}
|
|
|
|
// validLevel 校验 1-5 等级字符串是否合法。
|
|
func validLevel(s string) bool {
|
|
switch s {
|
|
case "1", "2", "3", "4", "5":
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|