Files
gen2d/backend/internal/service/pipeline_test.go
T

241 lines
6.5 KiB
Go
Raw Normal View History

package service
import (
"context"
"testing"
)
func TestPipeline_HappyPath(t *testing.T) {
QualityChecker = func(_ context.Context, _ []GeneratedImage, _ map[string]string) (bool, string, error) {
return true, "", nil
}
defer func() { QualityChecker = defaultCheckQuality }()
output, err := RunPipeline(context.Background(), PipelineInput{
Prompt: "一个拿剑的小人",
AssetType: "sprite",
ProjectStyle: map[string]string{
"artStyle": "pixel",
"palette": "warm",
},
Params: AssetParams{
Resolution: 64,
Format: "spritesheet",
},
})
if err != nil {
t.Fatalf("RunPipeline failed: %v", err)
}
if len(output.Assets) == 0 {
t.Fatal("expected non-empty assets")
}
if output.Metadata.FrameWidth != 64 {
t.Errorf("expected FrameWidth=64, got %d", output.Metadata.FrameWidth)
}
if output.Metadata.FrameHeight != 64 {
t.Errorf("expected FrameHeight=64, got %d", output.Metadata.FrameHeight)
}
}
func TestPipeline_RetryThenPass(t *testing.T) {
QualityChecker = NewCountedQualityChecker(3)
defer func() { QualityChecker = defaultCheckQuality }()
output, err := RunPipeline(context.Background(), PipelineInput{
Prompt: "一把火焰剑",
AssetType: "sprite",
Params: AssetParams{
Resolution: 32,
Frames: FrameParams{
Directions: 4,
FramesPerDirection: 2,
},
Format: "spritesheet",
},
})
if err != nil {
t.Fatalf("RunPipeline failed: %v", err)
}
if len(output.Assets) != 8 {
t.Errorf("expected 8 assets, got %d", len(output.Assets))
}
if output.Metadata.FrameCount != 8 {
t.Errorf("expected FrameCount=8, got %d", output.Metadata.FrameCount)
}
if output.Metadata.Directions != 4 {
t.Errorf("expected Directions=4, got %d", output.Metadata.Directions)
}
}
func TestPipeline_MaxRetryDegrade(t *testing.T) {
QualityChecker = AlwaysFailQualityChecker()
defer func() { QualityChecker = defaultCheckQuality }()
output, err := RunPipeline(context.Background(), PipelineInput{
Prompt: "一只飞龙",
AssetType: "sprite",
Params: AssetParams{
Resolution: 48,
},
})
if err != nil {
t.Fatalf("RunPipeline failed: %v", err)
}
if len(output.Assets) == 0 {
t.Fatal("expected non-empty assets even on degrade")
}
if output.Metadata.FrameWidth != 48 {
t.Errorf("expected FrameWidth=48, got %d", output.Metadata.FrameWidth)
}
}
func TestPipeline_StyleMerge(t *testing.T) {
QualityChecker = func(_ context.Context, _ []GeneratedImage, style map[string]string) (bool, string, error) {
if style["artStyle"] != "realistic" {
t.Errorf("expected artStyle=realistic (task override), got %s", style["artStyle"])
}
if style["palette"] != "warm" {
t.Errorf("expected palette=warm (from project), got %s", style["palette"])
}
return true, "", nil
}
defer func() { QualityChecker = defaultCheckQuality }()
_, err := RunPipeline(context.Background(), PipelineInput{
Prompt: "测试风格合并",
AssetType: "sprite",
ProjectStyle: map[string]string{
"artStyle": "pixel",
"palette": "warm",
},
TaskStyle: map[string]string{
"artStyle": "realistic",
},
Params: AssetParams{Resolution: 64},
})
if err != nil {
t.Fatalf("RunPipeline failed: %v", err)
}
}
func TestPipeline_WithPromptOptimizer(t *testing.T) {
// 带标签时 PromptOptimizer 应优化原始提示词
QualityChecker = func(_ context.Context, _ []GeneratedImage, _ map[string]string) (bool, string, error) {
return true, "", nil
}
defer func() { QualityChecker = defaultCheckQuality }()
output, err := RunPipeline(context.Background(), PipelineInput{
Prompt: "一个战士",
AssetType: "sprite",
Tags: []string{"像素", "战士", "持剑"},
Params: AssetParams{
Resolution: 64,
Format: "spritesheet",
},
})
if err != nil {
t.Fatalf("RunPipeline with tags failed: %v", err)
}
// 有标签时应有优化后的输出
if len(output.Assets) == 0 {
t.Fatal("expected non-empty assets")
}
}
func TestPipeline_WithoutTags(t *testing.T) {
// 无标签时 PromptOptimizer 应跳过,原始提示词直接进入 PromptBuilder
QualityChecker = func(_ context.Context, _ []GeneratedImage, _ map[string]string) (bool, string, error) {
return true, "", nil
}
defer func() { QualityChecker = defaultCheckQuality }()
rawPrompt := "一个原始提示词没有标签"
output, err := RunPipeline(context.Background(), PipelineInput{
Prompt: rawPrompt,
AssetType: "sprite",
Params: AssetParams{
Resolution: 32,
Format: "spritesheet",
},
})
if err != nil {
t.Fatalf("RunPipeline without tags failed: %v", err)
}
if len(output.Assets) == 0 {
t.Fatal("expected non-empty assets")
}
}
func TestPipeline_PromptOptimizerFallback(t *testing.T) {
// PromptAgent 优化失败(无 API key)也不应阻塞管线
QualityChecker = func(_ context.Context, _ []GeneratedImage, _ map[string]string) (bool, string, error) {
return true, "", nil
}
defer func() { QualityChecker = defaultCheckQuality }()
output, err := RunPipeline(context.Background(), PipelineInput{
Prompt: "一个火球术",
AssetType: "animation",
Tags: []string{"火焰", "魔法"},
Params: AssetParams{
Resolution: 64,
Format: "spritesheet",
},
})
if err != nil {
t.Fatalf("RunPipeline with fallback prompt agent failed: %v", err)
}
if len(output.Assets) == 0 {
t.Fatal("expected non-empty assets even on prompt agent fallback")
}
}
func TestPipeline_PromptOptimizerSkipsEmptyTags(t *testing.T) {
// 无标签时 prompt 保持原样
QualityChecker = func(_ context.Context, _ []GeneratedImage, _ map[string]string) (bool, string, error) {
return true, "", nil
}
defer func() { QualityChecker = defaultCheckQuality }()
output, err := RunPipeline(context.Background(), PipelineInput{
Prompt: "原始提示词",
AssetType: "sprite",
Params: AssetParams{Resolution: 32},
})
if err != nil {
t.Fatalf("RunPipeline failed: %v", err)
}
if len(output.Assets) == 0 {
t.Fatal("expected non-empty assets")
}
}
func TestPipeline_PromptOptimizerRefinesPrompt(t *testing.T) {
// 有标签时 pipeline 产出优化后的 prompt
QualityChecker = func(_ context.Context, _ []GeneratedImage, _ map[string]string) (bool, string, error) {
return true, "", nil
}
defer func() { QualityChecker = defaultCheckQuality }()
output, err := RunPipeline(context.Background(), PipelineInput{
Prompt: "一个战士",
AssetType: "sprite",
Tags: []string{"像素", "战士"},
Params: AssetParams{Resolution: 32},
})
if err != nil {
t.Fatalf("RunPipeline with tags failed: %v", err)
}
if len(output.Assets) == 0 {
t.Fatal("expected non-empty assets")
}
}