From a7298db970cca2114e8709b4194c4732af77fd9c Mon Sep 17 00:00:00 2001
From: Pavlo Golub <pavlo.golub@gmail.com>
Date: Fri, 19 Jan 2024 14:40:49 +0100
Subject: [PATCH] [-] fix `Deallocate()` to return `expected.deallocateErr`
 (#185)

---
 pgxmock.go      | 11 ++++++-----
 pgxmock_test.go |  6 ++++--
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/pgxmock.go b/pgxmock.go
index b3e4670..3f1ad2b 100644
--- a/pgxmock.go
+++ b/pgxmock.go
@@ -387,8 +387,11 @@ func (c *pgxmock) Deallocate(ctx context.Context, name string) error {
 	if expected == nil {
 		return fmt.Errorf("Deallocate: prepared statement name '%s' doesn't exist", name)
 	}
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
 	expected.deallocated = true
-	return expected.waitForDelay(ctx)
+	return expected.deallocateErr
 }
 
 func (c *pgxmock) Commit(ctx context.Context) error {
@@ -480,11 +483,9 @@ func (c *pgxmock) Ping(ctx context.Context) (err error) {
 }
 
 func (c *pgxmock) Reset() {
-	ex, err := findExpectation[*ExpectedReset](c, "Reset()")
-	if err != nil {
-		return
+	if ex, err := findExpectation[*ExpectedReset](c, "Reset()"); err == nil {
+		_ = ex.waitForDelay(context.Background())
 	}
-	_ = ex.waitForDelay(context.Background())
 }
 
 type expectationType[t any] interface {
diff --git a/pgxmock_test.go b/pgxmock_test.go
index 900adf1..3adc9ae 100644
--- a/pgxmock_test.go
+++ b/pgxmock_test.go
@@ -240,9 +240,9 @@ func TestPrepareExpectations(t *testing.T) {
 	t.Parallel()
 	mock, _ := NewConn()
 	a := assert.New(t)
-
+	expErr := errors.New("invaders must die")
 	mock.ExpectPrepare("foo", "SELECT (.+) FROM articles WHERE id = ?").
-		WillReturnCloseError(errors.New("invaders must die")).
+		WillReturnCloseError(expErr).
 		WillDelayFor(1 * time.Second)
 
 	stmt, err := mock.Prepare(context.Background(), "baz", "SELECT (.+) FROM articles WHERE id = ?")
@@ -252,6 +252,8 @@ func TestPrepareExpectations(t *testing.T) {
 	stmt, err = mock.Prepare(context.Background(), "foo", "SELECT (.+) FROM articles WHERE id = $1")
 	a.NoError(err)
 	a.NotNil(stmt)
+	err = mock.Deallocate(context.Background(), "foo")
+	a.EqualError(err, expErr.Error())
 
 	// expect something else, w/o ExpectPrepare()
 	var id int