junhong_cmp_fiber/tests/testutils/db.go

package testutils

import (
	"context"
	"fmt"
	"strings"
	"sync"
	"testing"

	"github.com/redis/go-redis/v9"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"

	"github.com/break/junhong_cmp_fiber/internal/model"
)

// 全局单例数据库和 Redis 连接
// 使用 sync.Once 确保整个测试套件只创建一次连接，显著提升测试性能
var (
	testDBOnce    sync.Once
	testDB        *gorm.DB
	testDBInitErr error

	testRedisOnce    sync.Once
	testRedis        *redis.Client
	testRedisInitErr error
)

// 测试数据库配置
// TODO: 未来可以从环境变量或配置文件加载
const (
	testDBDSN       = "host=cxd.whcxd.cn port=16159 user=erp_pgsql password=erp_2025 dbname=junhong_cmp_test sslmode=disable TimeZone=Asia/Shanghai"
	testRedisAddr   = "cxd.whcxd.cn:16299"
	testRedisPasswd = "cpNbWtAaqgo1YJmbMp3h"
	testRedisDB     = 15
)

// GetTestDB 获取全局单例测试数据库连接
//
// 特点:
//   - 使用 sync.Once 确保整个测试套件只创建一次连接
//   - AutoMigrate 只在首次连接时执行一次
//   - 连接失败会跳过测试（不是致命错误）
//
// 用法:
//
//	func TestXxx(t *testing.T) {
//	    db := testutils.GetTestDB(t)
//	    // db 是全局共享的连接，不要直接修改其状态
//	    // 如需事务隔离，使用 NewTestTransaction(t)
//	}
func GetTestDB(t *testing.T) *gorm.DB {
	t.Helper()

	testDBOnce.Do(func() {
		var err error
		testDB, err = gorm.Open(postgres.Open(testDBDSN), &gorm.Config{
			Logger: logger.Default.LogMode(logger.Silent),
		})
		if err != nil {
			testDBInitErr = fmt.Errorf("无法连接测试数据库: %w", err)
			return
		}

		err = testDB.AutoMigrate(
			&model.Account{},
			&model.Role{},
			&model.Permission{},
			&model.AccountRole{},
			&model.RolePermission{},
			&model.Shop{},
			&model.Enterprise{},
			&model.PersonalCustomer{},
			&model.PersonalCustomerPhone{},
			&model.PersonalCustomerICCID{},
			&model.PersonalCustomerDevice{},
			&model.IotCard{},
			&model.IotCardImportTask{},
			&model.Device{},
			&model.DeviceImportTask{},
			&model.DeviceSimBinding{},
			&model.Carrier{},
			&model.Tag{},
			&model.PackageSeries{},
			&model.Package{},
			&model.ShopPackageAllocation{},
			&model.ShopSeriesAllocation{},
			&model.EnterpriseCardAuthorization{},
			&model.EnterpriseDeviceAuthorization{},
			&model.AssetAllocationRecord{},
			&model.CommissionWithdrawalRequest{},
			&model.CommissionWithdrawalSetting{},
			&model.Order{},
			&model.OrderItem{},
			&model.PackageUsage{},
			&model.Wallet{},
		)
		if err != nil {
			errMsg := err.Error()
			if strings.Contains(errMsg, "does not exist") && (strings.Contains(errMsg, "constraint") || strings.Contains(errMsg, "column")) {
				// 忽略约束和列不存在的错误，这是由于约束名变更或迁移未应用导致的
			} else {
				testDBInitErr = fmt.Errorf("数据库迁移失败: %w", err)
				return
			}
		}

		// 确保所有必要的列都存在（处理迁移未应用的情况）
		ensureTestDBColumns(testDB)
	})

	if testDBInitErr != nil {
		t.Skipf("跳过测试：%v", testDBInitErr)
	}

	return testDB
}

// GetTestRedis 获取全局单例 Redis 连接
//
// 特点:
//   - 使用 sync.Once 确保整个测试套件只创建一次连接
//   - 连接失败会跳过测试（不是致命错误）
//
// 用法:
//
//	func TestXxx(t *testing.T) {
//	    rdb := testutils.GetTestRedis(t)
//	    // rdb 是全局共享的连接
//	    // 使用 CleanTestRedisKeys(t) 自动清理测试相关的 Redis 键
//	}
func GetTestRedis(t *testing.T) *redis.Client {
	t.Helper()

	testRedisOnce.Do(func() {
		testRedis = redis.NewClient(&redis.Options{
			Addr:     testRedisAddr,
			Password: testRedisPasswd,
			DB:       testRedisDB,
		})

		ctx := context.Background()
		if err := testRedis.Ping(ctx).Err(); err != nil {
			testRedisInitErr = fmt.Errorf("无法连接 Redis: %w", err)
			return
		}
	})

	if testRedisInitErr != nil {
		t.Skipf("跳过测试：%v", testRedisInitErr)
	}

	return testRedis
}

