NianAIGC/backend/internal/assets/image_format_test.go

package assets

import (
	"bytes"
	"context"
	"encoding/base64"
	"errors"
	"image"
	"image/jpeg"
	"image/png"
	"io"
	"strings"
	"testing"
)

func imageFormatFixtures(t *testing.T) (pngBytes, jpegBytes, webpBytes []byte) {
	t.Helper()
	img := image.NewRGBA(image.Rect(0, 0, 2, 2))
	var p, j bytes.Buffer
	if err := png.Encode(&p, img); err != nil {
		t.Fatal(err)
	}
	if err := jpeg.Encode(&j, img, nil); err != nil {
		t.Fatal(err)
	}
	w, err := base64.StdEncoding.DecodeString("UklGRiIAAABXRUJQVlA4IBYAAAAwAQCdASoBAAEADsD+JaQAA3AAAAAA")
	if err != nil {
		t.Fatal(err)
	}
	return p.Bytes(), j.Bytes(), w
}

func TestNormalizeImageFileUsesBytesAndPreservesCompatibleNames(t *testing.T) {
	pngData, jpegData, webpData := imageFormatFixtures(t)
	for _, test := range []struct {
		label, name, mime, wantName, wantMIME string
		data                                  []byte
	}{
		{"opaque PNG", "result.image", "application/octet-stream", "result.png", "image/png", pngData},
		{"JPEG", "result.image", "", "result.jpg", "image/jpeg", jpegData},
		{"WebP despite bad MIME", "result.image", "image/png", "result.webp", "image/webp", webpData},
		{"wrong extension", "result.png", "image/png", "result.jpg", "image/jpeg", jpegData},
		{"no suffix", "result", "", "result.png", "image/png", pngData},
		{"uppercase", "result.PNG", "image/png; charset=binary", "result.PNG", "image/png", pngData},
		{"JPEG alias", "result.JPEG", "application/octet-stream", "result.JPEG", "image/jpeg", jpegData},
		{"Chinese filename", "葫芦娃和爷爷.image", "", "葫芦娃和爷爷.png", "image/png", pngData},
		{"multi-dot filename", "result.v1.image", "", "result.v1.png", "image/png", pngData},
		{"empty filename", "", "", "image.png", "image/png", pngData},
		{"video", "result.mp4", "video/mp4", "result.mp4", "video/mp4", []byte("video content")},
		{"unknown does not trust MIME", "result.image", "image/png", "result.image", "image/png", []byte(`{"error":"not an image"}`)},
		{"empty content", "result.image", "", "result.image", "", nil},
	} {
		t.Run(test.label, func(t *testing.T) {
			name, mime := normalizeImageFile(test.name, test.mime, test.data)
			if name != test.wantName || mime != test.wantMIME {
				t.Fatalf("got %q / %q, want %q / %q", name, mime, test.wantName, test.wantMIME)
			}
		})
	}
}

func TestImportGeneratedNormalizesImageNameStorageAndContentType(t *testing.T) {
	pngData, jpegData, webpData := imageFormatFixtures(t)
	for _, test := range []struct {
		extension, contentType string
		data                   []byte
	}{{"png", "image/png", pngData}, {"jpg", "image/jpeg", jpegData}, {"webp", "image/webp", webpData}} {
		t.Run(test.extension, func(t *testing.T) {
			cat, blobs := &memoryCatalog{}, &memoryBlobs{}
			remote := &memoryRemote{blob: Blob{Body: io.NopCloser(bytes.NewReader(test.data)), ContentType: "application/octet-stream", Size: int64(len(test.data))}}
			svc := NewService(cat, blobs, remote, nil, func(prefix string) string { return prefix + "-format" })
			url := "https://provider.test/result.image?signature=fixture"
			asset, err := svc.ImportGenerated(context.Background(), PlatformScope("owner-a"), ImportGeneratedCommand{URL: url, Capability: "image.generate"})
			if err != nil {
				t.Fatal(err)
			}
			if asset.Name != "result."+test.extension || !strings.HasSuffix(asset.StoragePath, "-result."+test.extension) || asset.Kind != KindImage {
				t.Fatalf("asset=%#v", asset)
			}
			if asset.Metadata["contentType"] != test.contentType || asset.Metadata["importedFrom"] != url || blobs.putContentType != test.contentType {
				t.Fatalf("metadata=%#v stored MIME=%q", asset.Metadata, blobs.putContentType)
			}
			if !bytes.Equal(blobs.putBody, test.data) {
				t.Fatal("renaming changed the original image bytes")
			}
		})
	}
}

