improve test coverage

argument_test.go

@@ -99,9 +99,10 @@ func TestExpectQueryRewriterFail(t *testing.T) {
 		t.Errorf("an error '%s' was not expected when opening a stub database connection", err)
 	}
 
-	mock.ExpectExec(`INSERT INTO users\(username\) VALUES \(\@user\)`).
+	mock.ExpectQuery(`INSERT INTO users\(username\) VALUES \(\@user\)`).
+		WithRewrittenSQL(`INSERT INTO users\(username\) VALUES \(\$1\)`).
 		WithArgs(failQryRW{})
-	_, err = mock.Exec(context.Background(), "INSERT INTO users(username) VALUES (@user)", "baz")
+	_, err = mock.Query(context.Background(), "INSERT INTO users(username) VALUES (@user)", "baz")
 	assert.Error(t, err)
 }
 

expectations.go

@@ -280,7 +280,8 @@ func (e *ExpectedPrepare) WillReturnCloseError(err error) *ExpectedPrepare {
 }
 
 // WillBeClosed is for backward compatibility only and will be removed soon.
-// One should use WillBeDeallocated() instead
+//
+// Deprecated: One should use WillBeDeallocated() instead.
 func (e *ExpectedPrepare) WillBeClosed() *ExpectedPrepare {
 	return e.WillBeDeallocated()
 }

expectations_test.go

@@ -249,3 +249,22 @@ func TestMissingWithArgs(t *testing.T) {
 		t.Error("expectation was not matched error was expected")
 	}
 }
+
+func TestWithRewrittenSQL(t *testing.T) {
+	t.Parallel()
+	mock, err := NewConn()
+	a := assert.New(t)
+	a.NoError(err)
+
+	mock.ExpectQuery(`INSERT INTO users\(username\) VALUES \(\@user\)`).
+		WithArgs(pgx.NamedArgs{"user": "John"}).
+		WithRewrittenSQL(`INSERT INTO users\(username\) VALUES \(\$1\)`).
+		WillReturnRows()
+
+	_, err = mock.Query(context.Background(),
+		"INSERT INTO users(username) VALUES (@user)",
+		pgx.NamedArgs{"user": "John"},
+	)
+	a.NoError(err)
+	a.NoError(mock.ExpectationsWereMet())
+}

pgxmock_test.go

@@ -1019,7 +1019,7 @@ func TestPreparedStatementCloseExpectation(t *testing.T) {
 	mock, _ := NewConn()
 	a := assert.New(t)
 
-	ep := mock.ExpectPrepare("foo", "INSERT INTO ORDERS").WillBeDeallocated()
+	ep := mock.ExpectPrepare("foo", "INSERT INTO ORDERS").WillBeClosed()
 	ep.ExpectExec().WithArgs(AnyArg(), AnyArg()).WillReturnResult(NewResult("UPDATE", 1))
 
 	stmt, err := mock.Prepare(context.Background(), "foo", "INSERT INTO ORDERS(ID, STATUS) VALUES (?, ?)")