new-api/controller/channel_asset_credential_test.go

package controller

import (
	"testing"

	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/model"
	"github.com/glebarez/sqlite"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

func setupChannelAssetCredentialControllerDB(t *testing.T) *gorm.DB {
	t.Helper()
	db, err := gorm.Open(sqlite.Open("file:controller_channel_asset_credentials?mode=memory&cache=shared"), &gorm.Config{
		Logger: logger.Default.LogMode(logger.Silent),
	})
	require.NoError(t, err)
	sqlDB, err := db.DB()
	require.NoError(t, err)
	sqlDB.SetMaxOpenConns(1)
	originalDB := model.DB
	model.DB = db
	require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.ChannelAssetCredential{}))
	t.Cleanup(func() {
		model.DB = originalDB
		require.NoError(t, sqlDB.Close())
	})
	return db
}

func TestAttachChannelAssetCredentialSummariesDoesNotExposeSecrets(t *testing.T) {
	db := setupChannelAssetCredentialControllerDB(t)
	accessKey := "ak-" + t.Name()
	secretKey := "sk-" + t.Name()
	channel := &model.Channel{Id: 61, Type: constant.ChannelTypeChinaMobileSeedance, Key: "video-key", Name: "channel"}
	require.NoError(t, db.Create(channel).Error)
	require.NoError(t, model.UpsertChannelAssetCredential(&model.ChannelAssetCredential{
		ChannelId: 61,
		AccessKey: accessKey,
		SecretKey: secretKey,
		PoolID:    "pool-61",
	}))

	require.NoError(t, attachChannelAssetCredentialSummaries([]*model.Channel{channel}))
	assert.True(t, channel.AssetCredentialConfigured)
	assert.Equal(t, "pool-61", channel.AssetCredentialPoolID)
	assert.NotContains(t, channel.Key, accessKey)
	assert.NotContains(t, channel.Key, secretKey)
}

func TestAttachChannelAssetCredentialSummariesIncludesDoubaoVideo(t *testing.T) {
	db := setupChannelAssetCredentialControllerDB(t)
	channel := &model.Channel{Id: 54, Type: constant.ChannelTypeDoubaoVideo, Key: "video-key", Name: "doubao"}
	require.NoError(t, db.Create(channel).Error)
	require.NoError(t, model.UpsertChannelAssetCredential(&model.ChannelAssetCredential{
		ChannelId: 54,
		AccessKey: "ak-" + t.Name(),
		SecretKey: "sk-" + t.Name(),
		BaseURL:   "http://example.com/openApi/portrait",
	}))

	require.NoError(t, attachChannelAssetCredentialSummaries([]*model.Channel{channel}))
	assert.True(t, channel.AssetCredentialConfigured)
	assert.Equal(t, "http://example.com/openApi/portrait", channel.AssetCredentialBaseURL)
}