pgx driver mock for Golang

pgxmock is a mock library implementing pgx - PostgreSQL Driver and Toolkit. It's based on the well-known sqlmock library for sql/driver.

pgxmock has one and only purpose - to simulate pgx behavior in tests, without needing a real database connection. It helps to maintain correct TDD workflow.

• does not require any modifications to your source code;
• has strict by default expectation order matching;
• has no third party dependencies except pgx packages (testify is used only by its own tests).

Install

go get github.com/pashagolub/pgxmock/v6

Version support policy

pgxmock follows the official Go release policy: the two most recent major Go releases are supported. Older versions may work but are not tested or maintained.

Documentation and Examples

Visit godoc for general examples and public api reference.

See implementation examples:

• the simplest one
• blog API server
• copy from
• batch

Something you may want to test

package main

import (
	"context"

	pgx "github.com/jackc/pgx/v5"
	pgxpool "github.com/jackc/pgx/v5/pgxpool"
)

type PgxIface interface {
	Begin(context.Context) (pgx.Tx, error)
	Close()
}

func recordStats(db PgxIface, userID, productID int) (err error) {
	tx, err := db.Begin(context.Background())
	if err != nil {
		return
	}
	defer func() {
		switch err {
		case nil:
			err = tx.Commit(context.Background())
		default:
			_ = tx.Rollback(context.Background())
		}
	}()
	sql := "UPDATE products SET views = views + 1"
	if _, err = tx.Exec(context.Background(), sql); err != nil {
		return
	}
	sql = "INSERT INTO product_viewers (user_id, product_id) VALUES ($1, $2)"
	if _, err = tx.Exec(context.Background(), sql, userID, productID); err != nil {
		return
	}
	return
}

func main() {
	// @NOTE: the real connection is not required for tests
	db, err := pgxpool.New(context.Background(), "postgres://rolname@hostname/dbname")
	if err != nil {
		panic(err)
	}
	defer db.Close()

	if err = recordStats(db, 1 /*some user id*/, 5 /*some product id*/); err != nil {
		panic(err)
	}
}

Tests with pgxmock

package main

import (
	"context"
	"fmt"
	"testing"

	"github.com/pashagolub/pgxmock/v6"
)

// a successful case
func TestShouldUpdateStats(t *testing.T) {
	mock, err := pgxmock.NewPool()
	if err != nil {
		t.Fatal(err)
	}
	defer mock.Close()

	mock.ExpectBegin()
	mock.ExpectExec("UPDATE products").
		WillReturnResult(pgxmock.NewResult("UPDATE", 1))
	mock.ExpectExec("INSERT INTO product_viewers").
		WithArgs(2, 3).
		WillReturnResult(pgxmock.NewResult("INSERT", 1))
	mock.ExpectCommit()

	// now we execute our method
	if err = recordStats(mock, 2, 3); err != nil {
		t.Errorf("error was not expected while updating: %s", err)
	}

	// we make sure that all expectations were met
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Errorf("there were unfulfilled expectations: %s", err)
	}
}

// a failing test case
func TestShouldRollbackStatUpdatesOnFailure(t *testing.T) {
	mock, err := pgxmock.NewPool()
	if err != nil {
		t.Fatal(err)
	}
	defer mock.Close()

	mock.ExpectBegin()
	mock.ExpectExec("UPDATE products").
		WillReturnResult(pgxmock.NewResult("UPDATE", 1))
	mock.ExpectExec("INSERT INTO product_viewers").
		WithArgs(2, 3).
		WillReturnError(fmt.Errorf("some error"))
	mock.ExpectRollback()

	// now we execute our method
	if err = recordStats(mock, 2, 3); err == nil {
		t.Errorf("was expecting an error, but there was none")
	}

	// we make sure that all expectations were met
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Errorf("there were unfulfilled expectations: %s", err)
	}
}

Loading custom types

Code that uses PostgreSQL enums, composites, domains or ranges loads them at startup and registers them. ExpectLoadTypes (or ExpectLoadType) mocks the loading; as in pgx, registering stays with the code under test:

// the code under test
type typeLoader interface {
	LoadTypes(ctx context.Context, typeNames []string) ([]*pgtype.Type, error)
	TypeMap() *pgtype.Map
}

func registerTypes(ctx context.Context, conn typeLoader) error {
	types, err := conn.LoadTypes(ctx, []string{"status"})
	if err != nil {
		return err
	}
	conn.TypeMap().RegisterTypes(types)
	return nil
}

