package gormx import ( "context" "database/sql/driver" "errors" "io" "sync/atomic" "testing" "time" "git.yuntuo.hk/k2safe/go-omni/database" "git.yuntuo.hk/k2safe/go-omni/logx" "gorm.io/gorm" "gorm.io/gorm/clause" "gorm.io/gorm/schema" ) type testConnector struct { closed atomic.Int32 query func(context.Context) error } func (c *testConnector) Connect(context.Context) (driver.Conn, error) { return &testConn{owner: c}, nil } func (c *testConnector) Driver() driver.Driver { return testDriver{} } func (c *testConnector) Close() error { c.closed.Add(1); return nil } type testDriver struct{} func (testDriver) Open(string) (driver.Conn, error) { return nil, errors.New("unexpected new pool") } type testConn struct{ owner *testConnector } func (c *testConn) Prepare(string) (driver.Stmt, error) { return nil, errors.New("unused") } func (c *testConn) Close() error { return nil } func (c *testConn) Begin() (driver.Tx, error) { return nil, errors.New("unused") } func (c *testConn) Ping(context.Context) error { return nil } func (c *testConn) QueryContext(ctx context.Context, _ string, _ []driver.NamedValue) (driver.Rows, error) { if c.owner.query != nil { if err := c.owner.query(ctx); err != nil { return nil, err } } return &testRows{}, nil } type testRows struct{ done bool } func (r *testRows) Columns() []string { return []string{"version"} } func (r *testRows) Close() error { return nil } func (r *testRows) Next(values []driver.Value) error { if r.done { return io.EOF } r.done = true values[0] = "8.4.0" return nil } func poolFixture(t *testing.T) (*database.DB, *testConnector) { t.Helper() connector := &testConnector{} pool, err := database.OpenConnector(context.Background(), connector, database.DefaultPoolConfig()) if err != nil { t.Fatal(err) } t.Cleanup(func() { _ = pool.Close() }) return pool, connector } func TestBorrowedPoolFailureAndStartupDeadline(t *testing.T) { for _, mode := range []string{"driverError", "deadline", "canceled"} { t.Run(mode, func(t *testing.T) { pool, connector := poolFixture(t) cause := errors.New("private-driver-error") ctx, cancel := context.WithCancel(context.Background()) defer cancel() connector.query = func(ctx context.Context) error { if mode == "driverError" { return cause } <-ctx.Done() return ctx.Err() } if mode == "canceled" { cancel() cause = context.Canceled } else if mode == "deadline" { cause = context.DeadlineExceeded } start := time.Now() orm, err := MySQL(ctx, pool, WithInitTimeout(20*time.Millisecond)) if orm != nil || !errors.Is(err, cause) { t.Fatal("lost cause", err) } if time.Since(start) > time.Second { t.Fatal("startup deadline ignored") } if connector.closed.Load() != 0 { t.Fatal("borrowed pool closed after GORM failure") } if healthErr := pool.Health(context.Background()); healthErr != nil { t.Fatal("owner pool unusable", healthErr) } if err.Error() == cause.Error() { t.Fatal("unsafe initialization error") } }) } } func TestPoolIdentityNativeOptionsAndRequestContext(t *testing.T) { pool, connector := poolFixture(t) ctx, cancel := context.WithCancel(context.Background()) cfg := gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: "business_"}, ClauseBuilders: map[string]clause.ClauseBuilder{}} logger, _ := logx.New(logx.Config{Output: "discard"}) orm, err := MySQL(ctx, pool, WithConfig(cfg), WithLogger(logger)) if err != nil { t.Fatal(err) } cancel() // The startup context must not poison future queries. got, err := orm.DB() if err != nil || got != pool.DB { t.Fatal("pool identity changed", err) } if len(cfg.ClauseBuilders) != 0 || cfg.ConnPool != nil || cfg.Logger != nil { t.Fatal("native input config mutated") } if orm.NamingStrategy.TableName("Item") != "business_items" { t.Fatal("native option lost") } if orm.SkipDefaultTransaction || orm.AllowGlobalUpdate { t.Fatal("transaction safety defaults changed") } ctx, cancel = context.WithCancel(context.Background()) cancel() var v string if err = orm.WithContext(ctx).Raw("SELECT VERSION()").Scan(&v).Error; !errors.Is(err, context.Canceled) { t.Fatal(err) } if err = orm.WithContext(context.Background()).Raw("SELECT VERSION()").Scan(&v).Error; err != nil || v != "8.4.0" { t.Fatal(v, err) } if connector.closed.Load() != 0 { t.Fatal("pool unexpectedly closed") } if err = pool.Close(); err != nil { t.Fatal(err) } if err = orm.Raw("SELECT VERSION()").Scan(&v).Error; err == nil { t.Fatal("owner close did not affect GORM") } } func TestRejectedOptionsDoNotTouchPool(t *testing.T) { for _, option := range []Option{WithInitTimeout(0), WithLogConfig(LogConfig{Level: "bad"}), WithConfig(gorm.Config{ConnPool: &borrowedPool{}})} { pool, connector := poolFixture(t) if _, err := MySQL(context.Background(), pool, option); err == nil { t.Fatal("invalid option accepted") } if connector.closed.Load() != 0 { t.Fatal("pool closed") } } if _, err := MySQL(nil, nil); err == nil { t.Fatal("nil pool accepted") } pool, _ := poolFixture(t) if _, err := Open(nil, pool, func(gorm.ConnPool) gorm.Dialector { return nil }); err == nil { t.Fatal("nil dialect accepted") } }