[management] optimize test execution (#3204)

management/server/store/store.go

@@ -9,11 +9,16 @@ import (
 	"os"
 	"path"
 	"path/filepath"
+	"regexp"
 	"runtime"
+	"slices"
 	"strings"
 	"time"
 
+	"github.com/google/uuid"
 	log "github.com/sirupsen/logrus"
+	"gorm.io/driver/mysql"
+	"gorm.io/driver/postgres"
 	"gorm.io/driver/sqlite"
 	"gorm.io/gorm"
 
@@ -193,6 +198,8 @@ const (
 	mysqlDsnEnv    = "NETBIRD_STORE_ENGINE_MYSQL_DSN"
 )
 
+var supportedEngines = []Engine{SqliteStoreEngine, PostgresStoreEngine, MysqlStoreEngine}
+
 func getStoreEngineFromEnv() Engine {
 	// NETBIRD_STORE_ENGINE supposed to be used in tests. Otherwise, rely on the config file.
 	kind, ok := os.LookupEnv("NETBIRD_STORE_ENGINE")
@@ -201,7 +208,7 @@ func getStoreEngineFromEnv() Engine {
 	}
 
 	value := Engine(strings.ToLower(kind))
-	if value == SqliteStoreEngine || value == PostgresStoreEngine || value == MysqlStoreEngine {
+	if slices.Contains(supportedEngines, value) {
 		return value
 	}
 
@@ -349,51 +356,126 @@ func NewTestStoreFromSQL(ctx context.Context, filename string, dataDir string) (
 }
 
 func getSqlStoreEngine(ctx context.Context, store *SqlStore, kind Engine) (Store, func(), error) {
-	if kind == PostgresStoreEngine {
-		cleanUp, err := testutil.CreatePostgresTestContainer()
-		if err != nil {
-			return nil, nil, err
+	var cleanup func()
+	var err error
+	switch kind {
+	case PostgresStoreEngine:
+		store, cleanup, err = newReusedPostgresStore(ctx, store, kind)
+	case MysqlStoreEngine:
+		store, cleanup, err = newReusedMysqlStore(ctx, store, kind)
+	default:
+		cleanup = func() {
+			// sqlite doesn't need to be cleaned up
 		}
-
-		dsn, ok := os.LookupEnv(postgresDsnEnv)
-		if !ok {
-			return nil, nil, fmt.Errorf("%s is not set", postgresDsnEnv)
-		}
-
-		store, err = NewPostgresqlStoreFromSqlStore(ctx, store, dsn, nil)
-		if err != nil {
-			return nil, nil, err
-		}
-
-		return store, cleanUp, nil
 	}
-
-	if kind == MysqlStoreEngine {
-		cleanUp, err := testutil.CreateMysqlTestContainer()
-		if err != nil {
-			return nil, nil, err
-		}
-
-		dsn, ok := os.LookupEnv(mysqlDsnEnv)
-		if !ok {
-			return nil, nil, fmt.Errorf("%s is not set", mysqlDsnEnv)
-		}
-
-		store, err = NewMysqlStoreFromSqlStore(ctx, store, dsn, nil)
-		if err != nil {
-			return nil, nil, err
-		}
-
-		return store, cleanUp, nil
+	if err != nil {
+		return nil, cleanup, fmt.Errorf("failed to create test store: %v", err)
 	}
 
 	closeConnection := func() {
+		cleanup()
 		store.Close(ctx)
 	}
 
 	return store, closeConnection, nil
 }
 
+func newReusedPostgresStore(ctx context.Context, store *SqlStore, kind Engine) (*SqlStore, func(), error) {
+	if envDsn, ok := os.LookupEnv(postgresDsnEnv); !ok || envDsn == "" {
+		var err error
+		_, err = testutil.CreatePostgresTestContainer()
+		if err != nil {
+			return nil, nil, err
+		}
+	}
+
+	dsn, ok := os.LookupEnv(postgresDsnEnv)
+	if !ok {
+		return nil, nil, fmt.Errorf("%s is not set", postgresDsnEnv)
+	}
+
+	db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to open postgres connection: %v", err)
+	}
+
+	dsn, cleanup, err := createRandomDB(dsn, db, kind)
+	if err != nil {
+		return nil, cleanup, err
+	}
+
+	store, err = NewPostgresqlStoreFromSqlStore(ctx, store, dsn, nil)
+	if err != nil {
+		return nil, cleanup, err
+	}
+
+	return store, cleanup, nil
+}
+
+func newReusedMysqlStore(ctx context.Context, store *SqlStore, kind Engine) (*SqlStore, func(), error) {
+	if envDsn, ok := os.LookupEnv(mysqlDsnEnv); !ok || envDsn == "" {
+		var err error
+		_, err = testutil.CreateMysqlTestContainer()
+		if err != nil {
+			return nil, nil, err
+		}
+	}
+
+	dsn, ok := os.LookupEnv(mysqlDsnEnv)
+	if !ok {
+		return nil, nil, fmt.Errorf("%s is not set", mysqlDsnEnv)
+	}
+
+	db, err := gorm.Open(mysql.Open(dsn+"?charset=utf8&parseTime=True&loc=Local"), &gorm.Config{})
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to open mysql connection: %v", err)
+	}
+
+	dsn, cleanup, err := createRandomDB(dsn, db, kind)
+	if err != nil {
+		return nil, cleanup, err
+	}
+
+	store, err = NewMysqlStoreFromSqlStore(ctx, store, dsn, nil)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return store, cleanup, nil
+}
+
+func createRandomDB(dsn string, db *gorm.DB, engine Engine) (string, func(), error) {
+	dbName := fmt.Sprintf("test_db_%s", strings.ReplaceAll(uuid.New().String(), "-", "_"))
+
+	if err := db.Exec(fmt.Sprintf("CREATE DATABASE %s", dbName)).Error; err != nil {
+		return "", nil, fmt.Errorf("failed to create database: %v", err)
+	}
+
+	var err error
+	cleanup := func() {
+		switch engine {
+		case PostgresStoreEngine:
+			err = db.Exec(fmt.Sprintf("DROP DATABASE %s WITH (FORCE)", dbName)).Error
+		case MysqlStoreEngine:
+			// err = killMySQLConnections(dsn, dbName)
+			err = db.Exec(fmt.Sprintf("DROP DATABASE %s", dbName)).Error
+		}
+		if err != nil {
+			log.Errorf("failed to drop database %s: %v", dbName, err)
+			panic(err)
+		}
+		sqlDB, _ := db.DB()
+		_ = sqlDB.Close()
+	}
+
+	return replaceDBName(dsn, dbName), cleanup, nil
+}
+
+func replaceDBName(dsn, newDBName string) string {
+	re := regexp.MustCompile(`(?P<pre>[:/@])(?P<dbname>[^/?]+)(?P<post>\?|$)`)
+	return re.ReplaceAllString(dsn, `${pre}`+newDBName+`${post}`)
+}
+
 func loadSQL(db *gorm.DB, filepath string) error {
 	sqlContent, err := os.ReadFile(filepath)
 	if err != nil {