From 5e004fdd64e2f25c648dd5aa37148ca90161e665 Mon Sep 17 00:00:00 2001
From: E99p1ant <i@github.red>
Date: Sat, 13 Jan 2024 00:48:29 +0800
Subject: [PATCH] update for mysql

---
 internal/dbutil/testing.go | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/internal/dbutil/testing.go b/internal/dbutil/testing.go
index 1308ed2..95075f7 100644
--- a/internal/dbutil/testing.go
+++ b/internal/dbutil/testing.go
@@ -31,7 +31,7 @@ func NewTestDB(t *testing.T, migrationTables ...interface{}) (testDB *gorm.DB, c
 
 	switch dbType {
 	case "mysql":
-		dsn = os.ExpandEnv("mysql://$DB_USER:$DB_PASSWORD@tcp($DB_HOST:$DB_PORT)?charset=utf8mb4&parseTime=True&loc=Local")
+		dsn = os.ExpandEnv("$DB_USER:$DB_PASSWORD@tcp($DB_HOST:$DB_PORT)/DB_DATABASE?charset=utf8mb4&parseTime=True&loc=Local")
 		dialectFunc = mysql.Open
 	case "postgres":
 		dsn = os.ExpandEnv("postgres://$DB_USER:$DB_PASSWORD@$DB_HOST:$DB_PORT?sslmode=disable")
@@ -57,15 +57,21 @@ func NewTestDB(t *testing.T, migrationTables ...interface{}) (testDB *gorm.DB, c
 		t.Fatalf("Failed to create test database: %v", err)
 	}
 
-	cfg, err := url.Parse(dsn)
-	if err != nil {
-		t.Fatalf("Failed to parse DSN")
+	switch dbType {
+	case "mysql":
+		dsn = strings.ReplaceAll(dsn, os.Getenv("DB_DATABASE"), dbname)
+	case "postgres":
+		cfg, err := url.Parse(dsn)
+		if err != nil {
+			t.Fatalf("Failed to parse DSN")
+		}
+		cfg.Path = "/" + dbname
+		dsn = cfg.String()
 	}
-	cfg.Path = "/" + dbname
 
 	flagParseOnce.Do(flag.Parse)
 
-	testDB, err = gorm.Open(dialectFunc(cfg.String()), &gorm.Config{
+	testDB, err = gorm.Open(dialectFunc(dsn), &gorm.Config{
 		NowFunc:                Now,
 		SkipDefaultTransaction: true,
 	})