go-omni/database/gormx/gorm_test.go

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")
	}
}