0904

2026-09-08 10:10:44 +08:00
commit 2ed27dd55e
288 changed files with 67061 additions and 0 deletions
@@ -0,0 +1,145 @@
+// Copyright 2023 E99p1ant. All rights reserved.
+// Use of this source code is governed by a MIT-style
+// license that can be found in the LICENSE file.
+
+package dbutil
+
+import (
+	"context"
+	"flag"
+	"math/rand"
+	"net/url"
+	"os"
+	"strconv"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"gorm.io/driver/mysql"
+	"gorm.io/driver/postgres"
+	"gorm.io/gorm"
+)
+
+var flagParseOnce sync.Once
+
+func NewTestDB(t *testing.T, migrationTables ...interface{}) (testDB *gorm.DB, cleanup func(...string) error) {
+	dbType := os.Getenv("DB_TYPE")
+
+	var dsn string
+	var dialectFunc func(string) gorm.Dialector
+
+	switch dbType {
+	case "mysql":
+		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")
+		dialectFunc = postgres.Open
+	default:
+		t.Fatalf("Unknown database type: %q", dbType)
+	}
+
+	db, err := gorm.Open(dialectFunc(dsn), &gorm.Config{
+		NowFunc:                Now,
+		SkipDefaultTransaction: true,
+	})
+	if err != nil {
+		t.Fatalf("Failed to open connection: %v", err)
+	}
+
+	ctx := context.Background()
+	rng := rand.New(rand.NewSource(time.Now().UnixNano()))
+	dbname := "tamabox-test-" + strconv.FormatUint(rng.Uint64(), 10)
+
+	err = db.WithContext(ctx).Exec(`CREATE DATABASE ` + QuoteIdentifier(dbType, dbname)).Error
+	if err != nil {
+		t.Fatalf("Failed to create test database: %v", err)
+	}
+
+	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()
+	}
+
+	flagParseOnce.Do(flag.Parse)
+
+	testDB, err = gorm.Open(dialectFunc(dsn), &gorm.Config{
+		NowFunc:                Now,
+		SkipDefaultTransaction: true,
+	})
+	if err != nil {
+		t.Fatalf("Failed to open test connection: %v", err)
+	}
+
+	err = testDB.AutoMigrate(migrationTables...)
+	if err != nil {
+		t.Fatalf("Failed to auto migrate tables: %v", err)
+	}
+
+	t.Cleanup(func() {
+		defer func() {
+			if database, err := db.DB(); err == nil {
+				_ = database.Close()
+			}
+		}()
+
+		if t.Failed() {
+			t.Logf("DATABASE %s left intact for inspection", dbname)
+			return
+		}
+
+		database, err := testDB.DB()
+		if err != nil {
+			t.Fatalf("Failed to get currently open database: %v", err)
+		}
+
+		err = database.Close()
+		if err != nil {
+			t.Fatalf("Failed to close currently open database: %v", err)
+		}
+
+		err = db.WithContext(ctx).Exec(`DROP DATABASE ` + QuoteIdentifier(dbType, dbname)).Error
+		if err != nil {
+			t.Fatalf("Failed to drop test database: %v", err)
+		}
+	})
+
+	return testDB, func(tables ...string) error {
+		if t.Failed() {
+			return nil
+		}
+
+		for _, table := range tables {
+			var query string
+			switch dbType {
+			case "mysql":
+				query = `TRUNCATE TABLE ` + QuoteIdentifier(dbType, table)
+			case "postgres":
+				query = `TRUNCATE TABLE ` + QuoteIdentifier(dbType, table) + ` RESTART IDENTITY CASCADE`
+			}
+
+			err := testDB.WithContext(ctx).Exec(query).Error
+			if err != nil {
+				return err
+			}
+		}
+		return nil
+	}
+}
+
+// QuoteIdentifier quotes an "identifier" (e.g. a table or a column name) to be
+// used as part of an SQL statement.
+func QuoteIdentifier(typ, s string) string {
+	if typ == "postgres" {
+		return `"` + strings.ReplaceAll(s, `"`, `""`) + `"`
+	}
+	return "`" + strings.ReplaceAll(s, "`", "``") + "`"
+}