杭州明婳科技
278 lines
9.2 KiB
Go
278 lines
9.2 KiB
Go
package database_test
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/EthanCodeCraft/xlgo-core/config"
|
|
"github.com/EthanCodeCraft/xlgo-core/database"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/clause"
|
|
"gorm.io/gorm/schema"
|
|
)
|
|
|
|
func TestCloseAllWithoutInit(t *testing.T) {
|
|
if err := database.CloseAll(); err != nil {
|
|
t.Fatalf("CloseAll without init should not error: %v", err)
|
|
}
|
|
if database.GetDB() != nil {
|
|
t.Fatal("expected DB nil")
|
|
}
|
|
if database.GetReadDB() != nil {
|
|
t.Fatal("expected read DB nil")
|
|
}
|
|
if len(database.GetReplicas()) != 0 {
|
|
t.Fatal("expected replicas empty")
|
|
}
|
|
}
|
|
|
|
func TestDBContextHelpersWithoutInit(t *testing.T) {
|
|
ctx := database.UseMaster(context.Background())
|
|
if db := database.GetDBFromContext(ctx); db != nil {
|
|
t.Fatal("expected nil DB without init")
|
|
}
|
|
|
|
ctx = database.UseReplica(context.Background())
|
|
if db := database.GetDBFromContext(ctx); db != nil {
|
|
t.Fatal("expected nil read DB without init")
|
|
}
|
|
}
|
|
|
|
func TestRoundRobinPicker(t *testing.T) {
|
|
replicas := []*gorm.DB{{}, {}, {}}
|
|
p := &database.RoundRobinPicker{}
|
|
|
|
first := p.Pick(replicas)
|
|
second := p.Pick(replicas)
|
|
third := p.Pick(replicas)
|
|
fourth := p.Pick(replicas)
|
|
|
|
if first == nil || second == nil || third == nil {
|
|
t.Fatal("Picker returned nil for non-empty replicas")
|
|
}
|
|
if first != replicas[0] || second != replicas[1] || third != replicas[2] {
|
|
t.Fatal("RoundRobinPicker should cycle through replicas in order")
|
|
}
|
|
if fourth != replicas[0] {
|
|
t.Fatal("RoundRobinPicker should wrap around to the first replica")
|
|
}
|
|
if p.Pick(nil) != nil || p.Pick([]*gorm.DB{}) != nil {
|
|
t.Fatal("Picker should return nil for empty replicas")
|
|
}
|
|
}
|
|
|
|
func TestRandomPicker(t *testing.T) {
|
|
replicas := []*gorm.DB{{}, {}}
|
|
p := &database.RandomPicker{}
|
|
|
|
picked := p.Pick(replicas)
|
|
if picked == nil {
|
|
t.Fatal("RandomPicker returned nil for non-empty replicas")
|
|
}
|
|
if picked != replicas[0] && picked != replicas[1] {
|
|
t.Fatal("RandomPicker returned a replica not in the slice")
|
|
}
|
|
if p.Pick(nil) != nil || p.Pick([]*gorm.DB{}) != nil {
|
|
t.Fatal("RandomPicker should return nil for empty replicas")
|
|
}
|
|
}
|
|
|
|
func TestManagerReplicaFallbackToMaster(t *testing.T) {
|
|
mgr := database.NewManager(&config.Config{})
|
|
if mgr.Master() != nil {
|
|
t.Fatal("expected nil master before init")
|
|
}
|
|
if mgr.Replicas() != nil {
|
|
t.Fatal("expected nil replicas before init")
|
|
}
|
|
// 无从库时 Replica 应返回 master（此处均为 nil）
|
|
if mgr.Replica() != nil {
|
|
t.Fatal("expected Replica to fall back to master when no replicas")
|
|
}
|
|
}
|
|
|
|
func TestManagerSetPicker(t *testing.T) {
|
|
mgr := database.NewManager(&config.Config{})
|
|
rr := &database.RoundRobinPicker{}
|
|
mgr.SetPicker(rr)
|
|
if mgr.Picker() != rr {
|
|
t.Fatal("SetPicker did not install the picker")
|
|
}
|
|
// nil 不应覆盖已有 picker
|
|
mgr.SetPicker(nil)
|
|
if mgr.Picker() != rr {
|
|
t.Fatal("SetPicker(nil) should not clear the existing picker")
|
|
}
|
|
}
|
|
|
|
func TestDefaultManagerHealthCheckWithoutInit(t *testing.T) {
|
|
if err := database.GetDefaultManager().HealthCheck(context.Background()); err == nil {
|
|
t.Fatal("expected error when health checking uninitialized master")
|
|
}
|
|
}
|
|
|
|
func TestNilContextHelpersDoNotPanic(t *testing.T) {
|
|
_ = database.CloseAll()
|
|
|
|
if ctx := database.UseMaster(nil); ctx == nil {
|
|
t.Fatal("UseMaster(nil) 应返回可用 context")
|
|
}
|
|
if ctx := database.UseReplica(nil); ctx == nil {
|
|
t.Fatal("UseReplica(nil) 应返回可用 context")
|
|
}
|
|
if db := database.GetDBFromContext(nil); db != nil {
|
|
t.Fatal("未初始化数据库时 GetDBFromContext(nil) 应返回 nil")
|
|
}
|
|
if ctx := database.WithTx(nil, nil); ctx == nil {
|
|
t.Fatal("WithTx(nil, nil) 应返回可用 context")
|
|
}
|
|
if tx := database.TxFromContext(nil); tx != nil {
|
|
t.Fatal("TxFromContext(nil) 应返回 nil")
|
|
}
|
|
if err := database.TransactionWithContext(nil, func(tx *gorm.DB) error { return nil }); err == nil {
|
|
t.Fatal("未初始化数据库时 TransactionWithContext(nil) 应返回错误")
|
|
}
|
|
if err := database.ReadQuery(nil, &struct{}{}, "1=1"); err == nil {
|
|
t.Fatal("未初始化数据库时 ReadQuery(nil) 应返回错误")
|
|
}
|
|
if err := database.WriteQuery(nil, &struct{}{}, "1=1"); err == nil {
|
|
t.Fatal("未初始化数据库时 WriteQuery(nil) 应返回错误")
|
|
}
|
|
if got := database.HealthCheck(); !got["master"] {
|
|
// HealthCheck 内部使用 background context；未初始化时 master=false 即可。
|
|
return
|
|
}
|
|
t.Fatal("未初始化数据库时 HealthCheck master 不应为 true")
|
|
}
|
|
|
|
func TestNilConfigInitializationReturnsError(t *testing.T) {
|
|
mgr := database.NewManager(nil)
|
|
if err := mgr.InitDB(context.Background(), nil); err == nil {
|
|
t.Fatal("InitDB(nil) 应返回错误")
|
|
}
|
|
if err := mgr.InitDBWithReplicas(context.Background(), nil, nil); err == nil {
|
|
t.Fatal("InitDBWithReplicas(nil) 应返回错误")
|
|
}
|
|
if err := database.InitDB(context.Background(), nil); err == nil {
|
|
t.Fatal("包级 InitDB(nil) 应返回错误")
|
|
}
|
|
if err := database.InitDBWithReplicas(context.Background(), nil, nil); err == nil {
|
|
t.Fatal("包级 InitDBWithReplicas(nil) 应返回错误")
|
|
}
|
|
}
|
|
|
|
func TestDialectorNilConfigDoesNotPanic(t *testing.T) {
|
|
if d := database.Dialector(nil); d == nil {
|
|
t.Fatal("Dialector(nil) 应返回兜底 dialector")
|
|
}
|
|
}
|
|
|
|
func TestDialectorSelectsByDriver(t *testing.T) {
|
|
mysqlCfg := &config.Config{Database: config.DatabaseConfig{
|
|
Driver: config.DriverMySQL, Host: "localhost", Port: 3306,
|
|
User: "root", Password: "pass", Name: "db",
|
|
}}
|
|
if name := database.Dialector(mysqlCfg).Name(); name != "mysql" {
|
|
t.Fatalf("expected mysql dialector, got %q", name)
|
|
}
|
|
|
|
pgCfg := &config.Config{Database: config.DatabaseConfig{
|
|
Driver: config.DriverPostgres, Host: "localhost", Port: 5432,
|
|
User: "postgres", Password: "pass", Name: "db",
|
|
}}
|
|
if name := database.Dialector(pgCfg).Name(); name != "postgres" {
|
|
t.Fatalf("expected postgres dialector, got %q", name)
|
|
}
|
|
|
|
// 别名也应解析为 postgres
|
|
pgAliasCfg := &config.Config{Database: config.DatabaseConfig{
|
|
Driver: "PG", Host: "localhost", Port: 5432,
|
|
User: "postgres", Password: "pass", Name: "db",
|
|
}}
|
|
if name := database.Dialector(pgAliasCfg).Name(); name != "postgres" {
|
|
t.Fatalf("expected postgres dialector via alias, got %q", name)
|
|
}
|
|
|
|
// 未指定 Driver 时默认 mysql
|
|
defaultCfg := &config.Config{Database: config.DatabaseConfig{
|
|
Host: "localhost", Port: 3306, User: "root", Password: "pass", Name: "db",
|
|
}}
|
|
if name := database.Dialector(defaultCfg).Name(); name != "mysql" {
|
|
t.Fatalf("expected default mysql dialector, got %q", name)
|
|
}
|
|
}
|
|
|
|
// stubDialector 是一个用于测试 RegisterDialect 的占位 Dialector。
|
|
type stubDialector struct{ name string }
|
|
|
|
func (s stubDialector) Name() string { return s.name }
|
|
func (s stubDialector) Initialize(_ *gorm.DB) error { return nil }
|
|
func (s stubDialector) Migrator(db *gorm.DB) gorm.Migrator { return nil }
|
|
func (s stubDialector) DataTypeOf(*schema.Field) string { return "" }
|
|
func (s stubDialector) DefaultValueOf(*schema.Field) clause.Expression { return nil }
|
|
func (s stubDialector) BindVarTo(writer clause.Writer, _ *gorm.Statement, _ any) {}
|
|
func (s stubDialector) QuoteTo(writer clause.Writer, str string) { _, _ = writer.WriteString(str) }
|
|
func (s stubDialector) Explain(sql string, _ ...any) string { return sql }
|
|
|
|
func TestRegisterDialectAndCustomDriver(t *testing.T) {
|
|
const driver = "stubdb"
|
|
|
|
database.RegisterDialect(database.DialectSpec{
|
|
Name: driver,
|
|
Aliases: []string{"stub"},
|
|
Dialector: func(dsn string) gorm.Dialector { return stubDialector{name: "stubdb"} },
|
|
DSN: func(c *config.DatabaseConfig) string {
|
|
return "stub://" + c.Host
|
|
},
|
|
})
|
|
|
|
// Dialector 工厂可以解析主名和别名
|
|
if _, ok := database.LookupDialect(driver); !ok {
|
|
t.Fatal("expected stubdb dialector to be registered")
|
|
}
|
|
if _, ok := database.LookupDialect("STUB"); !ok {
|
|
t.Fatal("expected stub alias to be registered (case-insensitive)")
|
|
}
|
|
|
|
cfg := &config.Config{Database: config.DatabaseConfig{
|
|
Driver: driver, Host: "localhost",
|
|
}}
|
|
if name := database.Dialector(cfg).Name(); name != "stubdb" {
|
|
t.Fatalf("expected stubdb dialector, got %q", name)
|
|
}
|
|
|
|
// config.DSN() 应使用注册的 DSN 构建器
|
|
if dsn := cfg.Database.DSN(); dsn != "stub://localhost" {
|
|
t.Fatalf("expected DSN built by registered builder, got %q", dsn)
|
|
}
|
|
|
|
// 未知驱动应 fail-closed，避免拼写错误静默连向 MySQL。
|
|
unknownCfg := &config.Config{Database: config.DatabaseConfig{
|
|
Driver: "no-such-driver", Host: "localhost", Port: 3306,
|
|
User: "root", Password: "pass", Name: "db",
|
|
}}
|
|
d := database.Dialector(unknownCfg)
|
|
if name := d.Name(); name != "invalid" {
|
|
t.Fatalf("expected invalid dialector for unknown driver, got %q", name)
|
|
}
|
|
if _, err := gorm.Open(d, &gorm.Config{}); err == nil || !strings.Contains(err.Error(), "数据库驱动未注册") {
|
|
t.Fatalf("未知 driver 应在 GORM 初始化时报错，实际: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestRegisteredDialectsContainsBuiltins(t *testing.T) {
|
|
registered := database.RegisteredDialects()
|
|
want := map[string]bool{"mysql": false, "postgres": false, "pg": false}
|
|
for _, n := range registered {
|
|
if _, ok := want[n]; ok {
|
|
want[n] = true
|
|
}
|
|
}
|
|
for k, found := range want {
|
|
if !found {
|
|
t.Errorf("expected %q to be registered by default", k)
|
|
}
|
|
}
|
|
}
|