// NewTestTransaction 创建测试事务，自动在测试结束时回滚
//
// 特点:
//   - 每个测试用例获得独立的事务，互不干扰
//   - 使用 t.Cleanup() 确保即使测试 panic 也能回滚
//   - 回滚后数据库状态与测试前完全一致
//
// 用法:
//
//	func TestXxx(t *testing.T) {
//	    tx := testutils.NewTestTransaction(t)
//	    // 所有数据库操作使用 tx 而非 db
//	    store := postgres.NewXxxStore(tx, rdb)
//	    // 测试结束后自动回滚，无需手动清理
//	}
//
// 注意:
//   - 不要在子测试(t.Run)中调用此函数，因为子测试可能并行执行
//   - 如需在子测试中使用数据库，应在父测试中创建事务并传递
func NewTestTransaction(t *testing.T) *gorm.DB {
	t.Helper()

	db := GetTestDB(t)
	// 确保所有必要的列都存在
	ensureTestDBColumns(db)

	tx := db.Begin()
	if tx.Error != nil {
		t.Fatalf("开启测试事务失败: %v", tx.Error)
	}

	// 使用 t.Cleanup() 确保测试结束时自动回滚
	// 即使测试 panic 也能执行清理
	t.Cleanup(func() {
		tx.Rollback()
	})

	return tx
}

// CleanTestRedisKeys 清理当前测试的 Redis 键
//
// 特点:
//   - 使用测试名称作为键前缀，格式: test:{TestName}:*
//   - 测试开始时清理已有键（防止脏数据）
//   - 使用 t.Cleanup() 确保测试结束时自动清理
//
// 用法:
//
//	func TestXxx(t *testing.T) {
//	    rdb := testutils.GetTestRedis(t)
//	    testutils.CleanTestRedisKeys(t, rdb)
//	    // Redis 键使用测试专用前缀: test:TestXxx:your_key
//	}
//
// 键命名规范:
//   - 测试中创建的键应使用 GetTestRedisKeyPrefix(t) 作为前缀
//   - 例如: test:TestShopStore_Create:cache:shop:1
func CleanTestRedisKeys(t *testing.T, rdb *redis.Client) {
	t.Helper()

	ctx := context.Background()
	testPrefix := GetTestRedisKeyPrefix(t)

	// 测试开始前清理已有键
	cleanKeys(ctx, rdb, testPrefix)

	// 测试结束时自动清理
	t.Cleanup(func() {
		cleanKeys(ctx, rdb, testPrefix)
	})
}

// GetTestRedisKeyPrefix 获取当前测试的 Redis 键前缀
//
// 返回格式: test:{TestName}:
// 用于在测试中创建带前缀的 Redis 键，确保键不会与其他测试冲突
//
// 用法:
//
//	func TestXxx(t *testing.T) {
//	    prefix := testutils.GetTestRedisKeyPrefix(t)
//	    key := prefix + "my_cache_key"
//	    // key = "test:TestXxx:my_cache_key"
//	}
func GetTestRedisKeyPrefix(t *testing.T) string {
	t.Helper()
	return fmt.Sprintf("test:%s:", t.Name())
}

// cleanKeys 清理匹配前缀的所有 Redis 键
func cleanKeys(ctx context.Context, rdb *redis.Client, prefix string) {
	keys, err := rdb.Keys(ctx, prefix+"*").Result()
	if err != nil {
		// 忽略 Redis 错误，不影响测试
		return
	}
	if len(keys) > 0 {
		rdb.Del(ctx, keys...)
	}
}

// ensureTestDBColumns 确保测试数据库中所有必要的列都存在
// 处理迁移未应用导致的列缺失问题
func ensureTestDBColumns(db *gorm.DB) {
	// 添加 force_recharge_trigger_type 列到 tb_shop_series_allocation 表
	if !db.Migrator().HasColumn("tb_shop_series_allocation", "force_recharge_trigger_type") {
		db.Exec("ALTER TABLE tb_shop_series_allocation ADD COLUMN force_recharge_trigger_type int DEFAULT 2")
	}
}