easyai-ai-gateway/apps/api/internal/runner/upload_test.go

package runner

import (
	"encoding/base64"
	"strings"
	"testing"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

func TestGeneratedAssetDecisionSkipsURLResultAndStripsInlinePayload(t *testing.T) {
	item := map[string]any{
		"b64_json": base64.StdEncoding.EncodeToString([]byte("inline image")),
		"url":      "https://cdn.example.com/generated.png",
	}

	decision, err := generatedAssetDecisionForItem("images.generations", item, defaultGeneratedAssetUploadPolicy())
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if decision.Inline != nil {
		t.Fatalf("URL media should not be uploaded by the default policy")
	}
	if !containsString(decision.StripKeys, "b64_json") {
		t.Fatalf("inline payload should be stripped when URL is already available: %+v", decision.StripKeys)
	}
}

func TestGeneratedAssetDecisionUploadsInlineImageBase64(t *testing.T) {
	item := map[string]any{
		"b64_json":  base64.StdEncoding.EncodeToString([]byte("inline image")),
		"mime_type": "image/jpeg",
	}

	decision, err := generatedAssetDecisionForItem("images.generations", item, defaultGeneratedAssetUploadPolicy())
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if decision.Inline == nil {
		t.Fatalf("expected inline image to be uploaded")
	}
	if decision.Inline.Kind != "image" || decision.Inline.ContentType != "image/jpeg" {
		t.Fatalf("unexpected inline image metadata: %+v", decision.Inline)
	}
	if !containsString(decision.StripKeys, "b64_json") {
		t.Fatalf("uploaded inline payload should be stripped: %+v", decision.StripKeys)
	}
}

func TestGeneratedAssetDecisionUploadsInlineVideoBuffer(t *testing.T) {
	item := map[string]any{
		"type":         "video",
		"video_buffer": []any{float64(0), float64(1), float64(2), float64(3)},
	}

	decision, err := generatedAssetDecisionForItem("videos.generations", item, defaultGeneratedAssetUploadPolicy())
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if decision.Inline == nil {
		t.Fatalf("expected inline video buffer to be uploaded")
	}
	if decision.Inline.Kind != "video" || decision.Inline.ContentType != "video/mp4" {
		t.Fatalf("unexpected inline video metadata: %+v", decision.Inline)
	}
	if !containsString(decision.StripKeys, "video_buffer") {
		t.Fatalf("uploaded video buffer should be stripped: %+v", decision.StripKeys)
	}
}

func TestGeneratedAssetDecisionUploadsDataURL(t *testing.T) {
	item := map[string]any{
		"url": "data:image/webp;base64," + base64.StdEncoding.EncodeToString([]byte("inline webp")),
	}

	decision, err := generatedAssetDecisionForItem("images.generations", item, defaultGeneratedAssetUploadPolicy())
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if decision.Inline == nil {
		t.Fatalf("expected data URL to be uploaded")
	}
	if decision.Inline.SourceKey != "url" || decision.Inline.ContentType != "image/webp" {
		t.Fatalf("unexpected data URL metadata: %+v", decision.Inline)
	}
	if !containsString(decision.StripKeys, "url") {
		t.Fatalf("uploaded data URL field should be stripped: %+v", decision.StripKeys)
	}
}

func TestGeneratedAssetDecisionUploadsURLWhenPolicyUploadAll(t *testing.T) {
	item := map[string]any{
		"type":      "video",
		"video_url": "https://cdn.example.com/generated.mp4",
	}

	decision, err := generatedAssetDecisionForItem("videos.generations", item, generatedAssetUploadPolicy{UploadInlineMedia: true, UploadURLMedia: true})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if decision.URL == nil {
		t.Fatalf("expected URL media to be uploaded")
	}
	if decision.URL.Kind != "video" || decision.URL.SourceKey != "video_url" {
		t.Fatalf("unexpected URL media metadata: %+v", decision.URL)
	}
	if !containsString(decision.StripKeys, "video_url") {
		t.Fatalf("uploaded URL field should be stripped: %+v", decision.StripKeys)
	}
}

func TestGeneratedAssetDecisionSkipsAllWhenPolicyUploadNone(t *testing.T) {
	item := map[string]any{
		"b64_json": base64.StdEncoding.EncodeToString([]byte("inline image")),
	}

	decision, err := generatedAssetDecisionForItem("images.generations", item, generatedAssetUploadPolicy{UploadInlineMedia: false, UploadURLMedia: false})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if decision.Inline != nil || decision.URL != nil || len(decision.StripKeys) != 0 {
		t.Fatalf("upload_none should keep the result unchanged: %+v", decision)
	}
}

func TestGeneratedAssetUploadPolicyFromName(t *testing.T) {
	tests := []struct {
		name       string
		policyName string
		want       generatedAssetUploadPolicy
	}{
		{
			name:       "default",
			policyName: store.FileStorageResultUploadPolicyDefault,
			want:       generatedAssetUploadPolicy{UploadInlineMedia: true, UploadURLMedia: false},
		},
		{
			name:       "upload all",
			policyName: store.FileStorageResultUploadPolicyUploadAll,
			want:       generatedAssetUploadPolicy{UploadInlineMedia: true, UploadURLMedia: true},
		},
		{
			name:       "upload none",
			policyName: store.FileStorageResultUploadPolicyUploadNone,
			want:       generatedAssetUploadPolicy{UploadInlineMedia: false, UploadURLMedia: false},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := generatedAssetUploadPolicyFromName(tt.policyName)
			if got != tt.want {
				t.Fatalf("unexpected policy: got %+v, want %+v", got, tt.want)
			}
		})
	}
}

func TestResolvedGeneratedAssetContentTypePrefersDetectedMedia(t *testing.T) {
	pngPayload := []byte{0x89, 'P', 'N', 'G', 0x0d, 0x0a, 0x1a, 0x0a, 0, 0, 0, 0}

	contentType := resolvedGeneratedAssetContentType("image/jpeg", "image", pngPayload)
	if contentType != "image/png" {
		t.Fatalf("expected detected PNG content type, got %s", contentType)
	}
	if extension := fileExtensionForContentType(contentType, "image"); extension != ".png" {
		t.Fatalf("expected PNG extension, got %s", extension)
	}
}

func TestResolvedGeneratedAssetContentTypeKeepsDeclaredMediaWhenDetectionIsGeneric(t *testing.T) {
	contentType := resolvedGeneratedAssetContentType("image/webp", "image", []byte("not enough media bytes"))
	if contentType != "image/webp" {
		t.Fatalf("expected declared webp content type, got %s", contentType)
	}
}

func TestGeneratedAssetFileNameIsUniqueAndTyped(t *testing.T) {
	first := generatedAssetFileName("663e19cd4fa9d8078385c7c9", 0, "image/png", "image")
	second := generatedAssetFileName("663e19cd4fa9d8078385c7c9", 0, "image/png", "image")
	if first == second {
		t.Fatalf("expected generated file names to be unique, both were %s", first)
	}
	if !strings.HasPrefix(first, "gateway-result-663e19cd4fa9d8078385c7c9-01-") || !strings.HasSuffix(first, ".png") {
		t.Fatalf("unexpected generated file name: %s", first)
	}
}