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

package runner

import (
	"bytes"
	"context"
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base64"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

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

func TestDirectOSSUploadSignsRequestAndReturnsPublicURL(t *testing.T) {
	payload := []byte("production-isomorphic-image")
	var calls atomic.Int64
	var uploadedPath string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		call := calls.Add(1)
		if r.Method != http.MethodPut {
			t.Errorf("method=%s", r.Method)
		}
		if r.Header.Get("Content-Type") != "image/png" || r.Header.Get("Date") == "" {
			t.Errorf("missing OSS upload headers")
		}
		stringToSign := "PUT\n\nimage/png\n" + r.Header.Get("Date") + "\n/media-bucket" + r.URL.EscapedPath()
		mac := hmac.New(sha1.New, []byte("access-secret"))
		_, _ = mac.Write([]byte(stringToSign))
		expectedAuthorization := "OSS access-key:" + base64.StdEncoding.EncodeToString(mac.Sum(nil))
		if r.Header.Get("Authorization") != expectedAuthorization {
			t.Errorf("invalid OSS request signature")
		}
		body, err := io.ReadAll(r.Body)
		if err != nil || !bytes.Equal(body, payload) {
			t.Errorf("uploaded payload differs: err=%v", err)
		}
		uploadedPath = r.URL.EscapedPath()
		if call == 1 {
			http.Error(w, "retry", http.StatusBadGateway)
			return
		}
		w.WriteHeader(http.StatusOK)
	}))
	defer server.Close()

	cfg := config.Config{
		MediaOSSDirectEnabled:   true,
		MediaOSSEndpoint:        server.URL,
		MediaOSSBucket:          "media-bucket",
		MediaOSSAccessKeyID:     "access-key",
		MediaOSSAccessKeySecret: "access-secret",
		MediaOSSPublicBaseURL:   "https://cdn.example.com",
		MediaOSSObjectPrefix:    "easyai-ai-gateway/production/media",
	}
	uploader := newDirectOSSUploader(cfg)
	uploader.now = func() time.Time {
		return time.Date(2026, time.July, 30, 12, 0, 0, 0, time.UTC)
	}
	service := &Service{cfg: cfg, directOSS: uploader}
	uploaded, err := service.UploadFile(context.Background(), FileUploadPayload{
		Bytes:       payload,
		ContentType: "image/png",
		FileName:    "input.png",
		Scene:       store.FileStorageSceneRequestAsset,
		Source:      "acceptance",
	})
	if err != nil {
		t.Fatalf("direct OSS upload: %v", err)
	}
	if calls.Load() != 2 {
		t.Fatalf("upload calls=%d, want one retry", calls.Load())
	}
	if !strings.HasPrefix(uploadedPath, "/easyai-ai-gateway/production/media/request_asset/2026/07/30/input-") ||
		!strings.HasSuffix(uploadedPath, ".png") {
		t.Fatalf("unexpected object path=%q", uploadedPath)
	}
	if got := stringFromAny(uploaded["url"]); got != "https://cdn.example.com"+uploadedPath {
		t.Fatalf("public URL=%q", got)
	}
	channel, _ := uploaded["storageChannel"].(map[string]any)
	if stringFromAny(channel["provider"]) != "aliyun_oss_direct" {
		t.Fatalf("storage channel=%+v", channel)
	}
}

func TestDirectOSSDoesNotReplaceGeneralUploadScene(t *testing.T) {
	var calls atomic.Int64
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		calls.Add(1)
		w.WriteHeader(http.StatusOK)
	}))
	defer server.Close()
	storageDir := t.TempDir()
	cfg := config.Config{
		LocalUploadedStorageDir: storageDir,
		MediaOSSDirectEnabled:   true,
		MediaOSSEndpoint:        server.URL,
		MediaOSSBucket:          "media-bucket",
		MediaOSSAccessKeyID:     "access-key",
		MediaOSSAccessKeySecret: "access-secret",
		MediaOSSPublicBaseURL:   "https://cdn.example.com",
		MediaOSSObjectPrefix:    "easyai-ai-gateway/production/media",
	}
	service := &Service{cfg: cfg, directOSS: newDirectOSSUploader(cfg)}
	uploaded, err := service.UploadFile(context.Background(), FileUploadPayload{
		Bytes:       []byte("general-upload"),
		ContentType: "text/plain",
		FileName:    "general.txt",
		Scene:       store.FileStorageSceneUpload,
	})
	if err != nil {
		t.Fatalf("general upload: %v", err)
	}
	if calls.Load() != 0 {
		t.Fatalf("direct OSS calls=%d for general upload scene", calls.Load())
	}
	channel, _ := uploaded["storageChannel"].(map[string]any)
	if stringFromAny(channel["provider"]) != "local_static" {
		t.Fatalf("general upload channel=%+v", channel)
	}
}

func TestGeneratedInlineAssetUsesDirectOSS(t *testing.T) {
	var calls atomic.Int64
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
		calls.Add(1)
		if request.Method != http.MethodPut {
			t.Errorf("method=%s, want PUT", request.Method)
		}
		w.WriteHeader(http.StatusOK)
	}))
	defer server.Close()

	cfg := config.Config{
		MediaOSSDirectEnabled:   true,
		MediaOSSEndpoint:        server.URL,
		MediaOSSBucket:          "media-bucket",
		MediaOSSAccessKeyID:     "access-key",
		MediaOSSAccessKeySecret: "access-secret",
		MediaOSSPublicBaseURL:   "https://cdn.example.com",
		MediaOSSObjectPrefix:    "easyai-ai-gateway/production/media",
	}
	service := &Service{cfg: cfg, directOSS: newDirectOSSUploader(cfg)}
	upload, contentType, kind, strategy, err := service.uploadGeneratedAsset(
		context.Background(),
		"task-direct-oss",
		&generatedInlineAsset{
			Bytes:       []byte{0x89, 'P', 'N', 'G', 0x0d, 0x0a, 0x1a, 0x0a},
			ContentType: "image/png",
			Kind:        "image",
			SourceKey:   "b64_json",
		},
		0,
		[]store.FileStorageChannel{{ID: "unused-channel"}},
	)
	if err != nil {
		t.Fatalf("direct generated upload: %v", err)
	}
	if calls.Load() != 1 {
		t.Fatalf("direct OSS calls=%d, want 1", calls.Load())
	}
	if contentType != "image/png" || kind != "image" || strategy != "direct_aliyun_oss" {
		t.Fatalf("generated metadata=%s/%s/%s", contentType, kind, strategy)
	}
	channel, _ := upload["storageChannel"].(map[string]any)
	if stringFromAny(channel["provider"]) != "aliyun_oss_direct" {
		t.Fatalf("storage channel=%+v", channel)
	}
}