platform/pkg/testutil/db.go

package testutil

import (
	"platform/pkg/orm"
	"platform/web/models"
	q "platform/web/queries"
	"testing"

	"github.com/glebarez/sqlite"
	"gorm.io/gorm"
)

// SetupDBTest 创建一个基于 SQLite 内存数据库的 GORM 连接
func SetupDBTest(t *testing.T) *gorm.DB {
	// 使用 SQLite 内存数据库
	gormDB, err := gorm.Open(sqlite.Open("file::memory:?cache=shared"), &gorm.Config{})
	if err != nil {
		t.Fatalf("gorm 打开 SQLite 内存数据库失败: %v", err)
	}

	// 自动迁移数据表结构
	err = gormDB.AutoMigrate(
		&models.User{},
		&models.Whitelist{},
		&models.Resource{},
		&models.ResourcePss{},
		&models.Proxy{},
		&models.Channel{},
	)
	if err != nil {
		t.Fatalf("自动迁移表结构失败: %v", err)
	}

	// 设置全局数据库连接
	q.SetDefault(gormDB)
	orm.DB = gormDB

	return gormDB
}