// the test
func TestRegisterTypes(t *testing.T) {
	mock, err := pgxmock.NewConn()
	require.NoError(t, err)
	mock.ExpectLoadTypes("status").
		WillReturnTypes(&pgtype.Type{Name: "status", OID: 100000, Codec: &pgtype.EnumCodec{}})

	require.NoError(t, registerTypes(context.Background(), mock))
	_, ok := mock.TypeMap().TypeForName("status")
	assert.True(t, ok)
	assert.NoError(t, mock.ExpectationsWereMet())
}

Customize SQL query matching

There were plenty of requests from users regarding SQL query string validation or different matching option. We have implemented the QueryMatcher interface, which can be passed through an option when creating a mock with pgxmock.NewPool or pgxmock.NewConn.

This now allows to include some library, which would allow for example to parse and validate SQL AST. And create a custom QueryMatcher in order to validate SQL in sophisticated ways.

By default, pgxmock is preserving backward compatibility and default query matcher is pgxmock.QueryMatcherRegexp which uses expected SQL string as a regular expression to match incoming query string. There are also:

• QueryMatcherEqual, which does a full case sensitive match, ignoring differences in whitespace;
• QueryMatcherSubstring, which checks that the query contains the expected string, without the escaping a regular expression needs;
• QueryMatcherAny, which disables SQL matching altogether.

In order to customize the QueryMatcher, use the following:

	mock, err := pgxmock.NewPool(pgxmock.QueryMatcherOption(pgxmock.QueryMatcherEqual))

The query matcher can be fully customized based on user needs. pgxmock will not provide a standard sql parsing matchers.

Reporting a connection configuration

By default Config() reports a zero configuration. To test code that reads its own settings back off the pool, hand the mock the configuration production uses:

// the code under test
type configured interface {
	Config() *pgxpool.Config
}

func workers(db configured) int {
	return int(db.Config().MaxConns) / 2
}

// the test
func TestWorkers(t *testing.T) {
	cfg, err := pgxpool.ParseConfig("postgres://bob@localhost:5432/orders?pool_max_conns=8")
	require.NoError(t, err)
	mock, err := pgxmock.NewPool(pgxmock.PoolConfigOption(cfg))
	require.NoError(t, err)

	assert.Equal(t, 4, workers(mock))
}

For a connection mock use pgxmock.ConnConfigOption with a *pgx.ConnConfig.

Matching arguments like time.Time

There may be arguments which are of struct type and cannot be compared easily by value like time.Time. In this case pgxmock provides an Argument interface which can be used in more sophisticated matching. Here is a simple example of time argument matching:

import (
	"context"
	"testing"
	"time"

	"github.com/pashagolub/pgxmock/v6"
)

type AnyTime struct{}

// Match satisfies pgxmock.Argument interface
func (a AnyTime) Match(v any) bool {
	_, ok := v.(time.Time)
	return ok
}

func TestAnyTimeArgument(t *testing.T) {
	t.Parallel()
	mock, err := pgxmock.NewPool()
	if err != nil {
		t.Errorf("an error '%s' was not expected when creating a mock pool", err)
	}
	defer mock.Close()

	mock.ExpectExec("INSERT INTO users").
		WithArgs("john", AnyTime{}).
		WillReturnResult(pgxmock.NewResult("INSERT", 1))

	_, err = mock.Exec(context.Background(),
		"INSERT INTO users(name, created_at) VALUES ($1, $2)", "john", time.Now())
	if err != nil {
		t.Errorf("error '%s' was not expected, while inserting a row", err)
	}

	if err := mock.ExpectationsWereMet(); err != nil {
		t.Errorf("there were unfulfilled expectations: %s", err)
	}
}

It only asserts that the argument is of time.Time type. The same is built in as pgxmock.OfType[time.Time](), so a custom Argument is rarely needed:

Matcher	Matches
pgxmock.AnyArg()	any value
pgxmock.NotNil()	any value that is not nil; a typed nil pointer, slice or map counts as nil
pgxmock.OfType[T]()	any value of type T
pgxmock.AnyOf(values...)	any of the values, which may themselves be matchers
pgxmock.ArgumentFunc(f)	whatever the function f(any) bool accepts

	mock.ExpectExec("UPDATE orders").
		WithArgs(pgxmock.AnyOf("pending", "active"), pgxmock.OfType[time.Time](), pgxmock.NotNil()).
		WillReturnResult(pgxmock.NewResult("UPDATE", 1))

pgx.NamedArgs are matched as well. Use WithRewrittenSQL to also check the SQL they are rewritten to.

Simulating server and routing errors

Any expectation can be told to fail with WillReturnError, and the error worth returning is usually the one production code branches on.

