Mocking errors with Batch queries

[svennis94]

Describe the bug
We are implementing batch queries as in some cases we execute a batch of queries that can benefit from a single round trip.

When writing tests we wanted to mock errors as well when we noticed that when we do:

eb := mock.Exp.ExpectBatch()
eb.ExpectQuery("SELECT").WillReturnError(errors.New("Some Error"))
eb.ExpectQuery("SELECT").WillReturnRows(pgxmock.NewRows([]string{"second"}).AddRow(2))

The tests would fail with the message but it did return the expected error from the batch!

    Returned error: "Some Error" // Correct error returned from the batch!

    db_test.go:154: there were unfulfilled expectations: there is a remaining expectation which was not matched: ExpectedQuery => expecting call to Query() or to QueryRow():
        	- matches sql: 'SELECT'
        	- is without arguments
        	- returns data:
        		row 0 - [2]

So we thought, this can be an expected error since the first query returns an error, the second is not executed (even though a batch should execute all of them I believe). When we removed the second ExpectQuery we got a different error:

    Returned error: "call to method Batch() was not expected"

    db_test.go:154: there were unfulfilled expectations: there is a remaining expectation which was not matched: ExpectedBatch => expecting call to SendBatch()

This error is expected I believe as the batch has two select statements, to the "expector" is not expecting a Batch() as the second query was not expected.

The final attempt we did was to return the error on the second query instead:

eb := mock.Exp.ExpectBatch()
eb.ExpectQuery("SELECT").WillReturnRows(pgxmock.NewRows([]string{"first"}).AddRow(1))
eb.ExpectQuery("SELECT").WillReturnError(errors.New("Some Error"))

This made the testcase pass without expectation errors.

To Reproduce
A bit of pseudo code to simplify a testcase

Implementation on the database layer:

func (db *db) TestBatch() (int, error) {
	batch := &pgx.Batch{}

	first := 0
	batch.Queue("SELECT 1;").QueryRow(func(rows pgx.Row) error {
		return rows.Scan(&first)
	})
	second := 0
	batch.Queue("SELECT 2;").QueryRow(func(rows pgx.Row) error {
		return rows.Scan(&second)
	})

	return first + second, db.pool.SendBatch(context.Background(), batch).Close()
}

Testcases:

func TestNewBatchErrSuccess(t *testing.T) {
	dbl := tearUpMock(t)
	defer tearDownMock(t, dbl)

	// Test success
	eb := mock.Exp.ExpectBatch()
	eb.ExpectQuery("SELECT").WillReturnRows(pgxmock.NewRows([]string{"first"}).AddRow(1))
	eb.ExpectQuery("SELECT").WillReturnRows(pgxmock.NewRows([]string{"second"}).AddRow(2))

	data, err := dbl.TestBatch()
	assert.Nil(t, err)
	assert.Equal(t, 3, data)
}

// Great success!
func TestNewBatchErrorSuccess(t *testing.T) {
	dbl := tearUpMock(t)
	defer tearDownMock(t, dbl)

	// Test error
	eb := mock.Exp.ExpectBatch()
	eb.ExpectQuery("SELECT").WillReturnRows(pgxmock.NewRows([]string{"first"}).AddRow(1))
	eb.ExpectQuery("SELECT").WillReturnError(errors.New("Some Error"))

	data, err := dbl.TestBatch()
	assert.Equal(t, "Some Error", err.Error())
	assert.Equal(t, 1, data)
}

// Expected to succeed, possible bug.
func TestNewBatchErrorFail1(t *testing.T) {
	dbl := tearUpMock(t)
	defer tearDownMock(t, dbl)

	// Test error
	eb := mock.Exp.ExpectBatch()
	eb.ExpectQuery("SELECT").WillReturnError(errors.New("Some Error"))
	eb.ExpectQuery("SELECT").WillReturnRows(pgxmock.NewRows([]string{"second"}).AddRow(2))

	data, err := dbl.TestBatch()
	assert.Equal(t, "Some Error", err.Error())
	assert.Equal(t, 0, data)
}

// Expected to fail due to missing ExpectQuery, not a bug.
func TestNewBatchErrorFail2(t *testing.T) {
	dbl := tearUpMock(t)
	defer tearDownMock(t, dbl)

	// Test error
	eb := mock.Exp.ExpectBatch()
	eb.ExpectQuery("SELECT").WillReturnError(errors.New("Some Error"))

	data, err := dbl.TestBatch()
	assert.Equal(t, "Some Error", err.Error())
	assert.Equal(t, 0, data)
}

Expected behavior
I would expect the batch to return the expected error but not complain about missing expectations. I would expect testcase TestNewBatchErrorFail1 to succeed. But this might be a limitation of batch queries with PGX as well, I am not sure about this 🙂

[pashagolub]

Thanks! I will investigate the issue!