Files

91 lines
3.4 KiB
Go

package version
import (
"os"
"testing"
"gorm.io/driver/postgres"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"git.ilapage.cn/ila/yovision/Sense/server/app/sense/quota"
migrationModels "git.ilapage.cn/ila/yovision/Sense/server/cmd/migrate/migration/models"
common "git.ilapage.cn/ila/yovision/Sense/server/common/models"
)
func prepareQuotaMigrationDB(t *testing.T, db *gorm.DB) {
t.Helper()
if err := db.AutoMigrate(&migrationModels.SysRole{}, &migrationModels.SysMenu{}, &deviceCasbinRule{}, &common.Migration{}); err != nil {
t.Fatal(err)
}
for _, role := range []string{"implementation_operator", "site_admin", "viewer"} {
if err := db.Create(&migrationModels.SysRole{RoleName: role, RoleKey: role, Status: "2"}).Error; err != nil {
t.Fatal(err)
}
}
if err := db.Create(&migrationModels.SysMenu{MenuName: "SenseDeviceManage", Title: "设备管理", Path: "devices", MenuType: "C", Component: "/sense/device/index"}).Error; err != nil {
t.Fatal(err)
}
}
func assertQuotaMigration(t *testing.T, db *gorm.DB, version string) {
t.Helper()
var setting quota.Setting
if err := db.First(&setting, "id = ?", quota.SettingID).Error; err != nil {
t.Fatal(err)
}
var menus, reads, updates, enables, viewerWrites, applied int64
db.Model(&migrationModels.SysMenu{}).Where("menu_name IN ?", []string{"SenseQuota", "SenseQuotaUpdate", "SenseDeviceEnable"}).Count(&menus)
db.Model(&deviceCasbinRule{}).Where("v1 = ? AND v2 = ?", "/api/v1/quota", "GET").Count(&reads)
db.Model(&deviceCasbinRule{}).Where("v1 = ? AND v2 = ?", "/api/v1/quota", "PUT").Count(&updates)
db.Model(&deviceCasbinRule{}).Where("v1 = ? AND v2 = ?", "/api/v1/devices/:id/enable", "PUT").Count(&enables)
db.Model(&deviceCasbinRule{}).Where("v0 = ? AND v2 = ?", "viewer", "PUT").Count(&viewerWrites)
db.Model(&common.Migration{}).Where("version = ?", version).Count(&applied)
if setting.Limit != 16 || menus != 3 || reads != 3 || updates != 1 || enables != 1 || viewerWrites != 0 || applied != 1 {
t.Fatalf("setting=%#v menus=%d reads=%d updates=%d enables=%d viewerWrites=%d applied=%d", setting, menus, reads, updates, enables, viewerWrites, applied)
}
}
func TestQuotaMigrationAddsDefaultAndLeastPrivilegeMenus(t *testing.T) {
t.Setenv("SENSE_PROVISIONING_QUOTA", "")
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
if err != nil {
t.Fatal(err)
}
prepareQuotaMigrationDB(t, db)
const version = "2026082812000_quota.go"
if err = migrateSenseQuota(db, version); err != nil {
t.Fatal(err)
}
assertQuotaMigration(t, db, version)
}
func TestQuotaMigrationOnPostgres(t *testing.T) {
dsn := os.Getenv("SENSE_QUOTA_MIGRATION_TEST_DATABASE_URL")
if dsn == "" {
t.Skip("set SENSE_QUOTA_MIGRATION_TEST_DATABASE_URL to run the PostgreSQL migration test")
}
t.Setenv("SENSE_PROVISIONING_QUOTA", "")
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
if err != nil {
t.Fatal(err)
}
const schema = "sense_quota_75_test"
if err = db.Exec("DROP SCHEMA IF EXISTS " + schema + " CASCADE").Error; err != nil {
t.Fatal(err)
}
if err = db.Exec("CREATE SCHEMA " + schema).Error; err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = db.Exec("DROP SCHEMA IF EXISTS " + schema + " CASCADE").Error })
if err = db.Exec("SET search_path TO " + schema).Error; err != nil {
t.Fatal(err)
}
prepareQuotaMigrationDB(t, db)
const version = "2026082812000_quota.go"
if err = migrateSenseQuota(db, version); err != nil {
t.Fatal(err)
}
assertQuotaMigration(t, db, version)
}