lzt/new-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Protect external token migration tests

Browse Source
This commit is contained in:
XiaoAI1024
2026-04-23 13:29:00 +08:00
parent 0feb6f2c3c
commit 49474520ec
1 changed files with 20 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
controller/token_test.go
+20 -12
View File
@@ -114,7 +114,7 @@ func setupTokenControllerTestDB(t *testing.T) *gorm.DB {
return db return db
} }
func openTokenControllerExternalDB(t *testing.T, dialect string, dsn string) *gorm.DB { func openTokenControllerExternalDB(t *testing.T, dialect string, dsn string) (*gorm.DB, *bool) {
t.Helper() t.Helper()
gin.SetMode(gin.TestMode) gin.SetMode(gin.TestMode)
@@ -142,15 +142,23 @@ func openTokenControllerExternalDB(t *testing.T, dialect string, dsn string) *go
model.DB = db model.DB = db
model.LOG_DB = db model.LOG_DB = db
if db.Migrator().HasTable("tokens") {
t.Skipf("refusing to run %s migration compatibility test against external database because tokens table already exists", dialect)
}
managedTokensTable := new(bool)
t.Cleanup(func() { t.Cleanup(func() {
_ = db.Exec("DROP TABLE IF EXISTS tokens").Error if *managedTokensTable && db.Migrator().HasTable("tokens") {
_ = db.Migrator().DropTable("tokens")
}
sqlDB, err := db.DB() sqlDB, err := db.DB()
if err == nil { if err == nil {
_ = sqlDB.Close() _ = sqlDB.Close()
} }
}) })
return db return db, managedTokensTable
} }
func seedToken(t *testing.T, db *gorm.DB, userID int, name string, rawKey string) *model.Token { func seedToken(t *testing.T, db *gorm.DB, userID int, name string, rawKey string) *model.Token {
@@ -266,18 +274,18 @@ func getTokenKeyColumnType(t *testing.T, db *gorm.DB, dialect string) string {
} }
} }
func runTokenMigrationCompatibilityTest(t *testing.T, db *gorm.DB, dialect string) { func runTokenMigrationCompatibilityTest(t *testing.T, db *gorm.DB, dialect string, managedTokensTable *bool) {
t.Helper() t.Helper()
legacyKey := strings.Repeat("a", 48) legacyKey := strings.Repeat("a", 48)
longKey := strings.Repeat("b", 64) longKey := strings.Repeat("b", 64)
if err := db.Exec("DROP TABLE IF EXISTS tokens").Error; err != nil {
t.Fatalf("failed to drop pre-existing token table: %v", err)
}
if err := db.AutoMigrate(&legacyToken{}); err != nil { if err := db.AutoMigrate(&legacyToken{}); err != nil {
t.Fatalf("failed to create legacy token schema: %v", err) t.Fatalf("failed to create legacy token schema: %v", err)
} }
if managedTokensTable != nil {
*managedTokensTable = true
}
if err := db.Create(&legacyToken{ if err := db.Create(&legacyToken{
UserId: 7, UserId: 7,
Key: legacyKey, Key: legacyKey,
@@ -359,7 +367,7 @@ func TestTokenAutoMigrateUsesVarchar128KeyColumn(t *testing.T) {
func TestTokenMigrationFromChar48ToVarchar128(t *testing.T) { func TestTokenMigrationFromChar48ToVarchar128(t *testing.T) {
db := openTokenControllerTestDB(t) db := openTokenControllerTestDB(t)
runTokenMigrationCompatibilityTest(t, db, "sqlite") runTokenMigrationCompatibilityTest(t, db, "sqlite", nil)
} }
func TestTokenMigrationFromChar48ToVarchar128MySQL(t *testing.T) { func TestTokenMigrationFromChar48ToVarchar128MySQL(t *testing.T) {
@@ -368,8 +376,8 @@ func TestTokenMigrationFromChar48ToVarchar128MySQL(t *testing.T) {
t.Skip("set TEST_MYSQL_DSN to run mysql migration compatibility test") t.Skip("set TEST_MYSQL_DSN to run mysql migration compatibility test")
} }
db := openTokenControllerExternalDB(t, "mysql", dsn) db, managedTokensTable := openTokenControllerExternalDB(t, "mysql", dsn)
runTokenMigrationCompatibilityTest(t, db, "mysql") runTokenMigrationCompatibilityTest(t, db, "mysql", managedTokensTable)
} }
func TestTokenMigrationFromChar48ToVarchar128Postgres(t *testing.T) { func TestTokenMigrationFromChar48ToVarchar128Postgres(t *testing.T) {
@@ -378,8 +386,8 @@ func TestTokenMigrationFromChar48ToVarchar128Postgres(t *testing.T) {
t.Skip("set TEST_POSTGRES_DSN to run postgres migration compatibility test") t.Skip("set TEST_POSTGRES_DSN to run postgres migration compatibility test")
} }
db := openTokenControllerExternalDB(t, "postgres", dsn) db, managedTokensTable := openTokenControllerExternalDB(t, "postgres", dsn)
runTokenMigrationCompatibilityTest(t, db, "postgres") runTokenMigrationCompatibilityTest(t, db, "postgres", managedTokensTable)
} }
func TestGetAllTokensMasksKeyInResponse(t *testing.T) { func TestGetAllTokensMasksKeyInResponse(t *testing.T) {