438 lines
16 KiB
Go
438 lines
16 KiB
Go
package runner
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/base64"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
|
"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 TestGeneratedAssetDecisionUploadsAudioContentDataURL(t *testing.T) {
|
|
item := map[string]any{
|
|
"type": "audio",
|
|
"content": "data:audio/mpeg;base64," + base64.StdEncoding.EncodeToString([]byte("inline audio")),
|
|
"mime_type": "audio/mpeg",
|
|
}
|
|
|
|
decision, err := generatedAssetDecisionForItem("speech.generations", item, defaultGeneratedAssetUploadPolicy())
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
if decision.Inline == nil {
|
|
t.Fatalf("expected inline audio content to be uploaded")
|
|
}
|
|
if decision.Inline.Kind != "audio" || decision.Inline.ContentType != "audio/mpeg" || decision.Inline.SourceKey != "content" {
|
|
t.Fatalf("unexpected inline audio metadata: %+v", decision.Inline)
|
|
}
|
|
if !containsString(decision.StripKeys, "content") {
|
|
t.Fatalf("uploaded audio content should be stripped: %+v", decision.StripKeys)
|
|
}
|
|
}
|
|
|
|
func TestGeneratedAssetDecisionDoesNotUploadPlainTextContent(t *testing.T) {
|
|
item := map[string]any{
|
|
"type": "text",
|
|
"content": base64.StdEncoding.EncodeToString([]byte(strings.Repeat("plain text ", 20))),
|
|
}
|
|
|
|
decision, err := generatedAssetDecisionForItem("speech.generations", item, defaultGeneratedAssetUploadPolicy())
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
if decision.Inline != nil || len(decision.StripKeys) > 0 {
|
|
t.Fatalf("plain text content should not be treated as generated media: %+v", decision)
|
|
}
|
|
}
|
|
|
|
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 TestGeneratedAssetDecisionStoresInlineLocallyWhenPolicyUploadNone(t *testing.T) {
|
|
item := map[string]any{
|
|
"b64_json": base64.StdEncoding.EncodeToString([]byte("inline image")),
|
|
}
|
|
|
|
decision, err := generatedAssetDecisionForItem("images.generations", item, generatedAssetUploadPolicyFromName(store.FileStorageResultUploadPolicyUploadNone))
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
if decision.Inline == nil || decision.URL != nil {
|
|
t.Fatalf("upload_none should still turn inline payloads into static URLs: %+v", decision)
|
|
}
|
|
if !containsString(decision.StripKeys, "b64_json") {
|
|
t.Fatalf("inline payload should be stripped before persistence: %+v", decision.StripKeys)
|
|
}
|
|
}
|
|
|
|
func TestGeneratedAssetUploadPolicyFromName(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
policyName string
|
|
want generatedAssetUploadPolicy
|
|
}{
|
|
{
|
|
name: "default",
|
|
policyName: store.FileStorageResultUploadPolicyDefault,
|
|
want: generatedAssetUploadPolicy{UploadInlineMedia: true, UploadURLMedia: false, StoreInlineMediaLocally: false},
|
|
},
|
|
{
|
|
name: "upload all",
|
|
policyName: store.FileStorageResultUploadPolicyUploadAll,
|
|
want: generatedAssetUploadPolicy{UploadInlineMedia: true, UploadURLMedia: true, StoreInlineMediaLocally: false},
|
|
},
|
|
{
|
|
name: "upload none",
|
|
policyName: store.FileStorageResultUploadPolicyUploadNone,
|
|
want: generatedAssetUploadPolicy{UploadInlineMedia: true, UploadURLMedia: false, StoreInlineMediaLocally: true},
|
|
},
|
|
}
|
|
|
|
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)
|
|
}
|
|
}
|
|
|
|
func TestUploadGeneratedAssetStoresLocalWhenNoChannels(t *testing.T) {
|
|
storageDir := t.TempDir()
|
|
service := &Service{cfg: config.Config{LocalGeneratedStorageDir: storageDir}}
|
|
payload := []byte{0x89, 'P', 'N', 'G', 0x0d, 0x0a, 0x1a, 0x0a, 0, 0, 0, 0}
|
|
asset := &generatedInlineAsset{
|
|
Bytes: payload,
|
|
ContentType: "image/jpeg",
|
|
Kind: "image",
|
|
SourceKey: "b64_json",
|
|
}
|
|
|
|
upload, contentType, kind, strategy, err := service.uploadGeneratedAsset(context.Background(), "task-123", asset, 0, nil, false)
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
if contentType != "image/png" || kind != "image" || strategy != "local_static_inline_media" {
|
|
t.Fatalf("unexpected local upload metadata: contentType=%s kind=%s strategy=%s", contentType, kind, strategy)
|
|
}
|
|
urlValue := stringFromAny(upload["url"])
|
|
if !strings.HasPrefix(urlValue, "/static/generated/gateway-result-task-123-01-") || !strings.HasSuffix(urlValue, ".png") {
|
|
t.Fatalf("unexpected local static URL: %s", urlValue)
|
|
}
|
|
if stringFromAny(upload["expiresAt"]) == "" {
|
|
t.Fatalf("local static upload should expose expiresAt: %+v", upload)
|
|
}
|
|
entries, err := os.ReadDir(storageDir)
|
|
if err != nil {
|
|
t.Fatalf("failed to read local static dir: %v", err)
|
|
}
|
|
if len(entries) != 1 || !strings.HasSuffix(entries[0].Name(), ".png") {
|
|
t.Fatalf("expected one PNG file in local static dir, got %+v", entries)
|
|
}
|
|
stored, err := os.ReadFile(filepath.Join(storageDir, entries[0].Name()))
|
|
if err != nil {
|
|
t.Fatalf("failed to read local static file: %v", err)
|
|
}
|
|
if !bytes.Equal(stored, payload) {
|
|
t.Fatalf("stored payload does not match source payload")
|
|
}
|
|
}
|
|
|
|
func TestUploadGeneratedAssetStoresAudioLocalWhenNoChannels(t *testing.T) {
|
|
storageDir := t.TempDir()
|
|
service := &Service{cfg: config.Config{LocalGeneratedStorageDir: storageDir}}
|
|
asset := &generatedInlineAsset{
|
|
Bytes: []byte("inline audio payload"),
|
|
ContentType: "audio/mpeg",
|
|
Kind: "audio",
|
|
SourceKey: "content",
|
|
}
|
|
|
|
upload, contentType, kind, strategy, err := service.uploadGeneratedAsset(context.Background(), "task-tts", asset, 0, nil, false)
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
if contentType != "audio/mpeg" || kind != "audio" || strategy != "local_static_inline_media" {
|
|
t.Fatalf("unexpected local audio metadata: contentType=%s kind=%s strategy=%s", contentType, kind, strategy)
|
|
}
|
|
urlValue := stringFromAny(upload["url"])
|
|
if !strings.HasPrefix(urlValue, "/static/generated/gateway-result-task-tts-01-") || !strings.HasSuffix(urlValue, ".mp3") {
|
|
t.Fatalf("unexpected local audio URL: %s", urlValue)
|
|
}
|
|
if stringFromAny(upload["expiresAt"]) == "" {
|
|
t.Fatalf("local audio static upload should expose expiresAt: %+v", upload)
|
|
}
|
|
entries, err := os.ReadDir(storageDir)
|
|
if err != nil {
|
|
t.Fatalf("failed to read local static dir: %v", err)
|
|
}
|
|
if len(entries) != 1 || !strings.HasSuffix(entries[0].Name(), ".mp3") {
|
|
t.Fatalf("expected one MP3 file in local static dir, got %+v", entries)
|
|
}
|
|
}
|
|
|
|
func TestUploadGeneratedRawMediaValueReplacesGeminiInlineDataWithAssetRef(t *testing.T) {
|
|
storageDir := t.TempDir()
|
|
service := &Service{cfg: config.Config{LocalGeneratedStorageDir: storageDir}}
|
|
payload := append([]byte{0x89, 'P', 'N', 'G', 0x0d, 0x0a, 0x1a, 0x0a}, bytes.Repeat([]byte{0}, 160)...)
|
|
raw := map[string]any{
|
|
"candidates": []any{
|
|
map[string]any{
|
|
"content": map[string]any{
|
|
"parts": []any{
|
|
map[string]any{
|
|
"inlineData": map[string]any{
|
|
"mimeType": "image/png",
|
|
"data": base64.StdEncoding.EncodeToString(payload),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
index := 0
|
|
|
|
uploaded, changed, err := service.uploadGeneratedRawMediaValue(context.Background(), "task-raw", "chat.completions", raw, "", nil, defaultGeneratedAssetUploadPolicy(), nil, &index)
|
|
if err != nil {
|
|
t.Fatalf("upload raw media: %v", err)
|
|
}
|
|
if !changed {
|
|
t.Fatal("expected raw inlineData to be replaced")
|
|
}
|
|
uploadedRaw := uploaded.(map[string]any)
|
|
candidates := uploadedRaw["candidates"].([]any)
|
|
candidate := candidates[0].(map[string]any)
|
|
content := candidate["content"].(map[string]any)
|
|
parts := content["parts"].([]any)
|
|
part := parts[0].(map[string]any)
|
|
inlineData := part["inlineData"].(map[string]any)
|
|
data, ok := inlineData["data"].(map[string]any)
|
|
if !ok {
|
|
t.Fatalf("inlineData.data should be an asset reference, got %+v", inlineData["data"])
|
|
}
|
|
ref, _ := data["assetRef"].(map[string]any)
|
|
if ref["sha256"] == "" || ref["contentType"] != "image/png" || ref["size"] != len(payload) {
|
|
t.Fatalf("unexpected asset ref: %+v", ref)
|
|
}
|
|
if urlValue := stringFromAny(data["url"]); !strings.HasPrefix(urlValue, "/static/generated/gateway-result-task-raw-01-") || !strings.HasSuffix(urlValue, ".png") {
|
|
t.Fatalf("unexpected raw media URL: %s", urlValue)
|
|
}
|
|
if inlineData["data"] == base64.StdEncoding.EncodeToString(payload) {
|
|
t.Fatal("raw inlineData still contains base64 payload")
|
|
}
|
|
entries, err := os.ReadDir(storageDir)
|
|
if err != nil {
|
|
t.Fatalf("read generated storage: %v", err)
|
|
}
|
|
if len(entries) != 1 || !strings.HasSuffix(entries[0].Name(), ".png") {
|
|
t.Fatalf("expected one generated PNG, got %+v", entries)
|
|
}
|
|
}
|
|
|
|
func TestUploadFileStoresLocalWhenNoChannels(t *testing.T) {
|
|
storageDir := t.TempDir()
|
|
service := &Service{cfg: config.Config{
|
|
LocalUploadedStorageDir: storageDir,
|
|
ServerMainBaseURL: "http://127.0.0.1:1",
|
|
ServerMainInternalToken: "change-me",
|
|
}}
|
|
payload := []byte("%PDF-1.4")
|
|
|
|
upload, err := service.UploadFile(context.Background(), FileUploadPayload{
|
|
Bytes: payload,
|
|
ContentType: "application/pdf",
|
|
FileName: "用户文件.png",
|
|
Source: "playground",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
urlValue := stringFromAny(upload["url"])
|
|
if !strings.HasPrefix(urlValue, "/static/uploaded/") || !strings.HasSuffix(urlValue, ".pdf") {
|
|
t.Fatalf("unexpected uploaded local static URL: %s", urlValue)
|
|
}
|
|
if stringFromAny(upload["expiresAt"]) == "" {
|
|
t.Fatalf("local uploaded static file should expose expiresAt: %+v", upload)
|
|
}
|
|
storageChannel, _ := upload["storageChannel"].(map[string]any)
|
|
if stringFromAny(storageChannel["provider"]) != "local_static" {
|
|
t.Fatalf("expected local static provider metadata, got %+v", upload["storageChannel"])
|
|
}
|
|
assetStorage, _ := upload["assetStorage"].(map[string]any)
|
|
if stringFromAny(assetStorage["strategy"]) != "local_static_upload" || stringFromAny(assetStorage["scene"]) != store.FileStorageSceneUpload {
|
|
t.Fatalf("unexpected upload asset storage metadata: %+v", assetStorage)
|
|
}
|
|
entries, err := os.ReadDir(storageDir)
|
|
if err != nil {
|
|
t.Fatalf("failed to read uploaded static dir: %v", err)
|
|
}
|
|
if len(entries) != 1 || !strings.HasSuffix(entries[0].Name(), ".pdf") {
|
|
t.Fatalf("expected one PDF file in uploaded static dir, got %+v", entries)
|
|
}
|
|
stored, err := os.ReadFile(filepath.Join(storageDir, entries[0].Name()))
|
|
if err != nil {
|
|
t.Fatalf("failed to read uploaded static file: %v", err)
|
|
}
|
|
if !bytes.Equal(stored, payload) {
|
|
t.Fatalf("stored uploaded payload does not match source payload")
|
|
}
|
|
}
|
|
|
|
func TestRedactGeneratedResultRawDataAudioPayload(t *testing.T) {
|
|
audioHex := strings.Repeat("49443304", 40)
|
|
result := map[string]any{
|
|
"raw_data": map[string]any{
|
|
"base_resp": map[string]any{"status_code": float64(0)},
|
|
"data": map[string]any{"audio": audioHex},
|
|
},
|
|
}
|
|
|
|
if !redactGeneratedResultRawData(result) {
|
|
t.Fatalf("expected raw audio payload to be redacted")
|
|
}
|
|
rawData, _ := result["raw_data"].(map[string]any)
|
|
data, _ := rawData["data"].(map[string]any)
|
|
audio, _ := data["audio"].(map[string]any)
|
|
if audio["redacted"] != true || audio["encoding"] != "hex" || audio["contentType"] != "audio/mpeg" {
|
|
t.Fatalf("unexpected redacted audio payload: %+v", audio)
|
|
}
|
|
if _, ok := audio["sha256"].(string); !ok {
|
|
t.Fatalf("expected redacted audio payload to include sha256: %+v", audio)
|
|
}
|
|
}
|