From 005f8530f43c63be2697698214f71cbb824cdd3f Mon Sep 17 00:00:00 2001
From: Pavlo Golub <pavlo.golub@gmail.com>
Date: Wed, 18 Oct 2023 11:02:49 +0200
Subject: [PATCH] improve test coverage

---
 argument_test.go     |  5 +++--
 expectations.go      |  3 ++-
 expectations_test.go | 19 +++++++++++++++++++
 pgxmock_test.go      |  2 +-
 4 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/argument_test.go b/argument_test.go
index ea02a68..f268dbc 100644
--- a/argument_test.go
+++ b/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)
 }
 
diff --git a/expectations.go b/expectations.go
index bcfd550..c2e3107 100644
--- a/expectations.go
+++ b/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()
 }
diff --git a/expectations_test.go b/expectations_test.go
index e8e9eac..cc60034 100644
--- a/expectations_test.go
+++ b/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())
+}
diff --git a/pgxmock_test.go b/pgxmock_test.go
index 3ab286c..0e60170 100644
--- a/pgxmock_test.go
+++ b/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 (?, ?)")