ai-agent/internal/services/conversation_vision_asset_test.go

package services

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

	"code.tczkiot.com/wlw/ai-agent/internal/models"
	"code.tczkiot.com/wlw/ai-agent/internal/pkg/config"
	"code.tczkiot.com/wlw/ai-agent/internal/pkg/enums"
	"code.tczkiot.com/wlw/ai-agent/internal/services/storage"

	"github.com/glebarez/sqlite"
	"github.com/mlogclub/simple/sqls"
	"gorm.io/gorm"
	"gorm.io/gorm/schema"
)

func setupConversationVisionAssetTest(t *testing.T) *gorm.DB {
	t.Helper()
	database, err := gorm.Open(sqlite.Open("file:"+strings.ReplaceAll(t.Name(), "/", "_")+"?mode=memory&cache=shared"), &gorm.Config{
		NamingStrategy: schema.NamingStrategy{TablePrefix: "t_", SingularTable: true},
	})
	if err != nil {
		t.Fatalf("open sqlite: %v", err)
	}
	if err := database.AutoMigrate(&models.Asset{}); err != nil {
		t.Fatalf("migrate asset: %v", err)
	}
	sqls.SetDB(database)
	config.SetCurrent(&config.Config{Storage: config.StorageConfig{
		Default: enums.AssetProviderLocal, MaxUploadSizeMB: 20,
		Local: config.LocalStorageConfig{Root: t.TempDir(), BaseURL: "/storage"},
	}})
	storage.SetHostStorage(nil)
	t.Cleanup(func() { storage.SetHostStorage(nil) })
	return database
}

func TestConversationVisionImagesAreInlineAndConversationScoped(t *testing.T) {
	setupConversationVisionAssetTest(t)
	png := append([]byte("\x89PNG\r\n\x1a\n"), make([]byte, 520)...)
	asset, err := AssetService.UploadConversationBytes(png, "images", "device.png", 11, nil)
	if err != nil {
		t.Fatalf("upload conversation image: %v", err)
	}
	if asset.ConversationID != 11 || asset.MimeType != "image/png" {
		t.Fatalf("unexpected stored asset: %#v", asset)
	}
	payload := fmt.Sprintf(`{"asset_id":%q,"url":"https://attacker.invalid/ssrf.png","provider":"oss","storage_key":"other/customer.png"}`, asset.AssetID)
	message := models.Message{
		ID: 9, ConversationID: 11, SenderType: enums.IMSenderTypeCustomer,
		MessageType: enums.IMMessageTypeImage, Payload: payload,
		SendStatus: enums.IMMessageStatusSent,
	}
	images := AssetService.LoadConversationVisionImages(11, []models.Message{message}, 3)
	if len(images) != 1 {
		t.Fatalf("images = %#v, want one trusted image", images)
	}
	decoded, err := base64.StdEncoding.DecodeString(images[0].Base64Data)
	if err != nil || string(decoded) != string(png) {
		t.Fatalf("inline image mismatch: len=%d err=%v", len(decoded), err)
	}
	if images[0].MIMEType != "image/png" {
		t.Fatalf("mime type = %q", images[0].MIMEType)
	}
	if got := AssetService.LoadConversationVisionImages(12, []models.Message{{
		ID: 10, ConversationID: 12, SenderType: enums.IMSenderTypeCustomer,
		MessageType: enums.IMMessageTypeImage, Payload: payload, SendStatus: enums.IMMessageStatusSent,
	}}, 3); len(got) != 0 {
		t.Fatalf("cross-conversation asset leaked into model input: %#v", got)
	}
}

func TestConversationVisionImagesLoadsCompositeMessageInPayloadOrder(t *testing.T) {
	setupConversationVisionAssetTest(t)
	firstData := append([]byte("\x89PNG\r\n\x1a\n"), make([]byte, 520)...)
	secondData := append([]byte("\x89PNG\r\n\x1a\n"), make([]byte, 640)...)
	first, err := AssetService.UploadConversationBytes(firstData, "images", "front.png", 21, nil)
	if err != nil {
		t.Fatalf("upload first image: %v", err)
	}
	second, err := AssetService.UploadConversationBytes(secondData, "images", "label.png", 21, nil)
	if err != nil {
		t.Fatalf("upload second image: %v", err)
	}
	payload, err := buildIMMessageAssetBatchPayload([]*models.Asset{first, second})
	if err != nil {
		t.Fatalf("build batch payload: %v", err)
	}
	images := AssetService.LoadConversationVisionImages(21, []models.Message{{
		ID: 22, ConversationID: 21, SenderType: enums.IMSenderTypeCustomer,
		MessageType: enums.IMMessageTypeImage, Payload: payload, SendStatus: enums.IMMessageStatusSent,
	}}, 6)
	if len(images) != 2 || images[0].AssetID != first.AssetID || images[1].AssetID != second.AssetID {
		t.Fatalf("composite image order mismatch: %#v", images)
	}
}

func TestValidateConversationAssetRejectsCrossConversationAndFakeImage(t *testing.T) {
	asset := &models.Asset{ConversationID: 7, Status: enums.AssetStatusSuccess, MimeType: "image/png"}
	if err := validateConversationAsset(asset, 7, enums.IMMessageTypeImage); err != nil {
		t.Fatalf("valid scoped image rejected: %v", err)
	}
	if err := validateConversationAsset(asset, 8, enums.IMMessageTypeImage); err == nil {
		t.Fatal("cross-conversation asset must be rejected")
	}
	asset.MimeType = "text/html"
	if err := validateConversationAsset(asset, 7, enums.IMMessageTypeImage); err == nil {
		t.Fatal("non-image asset must not be sent as an image")
	}
}