moteva/server/internal/application/background_removal_test.go

package application

import (
	"bytes"
	"context"
	"image"
	"image/color"
	"image/png"
	"testing"
	"time"

	"img_infinite_canvas/internal/domain/design"
	"img_infinite_canvas/internal/infrastructure/memory"
)

type fakeBackgroundRemover struct {
	result design.BackgroundRemoval
	err    error
	calls  int
	last   design.BackgroundRemovalRequest
}

func (r *fakeBackgroundRemover) RemoveBackground(ctx context.Context, req design.BackgroundRemovalRequest) (design.BackgroundRemoval, error) {
	r.calls++
	r.last = req
	if r.err != nil {
		return design.BackgroundRemoval{}, r.err
	}
	return r.result, nil
}

func TestRunRemoveBackgroundCreatesAsyncPlaceholderBesideTarget(t *testing.T) {
	ctx := context.Background()
	repo := memory.NewProjectRepository()
	queue := &fakeJobQueue{}
	service := NewDesignService(repo, nil, nil, nil, nil)
	service.SetJobQueue(queue)

	source := testPNGDataURL(testColor(32, 96, 180, 255), 4, 3)
	project := design.Project{
		ID:        "project-remove-bg-placeholder",
		Title:     "Remove background",
		Status:    design.StatusReady,
		UpdatedAt: time.Now(),
		Nodes: []design.Node{{
			ID:      "image-1",
			Type:    design.NodeTypeImage,
			Title:   "Source",
			X:       100,
			Y:       120,
			Width:   200,
			Height:  150,
			Content: source,
			Status:  "success",
		}},
	}
	if err := repo.Save(ctx, project); err != nil {
		t.Fatal(err)
	}

	next, err := service.RunNodeActionAsync(ctx, project.ID, "image-1", design.NodeActionRequest{
		Action: "remove-background",
		Prompt: "去背景",
	})
	if err != nil {
		t.Fatal(err)
	}

	if len(next.Nodes) != 2 {
		t.Fatalf("expected source plus placeholder, got %#v", next.Nodes)
	}
	sourceNode := next.Nodes[0]
	placeholder := next.Nodes[1]
	if sourceNode.Status != "success" || sourceNode.Title != "Source" {
		t.Fatalf("expected source node to remain unchanged, got %#v", sourceNode)
	}
	if placeholder.Status != "generating" || placeholder.Title != "去背景中" {
		t.Fatalf("expected generating placeholder, got %#v", placeholder)
	}
	if placeholder.X != sourceNode.X+sourceNode.Width+32 || placeholder.Y != sourceNode.Y {
		t.Fatalf("expected placeholder beside source, got x=%f y=%f", placeholder.X, placeholder.Y)
	}
	if placeholder.Content != sourceNode.Content || placeholder.ParentID != sourceNode.ID {
		t.Fatalf("expected placeholder to reference source image, got %#v", placeholder)
	}
	if placeholder.LayerRole != "background-removed" || placeholder.Tone != "transparent-image" {
		t.Fatalf("expected transparent-image placeholder role, got %#v", placeholder)
	}
	if len(queue.jobs) != 1 {
		t.Fatalf("expected one async job, got %#v", queue.jobs)
	}
	if queue.jobs[0].Kind != design.JobNodeActionGeneration || queue.jobs[0].Request.Action != "remove-background" {
		t.Fatalf("expected remove-background job, got %#v", queue.jobs[0])
	}
	if queue.jobs[0].Request.Options[nodeActionResultNodeIDOption] != placeholder.ID {
		t.Fatalf("expected job to carry result node id %s, got %#v", placeholder.ID, queue.jobs[0].Request.Options)
	}
	task, err := service.GeneratorTask(ctx, project.ID, next.LastThreadID)
	if err != nil {
		t.Fatal(err)
	}
	if task.GeneratorName != removeBackgroundGeneratorName || task.OriginalInputArgs.Prompt != "" {
		t.Fatalf("expected removeBG task without prompt, got generator=%q prompt=%q", task.GeneratorName, task.OriginalInputArgs.Prompt)
	}
}