For a server-side failure, pgxmock.NewPgError builds a *pgconn.PgError with a SQLSTATE code, so the code under test takes the branch it takes in production:

	mock.ExpectExec("INSERT INTO users").
		WithArgs("john").
		WillReturnError(pgxmock.NewPgError("23505",
			`duplicate key value violates unique constraint "users_email_key"`))

	// in the code under test
	var pgErr *pgconn.PgError
	if errors.As(err, &pgErr) && pgErr.Code == "23505" { ... }

For code that routes reads and writes across a primary and its replicas with target_session_attrs, pgx v5.11 added sentinel errors that say why a connection was rejected. They are plain errors, so they are returned the same way and matched with errors.Is:

	mock.ExpectPing().WillReturnError(pgconn.ErrReadOnlyConnection)

The full set is pgconn.ErrReadOnlyConnection, pgconn.ErrReadWriteConnection, pgconn.ErrPrimaryConnection and pgconn.ErrStandbyConnection.

Returning rows and errors

	rows := mock.NewRows([]string{"id", "title"}).
		AddRow(1, "one").
		AddRow(2, "two").
		RowError(1, errors.New("row error")). // returned when reading the second row
		CloseError(errors.New("close error")) // returned by Rows.Err after the rows are read
	mock.ExpectQuery("SELECT id, title FROM articles").
		WillReturnRows(rows).
		RowsWillBeClosed() // fail ExpectationsWereMet if the code under test does not close them

Values assignable to the destination are stored as they are, sql.Scanner destinations and conversions are supported, and anything else goes through the pgtype codec of the column's DataTypeOID, see custom types.

Modifiers

Every expectation can be made optional, repeated, delayed, or made to fail or panic. The modifiers may be chained in any order with the builders specific to each expectation:

	mock.ExpectQuery("SELECT").
		Times(2).                         // must be called twice
		WillReturnRows(rows).
		WillDelayFor(100 * time.Millisecond) // honours the context deadline
	mock.ExpectPing().Maybe()               // may not be called at all

By default expectations must be met in the order they were declared. Call mock.MatchExpectationsInOrder(false) when the code under test runs queries concurrently.

Transactions

Begin and BeginTx return the mock itself as the pgx.Tx, so expectations for the transaction are set on the mock.

Batches

	eb := mock.ExpectBatch()
	eb.ExpectQuery("SELECT balance").WithArgs(1).WillReturnRows(mock.NewRows([]string{"balance"}).AddRow(100))
	eb.ExpectExec("UPDATE accounts").WithArgs(1).WillReturnResult(pgxmock.NewResult("UPDATE", 1))

Results are read in the order they were queued, through BatchResults or the QueuedQuery callbacks.

CopyFrom

ExpectCopyFrom matches the table and columns. Add WithRows to assert the rows that were copied too:

	mock.ExpectCopyFrom(pgx.Identifier{"users"}, []string{"name", "age"}).
		WithRows(pgxmock.NewCopyRows("name", "age").
			AddRow("alice", 30).
			AddRow("bob", 40).
			Unordered()). // the rows may arrive in any order
		WillReturnResult(2)

Values are compared by their encoded form, so an int matches the int32 a pgx.CopyFromSource yields.

Custom types

Every mock owns a pgtype.Map, returned by TypeMap() of a mocked connection just like pgx.Conn.TypeMap (for a mocked pool use mock.AsConn().TypeMap()). Register custom types on it and give the column a DataTypeOID, and scanning goes through their codec:

	mock.TypeMap().RegisterType(&pgtype.Type{Name: "status", OID: statusOID, Codec: &pgtype.EnumCodec{}})
	col := mock.NewColumn("status")
	col.DataTypeOID = statusOID
	mock.ExpectQuery("SELECT status").WillReturnRows(mock.NewRowsWithColumnDefinition(*col).AddRow("active"))

LISTEN/NOTIFY

	mock.ExpectExec("LISTEN chat").WillReturnResult(pgxmock.NewResult("LISTEN", 0))
	mock.ExpectWaitForNotification().
		WillReturnNotification(&pgconn.Notification{Channel: "chat", Payload: "hello"})

Closed connections

By default a closed mock keeps serving expectations. Create it with pgxmock.ErrorOnClosedConnOption() to have every operation after Close fail with pgconn.ErrConnClosed, as pgx does.

What cannot be mocked

Conn(), PgConn(), LargeObjects() and the pool's Acquire, AcquireFunc and AcquireAllIdle hand out concrete pgx types that cannot be built outside of pgx. Have the code under test accept an interface of its own instead, or use PgxPoolIface.AsConn() to get a mocked connection from a mocked pool.

Run tests

go test -race ./...

Contributions

Feel free to open a pull request. Note, if you wish to contribute an extension to public (exported methods or types) - please open an issue before, to discuss whether these changes can be accepted. All backward incompatible changes are and will be treated cautiously

License

The three clause BSD license

Star History