Help with mocking returning pgx.Row from QueryRow method

[stori-rusher]

We are using separate database connections for read and write proxies, so I've created a DBConnector interface and a struct that implements it that look like this:

type DBConnector interface {
	Begin(ctx context.Context) (pgx.Tx, error)
	BeginFunc(ctx context.Context, f func(pgx.Tx) error) error
	CloseConnection()
	Exec(ctx context.Context, sql string, args ...interface{}) (pgconn.CommandTag, error)
	Query(ctx context.Context, sql string, args ...interface{}) (pgx.Rows, error)
	QueryRow(ctx context.Context, sql string, args ...interface{}) (pgx.Row, error)
	NewConnection(ctx context.Context, dp DBProxies) error
}

type AWSDB struct {
	reader *pgxpool.Pool
	writer *pgxpool.Pool
}

The AWSBD struct has methods like this to wrap the pgxpool.Pool methods using the intended reader or writer connection.

// QueryRow wraps the *pgxpool.Pool QueryRow method using the AWSDB reader connection
func (d *AWSDB) QueryRow(ctx context.Context, sql string, args ...interface{}) (pgx.Row, error) {
	if d.reader == nil {
		return nil, errors.New("reader pool not created")
	}

	return d.reader.QueryRow(ctx, sql, args...), nil
}

And so some code that implements this would look like this:

func getUserData(ctx context.Context, dbc db.DBConnector, c string) (row, error) {
	query := `
	SELECT
		ap.external_acct_id as external_acct_id,
		apt.user_status as user_status,
		ap.acct_uuid as acct_uuid,
		ap.user_id as user_id
	FROM account_profile ap
	JOIN app_user_type apt ON apt.user_id = ap.user_id
	WHERE ap.id = $1;
	`

	pgRow, err := dbc.QueryRow(ctx, query, c)
	if err != nil {
		return row{}, err
	}

	var r row
	if err := pgRow.Scan(
		&r.ExternalAcctID,
		&r.UserStatus,
		&r.AcctUUID,
		&r.UserID,
	); err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return row{}, errors.New("no user found with id " + c)
		} else {
			return row{}, err
		}
	}

	return r, nil
}

I'd like to be able to write unit tests around this getUserData function to check responses for things like no rows found, etc. My approach has been to create a mock db connection that implements my db.DBConnector interface like this:

type mockAWSDB struct {
	db.DBConnector
	reader pgxmock.PgxPoolIface
}

func (md *mockAWSDB) NewConnection(ctx context.Context, dp db.DBProxies) error {
	pool, err := pgxmock.NewPool()
	if err != nil {
		return err
	}

	switch dp {
	case db.Read:
		md.reader = pool
	default:
		return fmt.Errorf("db proxy " + string(db.Read) + " only supported")
	}

	return nil
}


// This is where I'd like to create the mock `pgx.Row` value to return expected values
func (md *mockAWSDB) QueryRow(ctx context.Context, sql string, args ...interface{}) (pgx.Row, error) {

}

I'm falling short here and don't know how to make this QueryRow method to return the pgx.Row value that I want using the underlying pgxmock.PgxPoolIface QueryRow method. Any help would be appreciated, thanks.

[pashagolub]

Hello,

the key trick here is to use interface instead of the real pgx objects, e.g.

type AWSDB struct {
	reader DBConnector
	writer DBConnector
}

Then in the tests instead of creating a real pgxpool.Pool you should create a pgxmock object. Check here the example:
https://github.com/pashagolub/pgxmock/blob/master/examples/basic/

Instead of using the real pgxpool.Pool as an argument in the recordStats() function, we use the PgxIface interface. Which allows us to use pgxmock during tests and the real connection during work.

[stori-rusher]

@pashagolub Thanks for the quick reply, I've made some updates to define the reader and writer as DBConnector interfaces as you suggested. I'm now running into an error with the query matching.

In my code, I have this doing the query to the db:

query := `
SELECT
	ap.external_acct_id as external_acct_id,
	apt.user_status as user_status,
	ap.acct_uuid as acct_uuid,
	ap.user_id as user_id
FROM ccm.account_profile ap
JOIN usr.app_user_type apt ON apt.user_id = ap.user_id
WHERE ap.clabe = $1;
`

// dbc is an instance of the DBConnector interface
pgRow = dbc.QueryRow(newCTX, query, c)

And so in my unit test I've set it up like this:

mdb := &mockAWSDB{}
if err := mdb.NewConnection(ctx, db.Read); err != nil {
	t.Fatalf("error making new db connection: %v", err)
}
defer mdb.Close()

query := `
	SELECT
		ap.external_acct_id as external_acct_id,
		apt.user_status as user_status,
		ap.acct_uuid as acct_uuid,
		ap.user_id as user_id
	FROM ccm.account_profile ap
	JOIN usr.app_user_type apt ON apt.user_id = ap.user_id
	WHERE ap.clabe = $1;
`

columns := []string{"acct_uuid", "user_status", "external_acct_id", "user_id"}

rs := pgxmock.NewRows(columns).
	AddRow("d6560a84-1fc2-4a26-832c-b6e3f1d3f6b0", "000000002281", "da3ddc9a-ab76-4f6e-a1f4-9f6dc4e3095f", 3000)

mdb.reader.ExpectQuery(query).WithArgs("646180224400022816").WillReturnRows(rs)

// we're working with an AWS Lambda handler
event := events.APIGatewayProxyRequest{
	Body: `
	{
		"abono": {
			"cuentaBeneficiario": "646180224400022816"
		}
	}
	`,
}

// this paymentHandler function just wraps the AWS handler func so I can pass the DBConnector interface in
handler := paymentHandler(mdb)

// this calls fires off the lambda handler, which eventually calls the above query.
res, err := handler(ctx, event)

When I run this test, the error I get back is as follows:

Query: could not match actual sql: "SELECT ap.external_acct_id as external_acct_id, apt.user_status as user_status, ap.acct_uuid as acct_uuid, ap.user_id as user_id FROM ccm.account_profile ap JOIN usr.app_user_type apt ON apt.user_id = ap.user_id WHERE ap.clabe = $1;" with expected regexp "SELECT ap.external_acct_id as external_acct_id, apt.user_status as user_status, ap.acct_uuid as acct_uuid, ap.user_id as user_id FROM ccm.account_profile ap JOIN usr.app_user_type apt ON apt.user_id = ap.user_id WHERE ap.clabe = $1;"

The two SQL statements match, but the regex is failing. Should I be formatting this a different way?

[pashagolub]

Your strings are not equal. Please, see the answer in #53