From 4689a4afbfddf1c028e81cd800ad65f10ee1d439 Mon Sep 17 00:00:00 2001
From: Jamie Cuthill <jamie.cuthill@sainsburys.co.uk>
Date: Fri, 26 Feb 2016 16:07:19 +0000
Subject: [PATCH] When Exec/Query is expected to error correctly mark
 expectation as triggered

 - Fixes regression from issue 37 'fix'
---
 sqlmock.go      |  7 +++++--
 sqlmock_test.go | 36 ++++++++++++++++++++++++++++++++----
 2 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/sqlmock.go b/sqlmock.go
index 8c344cd..9053dcd 100644
--- a/sqlmock.go
+++ b/sqlmock.go
@@ -240,6 +240,8 @@ func (c *sqlmock) Exec(query string, args []driver.Value) (res driver.Result, er
 		return nil, fmt.Errorf("exec query '%s', arguments do not match: %s", query, err)
 	}
 
+	expected.triggered = true
+
 	if expected.err != nil {
 		return nil, expected.err // mocked to return error
 	}
@@ -247,7 +249,7 @@ func (c *sqlmock) Exec(query string, args []driver.Value) (res driver.Result, er
 	if expected.result == nil {
 		return nil, fmt.Errorf("exec query '%s' with args %+v, must return a database/sql/driver.result, but it was not set for expectation %T as %+v", query, args, expected, expected)
 	}
-	expected.triggered = true
+
 	return expected.result, err
 }
 
@@ -349,6 +351,8 @@ func (c *sqlmock) Query(query string, args []driver.Value) (rw driver.Rows, err
 		return nil, fmt.Errorf("exec query '%s', arguments do not match: %s", query, err)
 	}
 
+	expected.triggered = true
+
 	if expected.err != nil {
 		return nil, expected.err // mocked to return error
 	}
@@ -357,7 +361,6 @@ func (c *sqlmock) Query(query string, args []driver.Value) (rw driver.Rows, err
 		return nil, fmt.Errorf("query '%s' with args %+v, must return a database/sql/driver.rows, but it was not set for expectation %T as %+v", query, args, expected, expected)
 	}
 
-	expected.triggered = true
 	return expected.rows, err
 }
 
diff --git a/sqlmock_test.go b/sqlmock_test.go
index 31f57d4..d040213 100644
--- a/sqlmock_test.go
+++ b/sqlmock_test.go
@@ -679,14 +679,42 @@ func TestRunExecsWithOrderedShouldNotMeetAllExpectations(t *testing.T) {
 // see #37 issue
 func TestRunQueriesWithOrderedShouldNotMeetAllExpectations(t *testing.T) {
 	db, dbmock, _ := New()
-	dbmock.ExpectQuery("THE FIRST EXEC")
-	dbmock.ExpectQuery("THE SECOND EXEC")
+	dbmock.ExpectQuery("THE FIRST QUERY")
+	dbmock.ExpectQuery("THE SECOND QUERY")
 
-	_, _ = db.Query("THE FIRST EXEC")
-	_, _ = db.Query("THE WRONG EXEC")
+	_, _ = db.Query("THE FIRST QUERY")
+	_, _ = db.Query("THE WRONG QUERY")
 
 	err := dbmock.ExpectationsWereMet()
 	if err == nil {
 		t.Fatal("was expecting an error, but there wasn't any")
 	}
 }
+
+func TestRunExecsWithExpectedErrorMeetsExpectations(t *testing.T) {
+	db, dbmock, _ := New()
+	dbmock.ExpectExec("THE FIRST EXEC").WillReturnError(fmt.Errorf("big bad bug"))
+	dbmock.ExpectExec("THE SECOND EXEC").WillReturnResult(NewResult(0, 0))
+
+	_, _ = db.Exec("THE FIRST EXEC")
+	_, _ = db.Exec("THE SECOND EXEC")
+
+	err := dbmock.ExpectationsWereMet()
+	if err != nil {
+		t.Fatalf("all expectations should be met: %s", err)
+	}
+}
+
+func TestRunQueryWithExpectedErrorMeetsExpectations(t *testing.T) {
+	db, dbmock, _ := New()
+	dbmock.ExpectQuery("THE FIRST QUERY").WillReturnError(fmt.Errorf("big bad bug"))
+	dbmock.ExpectQuery("THE SECOND QUERY").WillReturnRows(NewRows([]string{"col"}).AddRow(1))
+
+	_, _ = db.Query("THE FIRST QUERY")
+	_, _ = db.Query("THE SECOND QUERY")
+
+	err := dbmock.ExpectationsWereMet()
+	if err != nil {
+		t.Fatalf("all expectations should be met: %s", err)
+	}
+}