2026-05-12 09:53:20 +08:00

181 lines
4.6 KiB
Go

// Package mysqlschema creates isolated MySQL schemas for service integration tests.
package mysqlschema
import (
"database/sql"
"fmt"
"hash/crc32"
"os"
"path/filepath"
"runtime"
"strings"
"testing"
"time"
mysqldriver "github.com/go-sql-driver/mysql"
)
// Config describes how to create one temporary schema from a service initdb script.
type Config struct {
EnvVar string
InitDBPath string
DatabasePrefix string
}
// Schema owns one temporary MySQL database initialized from production DDL.
type Schema struct {
DSN string
DB *sql.DB
t testing.TB
adminDB *sql.DB
database string
}
// New creates an isolated schema and registers cleanup with t.Cleanup.
func New(t testing.TB, cfg Config) *Schema {
t.Helper()
if strings.TrimSpace(cfg.EnvVar) == "" {
cfg.EnvVar = "MYSQL_TEST_DSN"
}
baseDSN := os.Getenv(cfg.EnvVar)
if strings.TrimSpace(baseDSN) == "" {
t.Skipf("set %s to run MySQL-backed tests", cfg.EnvVar)
}
if strings.TrimSpace(cfg.InitDBPath) == "" {
t.Fatalf("mysqlschema InitDBPath is required")
}
if strings.TrimSpace(cfg.DatabasePrefix) == "" {
cfg.DatabasePrefix = "hy_test"
}
dsnCfg, err := mysqldriver.ParseDSN(baseDSN)
if err != nil {
t.Fatalf("parse %s failed: %v", cfg.EnvVar, err)
}
database := testDatabaseName(cfg.DatabasePrefix, t.Name())
adminCfg := dsnCfg.Clone()
adminCfg.DBName = ""
adminDB, err := sql.Open("mysql", adminCfg.FormatDSN())
if err != nil {
t.Fatalf("open mysql admin connection failed: %v", err)
}
if _, err := adminDB.Exec("CREATE DATABASE " + quoteIdentifier(database) + " DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci"); err != nil {
_ = adminDB.Close()
t.Fatalf("create test database failed: %v", err)
}
testCfg := dsnCfg.Clone()
testCfg.DBName = database
testDSN := testCfg.FormatDSN()
rawDB, err := sql.Open("mysql", testDSN)
if err != nil {
_ = adminDB.Close()
t.Fatalf("open mysql test database failed: %v", err)
}
initSchema(t, rawDB, cfg.InitDBPath)
schema := &Schema{
DSN: testDSN,
DB: rawDB,
t: t,
adminDB: adminDB,
database: database,
}
t.Cleanup(schema.Close)
return schema
}
// Close closes connections and drops the temporary database.
func (s *Schema) Close() {
s.t.Helper()
if s.DB != nil {
_ = s.DB.Close()
}
if s.adminDB != nil && s.database != "" {
if _, err := s.adminDB.Exec("DROP DATABASE IF EXISTS " + quoteIdentifier(s.database)); err != nil {
s.t.Fatalf("drop test database %s failed: %v", s.database, err)
}
_ = s.adminDB.Close()
}
}
// InitDBPath resolves a service initdb file relative to the caller source file.
func InitDBPath(t testing.TB, callerFile string, relativePath ...string) string {
t.Helper()
parts := append([]string{filepath.Dir(callerFile)}, relativePath...)
return filepath.Clean(filepath.Join(parts...))
}
// CallerFile returns the current caller file for compact per-service test helpers.
func CallerFile(t testing.TB, skip int) string {
t.Helper()
_, file, _, ok := runtime.Caller(skip)
if !ok {
t.Fatalf("runtime caller unavailable")
}
return file
}
func initSchema(t testing.TB, db *sql.DB, initDBPath string) {
t.Helper()
script, err := os.ReadFile(initDBPath)
if err != nil {
t.Fatalf("read initdb %s failed: %v", initDBPath, err)
}
for statement := range strings.SplitSeq(stripSQLLineComments(string(script)), ";") {
statement = strings.TrimSpace(statement)
if statement == "" {
continue
}
upper := strings.ToUpper(statement)
if strings.HasPrefix(upper, "CREATE DATABASE") || strings.HasPrefix(upper, "USE ") {
continue
}
if _, err := db.Exec(statement); err != nil {
t.Fatalf("exec initdb statement failed: %v\nstatement:\n%s", err, statement)
}
}
}
func stripSQLLineComments(script string) string {
lines := strings.Split(script, "\n")
kept := make([]string, 0, len(lines))
for _, line := range lines {
if strings.HasPrefix(strings.TrimSpace(line), "--") {
continue
}
kept = append(kept, line)
}
return strings.Join(kept, "\n")
}
func testDatabaseName(prefix string, testName string) string {
checksum := crc32.ChecksumIEEE([]byte(testName))
return fmt.Sprintf("%s_%08x_%d", sanitizePrefix(prefix), checksum, time.Now().UnixNano())
}
func sanitizePrefix(prefix string) string {
prefix = strings.ToLower(strings.TrimSpace(prefix))
var builder strings.Builder
for _, char := range prefix {
if (char >= 'a' && char <= 'z') || (char >= '0' && char <= '9') || char == '_' {
builder.WriteRune(char)
}
}
if builder.Len() == 0 {
return "hy_test"
}
return builder.String()
}
func quoteIdentifier(identifier string) string {
return "`" + strings.ReplaceAll(identifier, "`", "``") + "`"
}