ai-agent/internal/pkg/utils/message_test.go

package utils

import (
	"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"
	"strings"
	"testing"
	"time"

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

func TestBuildIMMessageAssetPayloadForResponseAddsSignedURL(t *testing.T) {
	config.SetCurrent(&config.Config{
		Storage: config.StorageConfig{
			Default: enums.AssetProviderLocal,
			Local: config.LocalStorageConfig{
				BaseURL: "https://files.example.com",
			},
		},
	})

	payload := `{"asset_id":"asset_1","provider":"local","storage_key":"attachments/demo.png","filename":"demo.png"}`
	got := buildIMMessageAssetPayloadForResponse(payload)

	if !strings.Contains(got, `"provider":"local"`) {
		t.Fatalf("expected provider in payload, got: %s", got)
	}
	if !strings.Contains(got, `"storage_key":"attachments/demo.png"`) {
		t.Fatalf("expected storage_key in payload, got: %s", got)
	}
	if !strings.Contains(got, `"url":"https://files.example.com/attachments/demo.png"`) {
		t.Fatalf("expected signed url in payload, got: %s", got)
	}
}

func TestSanitizeMessageHTMLStripsStoredSrcForManagedImages(t *testing.T) {
	html := `<p><img src="https://files.example.com/demo.png" data-provider="local" data-storage-key="attachments/demo.png" alt="demo"></p>`

	got := SanitizeMessageHTML(html)

	if strings.Contains(got, `src=`) {
		t.Fatalf("expected src removed from stored html, got: %s", got)
	}
	if !strings.Contains(got, `data-provider="local"`) {
		t.Fatalf("expected data-provider kept, got: %s", got)
	}
	if !strings.Contains(got, `data-storage-key="attachments/demo.png"`) {
		t.Fatalf("expected data-storage-key kept, got: %s", got)
	}
}

func TestBuildMessageHTMLForResponseAddsSignedURL(t *testing.T) {
	config.SetCurrent(&config.Config{
		Storage: config.StorageConfig{
			Default: enums.AssetProviderLocal,
			Local: config.LocalStorageConfig{
				BaseURL: "https://files.example.com",
			},
		},
	})

	html := `<p><img data-provider="local" data-storage-key="attachments/demo.png" alt="demo"></p>`
	got := BuildMessageHTMLForResponse(html)

	if !strings.Contains(got, `src="https://files.example.com/attachments/demo.png"`) {
		t.Fatalf("expected signed src in response html, got: %s", got)
	}
}

func TestBuildRuntimeMessageTextForHTML(t *testing.T) {
	got := BuildRuntimeMessageText(enums.IMMessageTypeHTML, `<p>你好</p><p><img data-provider="local" data-storage-key="images/demo.png" alt="demo"></p>`)
	if got != "你好 [图片]" {
		t.Fatalf("expected html converted to plain text summary, got: %q", got)
	}
}

func TestBuildRuntimeMessageTextForAssetMessages(t *testing.T) {
	if got := BuildRuntimeMessageText(enums.IMMessageTypeImage, "demo.png"); got != "[图片] demo.png" {
		t.Fatalf("unexpected image runtime text: %q", got)
	}
	if got := BuildRuntimeMessageText(enums.IMMessageTypeAttachment, "spec.pdf"); got != "[附件] spec.pdf" {
		t.Fatalf("unexpected attachment runtime text: %q", got)
	}
}

func TestBuildRenderableMessageTransformsPayloadAndHTML(t *testing.T) {
	config.SetCurrent(&config.Config{
		Storage: config.StorageConfig{
			Default: enums.AssetProviderLocal,
			Local: config.LocalStorageConfig{
				BaseURL: "https://files.example.com",
			},
		},
	})

	image := &models.Message{
		MessageType: enums.IMMessageTypeImage,
		Payload:     `{"asset_id":"asset_1","provider":"local","storage_key":"attachments/demo.png","filename":"demo.png"}`,
	}
	_, imagePayload := BuildRenderableMessage(image)
	if !strings.Contains(imagePayload, `"url":"https://files.example.com/attachments/demo.png"`) {
		t.Fatalf("expected image payload signed url, got: %s", imagePayload)
	}

	htmlMsg := &models.Message{
		MessageType: enums.IMMessageTypeHTML,
		Content:     `<p><img data-provider="local" data-storage-key="attachments/demo.png"></p>`,
	}
	htmlContent, _ := BuildRenderableMessage(htmlMsg)
	if !strings.Contains(htmlContent, `src="https://files.example.com/attachments/demo.png"`) {
		t.Fatalf("expected html content signed src, got: %s", htmlContent)
	}
}

func TestNormalizeMessageHTMLAssetsKeepsValidAttrsAndRemovesSrc(t *testing.T) {
	setupMessageTestDB(t)
	config.SetCurrent(&config.Config{
		Storage: config.StorageConfig{
			Default: enums.AssetProviderLocal,
			Local: config.LocalStorageConfig{
				BaseURL: "https://files.example.com",
			},
		},
	})
	createTestAsset(t, &models.Asset{
		AssetID:    "asset_local_1",
		Provider:   enums.AssetProviderLocal,
		StorageKey: "images/demo.png",
		Filename:   "demo.png",
		FileSize:   123,
		MimeType:   "image/png",
		Status:     enums.AssetStatusSuccess,
	})

	got, err := NormalizeMessageHTMLAssets(`<p><img src="https://files.example.com/images/demo.png" data-asset-id="asset_local_1" data-provider="local" data-storage-key="images/demo.png" alt="demo"></p>`)
	if err != nil {
		t.Fatalf("expected normalization success, got error: %v", err)
	}

	if !strings.Contains(got, `data-asset-id="asset_local_1"`) {
		t.Fatalf("expected data-asset-id added, got: %s", got)
	}
	if !strings.Contains(got, `data-provider="local"`) {
		t.Fatalf("expected data-provider added, got: %s", got)
	}
	if !strings.Contains(got, `data-storage-key="images/demo.png"`) {
		t.Fatalf("expected data-storage-key added, got: %s", got)
	}
	if strings.Contains(got, `src=`) {
		t.Fatalf("expected src removed after asset binding, got: %s", got)
	}
}

func TestNormalizeMessageHTMLAssetsRejectsMissingAssetMetadata(t *testing.T) {
	setupMessageTestDB(t)
	config.SetCurrent(&config.Config{
		Storage: config.StorageConfig{
			Default: enums.AssetProviderLocal,
			Local: config.LocalStorageConfig{
				BaseURL: "https://files.example.com",
			},
		},
	})

	_, err := NormalizeMessageHTMLAssets(`<p><img src="https://unknown.example.com/demo.png" alt="demo"></p>`)
	if err == nil {
		t.Fatalf("expected missing image asset metadata rejected")
	}
}

func TestNormalizeMessageHTMLAssetsRejectsIncompleteAttrs(t *testing.T) {
	setupMessageTestDB(t)
	config.SetCurrent(&config.Config{
		Storage: config.StorageConfig{
			Default: enums.AssetProviderLocal,
			Local: config.LocalStorageConfig{
				BaseURL: "https://files.example.com",
			},
		},
	})

	_, err := NormalizeMessageHTMLAssets(`<p><img data-asset-id="asset1" data-provider="local" alt="demo"></p>`)
	if err == nil {
		t.Fatalf("expected incomplete asset attrs rejected")
	}
}

func TestNormalizeMessageHTMLAssetsRejectsMismatchedAttrs(t *testing.T) {
	setupMessageTestDB(t)
	config.SetCurrent(&config.Config{
		Storage: config.StorageConfig{
			Default: enums.AssetProviderLocal,
			Local: config.LocalStorageConfig{
				BaseURL: "https://files.example.com",
			},
		},
	})
	createTestAsset(t, &models.Asset{
		AssetID:    "asset_local_2",
		Provider:   enums.AssetProviderLocal,
		StorageKey: "images/real.png",
		Filename:   "real.png",
		FileSize:   456,
		MimeType:   "image/png",
		Status:     enums.AssetStatusSuccess,
	})

	_, err := NormalizeMessageHTMLAssets(`<p><img data-asset-id="asset_local_2" data-provider="local" data-storage-key="images/wrong.png" alt="demo"></p>`)
	if err == nil {
		t.Fatalf("expected mismatched asset attrs rejected")
	}
}

func setupMessageTestDB(t *testing.T) {
	t.Helper()
	db, err := gorm.Open(sqlite.Open("file::memory:?cache=shared"), &gorm.Config{})
	if err != nil {
		t.Fatalf("open sqlite failed: %v", err)
	}
	if err := db.AutoMigrate(&models.Asset{}); err != nil {
		t.Fatalf("auto migrate asset failed: %v", err)
	}
	sqls.SetDB(db)
}

func createTestAsset(t *testing.T, item *models.Asset) {
	t.Helper()
	now := time.Now()
	if item.CreatedAt.IsZero() {
		item.CreatedAt = now
	}
	if item.UpdatedAt.IsZero() {
		item.UpdatedAt = now
	}
	if err := sqls.DB().Create(item).Error; err != nil {
		t.Fatalf("create asset failed: %v", err)
	}
}