go-sqlmock/driver.go

package sqlmock

import (
	"database/sql"
	"database/sql/driver"
	"fmt"
	"sync"
)

var pool *mockDriver

func init() {
	pool = &mockDriver{
		conns: make(map[string]*sqlmock),
	}
	sql.Register("sqlmock", pool)
}

type mockDriver struct {
	sync.Mutex
	counter int
	conns   map[string]*sqlmock
}

func (d *mockDriver) Open(dsn string) (driver.Conn, error) {
	d.Lock()
	defer d.Unlock()

	c, ok := d.conns[dsn]
	if !ok {
		return c, fmt.Errorf("expected a connection to be available, but it is not")
	}

	c.opened++
	return c, nil
}

// New creates sqlmock database connection
// and a mock to manage expectations.
// Pings db so that all expectations could be
// asserted.
func New() (*sql.DB, Sqlmock, error) {
	pool.Lock()
	dsn := fmt.Sprintf("sqlmock_db_%d", pool.counter)
	pool.counter++

	smock := &sqlmock{dsn: dsn, drv: pool, ordered: true}
	pool.conns[dsn] = smock
	pool.Unlock()

	return smock.open()
}

// NewWithDSN creates sqlmock database connection
// with a specific DSN and a mock to manage expectations.
// Pings db so that all expectations could be asserted.
//
// This method is introduced because of sql abstraction
// libraries, which do not provide a way to initialize
// with sql.DB instance. For example GORM library.
//
// Note, it will error if attempted to create with an
// already used dsn
//
// It is not recommended to use this method, unless you
// really need it and there is no other way around.
func NewWithDSN(dsn string) (*sql.DB, Sqlmock, error) {
	pool.Lock()
	if _, ok := pool.conns[dsn]; ok {
		pool.Unlock()
		return nil, nil, fmt.Errorf("cannot create a new mock database with the same dsn: %s", dsn)
	}
	smock := &sqlmock{dsn: dsn, drv: pool, ordered: true}
	pool.conns[dsn] = smock
	pool.Unlock()

	return smock.open()
}