func TestUploadRecognizesReuploadedLegacyImage(t *testing.T) {
	pngData, _, _ := imageFormatFixtures(t)
	cat, blobs := &memoryCatalog{}, &memoryBlobs{}
	svc := NewService(cat, blobs, nil, nil, nil)
	asset, err := svc.Upload(context.Background(), PlatformScope("owner-a"), UploadCommand{FileName: "download.image", ContentType: "application/octet-stream", Bytes: pngData})
	if err != nil {
		t.Fatal(err)
	}
	if asset.Name != "download.png" || asset.Kind != KindImage || asset.Metadata["contentType"] != "image/png" || !strings.HasSuffix(asset.StoragePath, "-download.png") || !bytes.Equal(blobs.putBody, pngData) {
		t.Fatalf("upload=%#v", asset)
	}
}

func TestNormalizeImageDownloadStreamsAllBytesAndClosesOriginal(t *testing.T) {
	pngData, _, _ := imageFormatFixtures(t)
	for _, data := range [][]byte{pngData, append(append([]byte(nil), pngData...), bytes.Repeat([]byte{7}, 4096)...)} {
		body := &imageTrackedBody{Reader: bytes.NewReader(data)}
		name, blob, err := NormalizeImageDownload("legacy.image", Blob{Body: body, ContentType: "application/octet-stream", Size: int64(len(data))})
		if err != nil || name != "legacy.png" || blob.ContentType != "image/png" || blob.Size != int64(len(data)) {
			t.Fatalf("name=%q blob=%#v err=%v", name, blob, err)
		}
		if body.readBytes > 512 || body.closed {
			t.Fatalf("prefix read consumed %d bytes, closed=%v", body.readBytes, body.closed)
		}
		got, err := io.ReadAll(blob.Body)
		if err != nil || !bytes.Equal(got, data) {
			t.Fatalf("stream mismatch; error=%v", err)
		}
		if err := blob.Body.Close(); err != nil || !body.closed {
			t.Fatalf("original body not closed; error=%v", err)
		}
	}
}

func TestNormalizeImageDownloadUnknownDataIsNotRenamed(t *testing.T) {
	for _, data := range [][]byte{[]byte("video bytes"), nil} {
		name, blob, err := NormalizeImageDownload("legacy.image", Blob{Body: io.NopCloser(bytes.NewReader(data)), ContentType: "application/octet-stream", Size: int64(len(data))})
		if err != nil || name != "legacy.image" || blob.ContentType != "application/octet-stream" {
			t.Fatalf("name=%q blob=%#v err=%v", name, blob, err)
		}
		got, _ := io.ReadAll(blob.Body)
		_ = blob.Body.Close()
		if !bytes.Equal(got, data) {
			t.Fatal("unknown bytes changed")
		}
	}
}

func TestNormalizeImageDownloadReadError(t *testing.T) {
	failure := errors.New("read failed")
	body := &imageTrackedBody{Reader: imageErrorReader{err: failure}}
	_, _, err := NormalizeImageDownload("legacy.image", Blob{Body: body})
	if !errors.Is(err, failure) || body.closed {
		t.Fatalf("error=%v closed=%v", err, body.closed)
	}
	_ = body.Close()
	if _, _, err := NormalizeImageDownload("legacy.image", Blob{}); !errors.Is(err, ErrBlobNotFound) {
		t.Fatalf("nil body error=%v", err)
	}
}

type imageTrackedBody struct {
	io.Reader
	readBytes int
	closed    bool
}

func (b *imageTrackedBody) Read(p []byte) (int, error) {
	n, err := b.Reader.Read(p)
	b.readBytes += n
	return n, err
}

func (b *imageTrackedBody) Close() error {
	b.closed = true
	return nil
}

type imageErrorReader struct{ err error }

func (r imageErrorReader) Read([]byte) (int, error) { return 0, r.err }