Mocking errors with Batch queries

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 :slightly_smiling_face:

pashagolub / pgxmock

Mocking errors with Batch queries #222