refactor(tests): streamline database connection handling in dockertest integration tests

- Removed unused database connection variables for MySQL and PostgreSQL.
- Simplified database connection setup by directly using the connection returned from gorm.Open.
- Enhanced cleanup logic to ensure proper closure of database connections after tests.
- Improved error handling during database setup and migration processes for clearer test failure reporting.

Author: asjdf

cache/cache.go

@@ -149,11 +149,14 @@ func (c *Gorm2Cache) SearchKeyExists(ctx context.Context, tableName string, SQL
 }
 
 func (c *Gorm2Cache) BatchSetPrimaryKeyCache(ctx context.Context, tableName string, kvs []util.Kv) error {
-	for idx, kv := range kvs {
-		// kv.Key 已经是最终格式（单个值或已用":"连接的联合主键），直接传入
-		kvs[idx].Key = util.GenPrimaryCacheKey(c.InstanceId, tableName, kv.Key)
+	cacheKvs := make([]util.Kv, 0, len(kvs))
+	for _, kv := range kvs {
+		cacheKvs = append(cacheKvs, util.Kv{
+			Key:   util.GenPrimaryCacheKey(c.InstanceId, tableName, kv.Key),
+			Value: kv.Value,
+		})
 	}
-	return c.cache.BatchSetKeys(ctx, kvs)
+	return c.cache.BatchSetKeys(ctx, cacheKvs)
 }
 
 func (c *Gorm2Cache) SetSearchCache(ctx context.Context, cacheValue string, tableName string,
@@ -189,14 +192,16 @@ func (c *Gorm2Cache) BatchGetUniqueCache(ctx context.Context, tableName string,
 	return c.cache.BatchGetValues(ctx, cacheKeys)
 }
 
-// BatchSetUniqueCache 批量设置unique键缓存。
-// 注意：会原地修改 kvs 中每个元素的 Key 为完整缓存 key，调用方不应再使用传入的 kvs。
+// BatchSetUniqueCache 批量设置 unique 键缓存。不会修改调用方传入的 kvs。
 func (c *Gorm2Cache) BatchSetUniqueCache(ctx context.Context, tableName string, uniqueIndexName string, kvs []util.Kv) error {
-	for idx, kv := range kvs {
-		// kv.Key 已经是最终格式（单个值或已用":"连接的联合unique键），直接传入
-		kvs[idx].Key = util.GenUniqueCacheKey(c.InstanceId, tableName, uniqueIndexName, kv.Key)
+	cacheKvs := make([]util.Kv, 0, len(kvs))
+	for _, kv := range kvs {
+		cacheKvs = append(cacheKvs, util.Kv{
+			Key:   util.GenUniqueCacheKey(c.InstanceId, tableName, uniqueIndexName, kv.Key),
+			Value: kv.Value,
+		})
 	}
-	return c.cache.BatchSetKeys(ctx, kvs)
+	return c.cache.BatchSetKeys(ctx, cacheKvs)
 }
 
 // InvalidateUniqueCache 失效unique键缓存

cache/dockertest_integration_test.go

@@ -58,15 +58,13 @@ func (UserSession) TableName() string {
 var (
 	mysqlPool        *dockertest.Pool
 	mysqlResource    *dockertest.Resource
-	mysqlDB          *gorm.DB
 	mysqlDSN         string
 	setupMySQLOnce   sync.Once
 	cleanupMySQLOnce sync.Once
 	mysqlSetupErr    error
 
 	pgPool        *dockertest.Pool
 	pgResource    *dockertest.Resource
-	pgDB          *gorm.DB
 	pgDSN         string
 	setupPGOnce   sync.Once
 	cleanupPGOnce sync.Once
@@ -106,21 +104,19 @@ func setupMySQL(t *testing.T) *gorm.DB {
 
 		mysqlPool.MaxWait = 120 * time.Second
 		if err := mysqlPool.Retry(func() error {
-			var openErr error
-			mysqlDB, openErr = gorm.Open(mysql.Open(mysqlDSN), &gorm.Config{
+			conn, openErr := gorm.Open(mysql.Open(mysqlDSN), &gorm.Config{
 				Logger: logger.Default.LogMode(logger.Silent),
 			})
 			if openErr != nil {
 				return openErr
 			}
-			sqlDB, openErr := mysqlDB.DB()
+			sqlDB, openErr := conn.DB()
 			if openErr != nil {
 				return openErr
 			}
-			sqlDB.SetMaxOpenConns(10)
-			sqlDB.SetMaxIdleConns(5)
-			sqlDB.SetConnMaxLifetime(time.Hour)
-			return sqlDB.Ping()
+			err := sqlDB.Ping()
+			_ = sqlDB.Close()
+			return err
 		}); err != nil {
 			mysqlSetupErr = fmt.Errorf("could not connect to MySQL: %w", err)
 			return
@@ -129,22 +125,31 @@ func setupMySQL(t *testing.T) *gorm.DB {
 	if mysqlSetupErr != nil {
 		t.Fatalf("MySQL setup failed: %v", mysqlSetupErr)
 	}
-	if mysqlDB == nil {
-		t.Fatal("MySQL DB is nil after setup")
+
+	db, err := gorm.Open(mysql.Open(mysqlDSN), &gorm.Config{
+		Logger: logger.Default.LogMode(logger.Silent),
+	})
+	if err != nil {
+		t.Fatalf("open MySQL DB failed: %v", err)
 	}
+	sqlDB, err := db.DB()
+	if err != nil {
+		t.Fatalf("get sql.DB from gorm failed: %v", err)
+	}
+	sqlDB.SetMaxOpenConns(10)
+	sqlDB.SetMaxIdleConns(5)
+	sqlDB.SetConnMaxLifetime(time.Hour)
+	t.Cleanup(func() {
+		_ = sqlDB.Close()
+	})
 
-	// Clean tables before each test
-	if err := mysqlDB.Exec("DROP TABLE IF EXISTS user_roles, users, user_sessions").Error; err != nil {
+	if err := db.Exec("DROP TABLE IF EXISTS user_roles, users, user_sessions").Error; err != nil {
 		t.Fatalf("failed to drop MySQL tables: %v", err)
 	}
-
-	// Auto migrate
-	err := mysqlDB.AutoMigrate(&UserRole{}, &User{}, &UserSession{})
-	if err != nil {
+	if err := db.AutoMigrate(&UserRole{}, &User{}, &UserSession{}); err != nil {
 		t.Fatalf("Auto migrate error: %v", err)
 	}
-
-	return mysqlDB
+	return db
 }
 
 func setupPostgreSQL(t *testing.T) *gorm.DB {
@@ -188,21 +193,19 @@ func setupPostgreSQL(t *testing.T) *gorm.DB {
 
 		pgPool.MaxWait = 120 * time.Second
 		if err := pgPool.Retry(func() error {
-			var openErr error
-			pgDB, openErr = gorm.Open(postgres.Open(pgDSN), &gorm.Config{
+			conn, openErr := gorm.Open(postgres.Open(pgDSN), &gorm.Config{
 				Logger: logger.Default.LogMode(logger.Silent),
 			})
 			if openErr != nil {
 				return openErr
 			}
-			sqlDB, openErr := pgDB.DB()
+			sqlDB, openErr := conn.DB()
 			if openErr != nil {
 				return openErr
 			}
-			sqlDB.SetMaxOpenConns(10)
-			sqlDB.SetMaxIdleConns(5)
-			sqlDB.SetConnMaxLifetime(time.Hour)
-			return sqlDB.Ping()
+			err := sqlDB.Ping()
+			_ = sqlDB.Close()
+			return err
 		}); err != nil {
 			pgSetupErr = fmt.Errorf("could not connect to PostgreSQL: %w", err)
 			return
@@ -211,47 +214,43 @@ func setupPostgreSQL(t *testing.T) *gorm.DB {
 	if pgSetupErr != nil {
 		t.Fatalf("PostgreSQL setup failed: %v", pgSetupErr)
 	}
-	if pgDB == nil {
-		t.Fatal("PostgreSQL DB is nil after setup")
+
+	db, err := gorm.Open(postgres.Open(pgDSN), &gorm.Config{
+		Logger: logger.Default.LogMode(logger.Silent),
+	})
+	if err != nil {
+		t.Fatalf("open PostgreSQL DB failed: %v", err)
+	}
+	sqlDB, err := db.DB()
+	if err != nil {
+		t.Fatalf("get sql.DB from gorm failed: %v", err)
 	}
+	sqlDB.SetMaxOpenConns(10)
+	sqlDB.SetMaxIdleConns(5)
+	sqlDB.SetConnMaxLifetime(time.Hour)
+	t.Cleanup(func() {
+		_ = sqlDB.Close()
+	})
 
-	// Clean tables before each test
-	if err := pgDB.Exec("DROP TABLE IF EXISTS user_roles, users, user_sessions CASCADE").Error; err != nil {
+	if err := db.Exec("DROP TABLE IF EXISTS user_roles, users, user_sessions CASCADE").Error; err != nil {
 		t.Fatalf("failed to drop PostgreSQL tables: %v", err)
 	}
-
-	// Auto migrate
-	err := pgDB.AutoMigrate(&UserRole{}, &User{}, &UserSession{})
-	if err != nil {
+	if err := db.AutoMigrate(&UserRole{}, &User{}, &UserSession{}); err != nil {
 		t.Fatalf("Auto migrate error: %v", err)
 	}
-
-	return pgDB
+	return db
 }
 
 func TestMain(m *testing.M) {
 	code := m.Run()
 
-	// Cleanup
+	// Cleanup: 每测试独立 *gorm.DB 由 t.Cleanup 关闭，此处仅回收容器
 	cleanupMySQLOnce.Do(func() {
-		if mysqlDB != nil {
-			sqlDB, _ := mysqlDB.DB()
-			if sqlDB != nil {
-				sqlDB.Close()
-			}
-		}
 		if mysqlResource != nil && mysqlPool != nil {
 			_ = mysqlPool.Purge(mysqlResource)
 		}
 	})
-
 	cleanupPGOnce.Do(func() {
-		if pgDB != nil {
-			sqlDB, _ := pgDB.DB()
-			if sqlDB != nil {
-				sqlDB.Close()
-			}
-		}
 		if pgResource != nil && pgPool != nil {
 			_ = pgPool.Purge(pgResource)
 		}
@@ -622,7 +621,9 @@ func TestCacheStats_MySQL(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Failed to create cache: %v", err)
 	}
-	db.Use(cache)
+	if err := db.Use(cache); err != nil {
+		t.Fatalf("failed to register cache plugin: %v", err)
+	}
 
 	// 创建测试数据
 	user := User{Email: "test@example.com", Username: "test", Name: "Test User"}

cache/query.go

@@ -429,7 +429,17 @@ func (h *queryHandler) AfterQuery() func(db *gorm.DB) {
 						}
 					}
 				}()
-				if !cache.Config.AsyncWrite {
+				if cache.Config.AsyncWrite {
+					// 异步写时不在主路径 Wait，由后台 goroutine 在写完后 cancel，避免 fillCallAfterQuery 提前 cancel 导致写缓存被中止
+					if cancelObj, hasCancel := db.InstanceGet("gorm:cache:query:single_flight_cancel"); hasCancel {
+						if cancel, ok := cancelObj.(context.CancelFunc); ok {
+							go func() {
+								wg.Wait()
+								cancel()
+							}()
+						}
+					}
+				} else {
 					wg.Wait()
 				}
 				return
@@ -491,9 +501,12 @@ func (h *queryHandler) fillCallAfterQuery(db *gorm.DB) {
 		return
 	}
 	// 释放 singleflight leader 使用的 background context，避免 timer 泄漏
-	if cancelObj, hasCancel := db.InstanceGet("gorm:cache:query:single_flight_cancel"); hasCancel {
-		if cancel, ok := cancelObj.(context.CancelFunc); ok {
-			cancel()
+	// AsyncWrite 时由 AfterQuery 内启动的 goroutine 在 wg.Wait() 后 cancel，此处不 cancel 以免提前中止异步写缓存
+	if !h.cache.Config.AsyncWrite {
+		if cancelObj, hasCancel := db.InstanceGet("gorm:cache:query:single_flight_cancel"); hasCancel {
+			if cancel, ok := cancelObj.(context.CancelFunc); ok {
+				cancel()
+			}
 		}
 	}
 	c.dest = db.Statement.Dest