update pg dsn

NekoWheel · Jan 12, 2024 · d5188f8 · d5188f8
1 parent 1b3740e
commit d5188f8
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/internal/dbutil/testing.go b/internal/dbutil/testing.go
@@ -9,6 +9,7 @@ import (
 	"flag"
 	"fmt"
 	"math/rand"
+	"net/url"
 	"os"
 	"strconv"
 	"strings"
@@ -34,13 +35,12 @@ func NewTestDB(t *testing.T, migrationTables ...interface{}) (testDB *gorm.DB, c
 		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("host=$DB_HOST user=$DB_USER password=$DB_PASSWORD dbname=$DB_DATABASE port=$DB_PORT sslmode=disable TimeZone=Asia/Shanghai")
+		dsn = os.ExpandEnv("postgres://$DB_USER:$DB_PASSWORD@$DB_HOST:$DB_PORT/$DB_DATABASE?sslmode=disable")
 		dialectFunc = postgres.Open
 	default:
 		t.Fatalf("Unknown database type: %q", dbType)
 	}
 
-	fmt.Println(dsn)
 	db, err := gorm.Open(dialectFunc(dsn), &gorm.Config{
 		NowFunc:                Now,
 		SkipDefaultTransaction: true,
@@ -59,12 +59,15 @@ func NewTestDB(t *testing.T, migrationTables ...interface{}) (testDB *gorm.DB, c
 		t.Fatalf("Failed to create test database: %v", err)
 	}
 
-	// HACK: replace the database name in the DSN.
-	dsn = strings.ReplaceAll(dsn, os.Getenv("DB_DATABASE"), dbname)
+	cfg, err := url.Parse(dsn)
+	if err != nil {
+		t.Fatalf("Failed to parse DSN")
+	}
+	cfg.Path = "/" + dbname
 
 	flagParseOnce.Do(flag.Parse)
 
-	testDB, err = gorm.Open(dialectFunc(dsn), &gorm.Config{
+	testDB, err = gorm.Open(dialectFunc(cfg.String()), &gorm.Config{
 		NowFunc:                Now,
 		SkipDefaultTransaction: true,
 	})