func TestCompleteRemoveBackgroundWritesResultToPlaceholder(t *testing.T) {
	ctx := context.Background()
	repo := memory.NewProjectRepository()
	assets := newFakeAssetStorage()
	queue := &fakeJobQueue{}
	remover := &fakeBackgroundRemover{
		result: design.BackgroundRemoval{
			Image:       testPNGBytes(testColor(32, 96, 180, 255), 4, 3),
			ContentType: "image/png",
			Width:       4,
			Height:      3,
		},
	}
	service := NewDesignService(repo, nil, assets, nil, nil)
	service.SetJobQueue(queue)
	service.SetBackgroundRemover(remover)

	project := design.Project{
		ID:        "project-remove-bg-complete",
		Title:     "Remove background",
		Status:    design.StatusReady,
		UpdatedAt: time.Now(),
		Nodes: []design.Node{{
			ID:      "image-1",
			Type:    design.NodeTypeImage,
			Title:   "Source",
			X:       100,
			Y:       120,
			Width:   200,
			Height:  150,
			Content: testPNGDataURL(testColor(32, 96, 180, 255), 4, 3),
			Status:  "success",
		}},
	}
	if err := repo.Save(ctx, project); err != nil {
		t.Fatal(err)
	}

	started, err := service.RunNodeActionAsync(ctx, project.ID, "image-1", design.NodeActionRequest{
		Action: "remove-background",
		Options: map[string]string{
			"imageWidth":  "1002",
			"imageHeight": "1002",
			"inputSize":   "1024",
		},
	})
	if err != nil {
		t.Fatal(err)
	}
	if len(queue.jobs) != 1 {
		t.Fatalf("expected one queued job, got %#v", queue.jobs)
	}
	if err := service.ProcessJob(ctx, queue.jobs[0]); err != nil {
		t.Fatal(err)
	}

	saved, err := repo.Get(ctx, started.ID)
	if err != nil {
		t.Fatal(err)
	}
	if remover.calls != 1 {
		t.Fatalf("expected one background removal call, got %d", remover.calls)
	}
	if remover.last.Width != 1002 || remover.last.Height != 1002 || remover.last.InputSize != 1024 {
		t.Fatalf("expected target dimensions in background removal request, got %#v", remover.last)
	}
	if len(saved.Nodes) != 2 {
		t.Fatalf("expected source plus result node, got %#v", saved.Nodes)
	}
	if saved.Nodes[0].Status != "success" || saved.Nodes[0].Content != project.Nodes[0].Content {
		t.Fatalf("expected source node unchanged, got %#v", saved.Nodes[0])
	}
	result := saved.Nodes[1]
	if result.Status != "success" || result.Title != "Background Removed" {
		t.Fatalf("expected completed result node, got %#v", result)
	}
	if result.LayerRole != "background-removed" || result.Tone != "transparent-image" {
		t.Fatalf("expected transparent PNG result role, got %#v", result)
	}
	if result.Content != "http://localhost:19000/canvas-assets/project-remove-bg-complete-image-1-remove-background.webp" {
		t.Fatalf("unexpected result URL %s", result.Content)
	}
	if saved.Status != design.StatusReady {
		t.Fatalf("expected ready project, got %s", saved.Status)
	}
	if len(saved.Messages) == 0 || saved.Messages[len(saved.Messages)-1].Status != "success" {
		t.Fatalf("expected success message, got %#v", saved.Messages)
	}
	select {
	case upload := <-assets.uploaded:
		if upload.ContentType != "image/webp" {
			t.Fatalf("expected webp upload, got %s", upload.ContentType)
		}
	default:
		t.Fatal("expected uploaded background removal asset")
	}
}

func testColor(r uint8, g uint8, b uint8, a uint8) color.NRGBA {
	return color.NRGBA{R: r, G: g, B: b, A: a}
}

func testPNGBytes(fill color.NRGBA, width int, height int) []byte {
	img := image.NewNRGBA(image.Rect(0, 0, width, height))
	for y := 0; y < height; y++ {
		for x := 0; x < width; x++ {
			img.SetNRGBA(x, y, fill)
		}
	}
	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		panic(err)
	}
	return buf.Bytes()
}