Adding feature to allow repeatable expectations

sqlmock.go

@@ -73,10 +73,23 @@ type Sqlmock interface {
 	// in any order. Or otherwise if switched to true, any unmatched
 	// expectations will be expected in order
 	MatchExpectationsInOrder(bool)
+
+	// AllowRepeatedExpectationMatching gives an option whether or not to
+	// allow expectations to be matched more than once.
+	//
+	// By default it is set to - false.
+	//
+	// This option may be turned on anytime during tests. As soon
+	// as it is switched to true, expectations will be allowed to match
+	// regardless of it has been previously matched against.
+	//
+	// When setting this true, consider if you will need to set MatchExpectationsInOrder(false)
+	AllowRepeatedExpectationMatching(bool)
 }
 
 type sqlmock struct {
 	ordered    bool
+	repeatable bool
 	dsn        string
 	opened     int
 	drv        *mockDriver
@@ -102,6 +115,10 @@ func (c *sqlmock) MatchExpectationsInOrder(b bool) {
 	c.ordered = b
 }
 
+func (c *sqlmock) AllowRepeatedExpectationMatching(b bool) {
+	c.repeatable = b
+}
+
 // Close a mock database driver connection. It may or may not
 // be called depending on the sircumstances, but if it is called
 // there must be an *ExpectedClose expectation satisfied.
@@ -121,10 +138,12 @@ func (c *sqlmock) Close() error {
 	for _, next := range c.expected {
 		next.Lock()
 		if next.fulfilled() {
-			next.Unlock()
 			fulfilled++
+			if !c.repeatable {
+				next.Unlock()
 				continue
 			}
+		}
 
 		if expected, ok = next.(*ExpectedClose); ok {
 			break
@@ -185,8 +204,10 @@ func (c *sqlmock) begin() (*ExpectedBegin, error) {
 		if next.fulfilled() {
 			next.Unlock()
 			fulfilled++
+			if !c.repeatable {
 				continue
 			}
+		}
 
 		if expected, ok = next.(*ExpectedBegin); ok {
 			break
@@ -246,8 +267,10 @@ func (c *sqlmock) exec(query string, args []namedValue) (*ExpectedExec, error) {
 		if next.fulfilled() {
 			next.Unlock()
 			fulfilled++
+			if !c.repeatable {
 				continue
 			}
+		}
 
 		if c.ordered {
 			if expected, ok = next.(*ExpectedExec); ok {
@@ -322,10 +345,13 @@ func (c *sqlmock) prepare(query string) (*ExpectedPrepare, error) {
 	for _, next := range c.expected {
 		next.Lock()
 		if next.fulfilled() {
-			next.Unlock()
 			fulfilled++
+
+			if !c.repeatable {
+				next.Unlock()
 				continue
 			}
+		}
 
 		if c.ordered {
 			if expected, ok = next.(*ExpectedPrepare); ok {
@@ -401,10 +427,12 @@ func (c *sqlmock) query(query string, args []namedValue) (*ExpectedQuery, error)
 	for _, next := range c.expected {
 		next.Lock()
 		if next.fulfilled() {
-			next.Unlock()
 			fulfilled++
+			if !c.repeatable {
+				next.Unlock()
 				continue
 			}
+		}
 
 		if c.ordered {
 			if expected, ok = next.(*ExpectedQuery); ok {
@@ -448,6 +476,14 @@ func (c *sqlmock) query(query string, args []namedValue) (*ExpectedQuery, error)
 	if expected.rows == nil {
 		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)
 	}
+
+	// reset rows for next use if allowed
+	if rs, ok := expected.rows.(*rowSets); ok {
+		rs.pos = 0
+		for _, set := range rs.sets {
+			set.pos = 0
+		}
+	}
 	return expected, nil
 }
 
@@ -481,8 +517,10 @@ func (c *sqlmock) Commit() error {
 		if next.fulfilled() {
 			next.Unlock()
 			fulfilled++
+			if !c.repeatable {
 				continue
 			}
+		}
 
 		if expected, ok = next.(*ExpectedCommit); ok {
 			break
@@ -516,8 +554,10 @@ func (c *sqlmock) Rollback() error {
 		if next.fulfilled() {
 			next.Unlock()
 			fulfilled++
+			if !c.repeatable {
 				continue
 			}
+		}
 
 		if expected, ok = next.(*ExpectedRollback); ok {
 			break

sqlmock_test.go

@@ -354,6 +354,57 @@ func TestPreparedQueryExecutions(t *testing.T) {
 	}
 }
 
+func TestPreparedQueryMultipleExecutions(t *testing.T) {
+	t.Parallel()
+	db, mock, err := New()
+	if err != nil {
+		t.Errorf("an error '%s' was not expected when opening a stub database connection", err)
+	}
+	defer db.Close()
+
+	mock.MatchExpectationsInOrder(false)
+	mock.AllowRepeatedExpectationMatching(true)
+
+	rs := NewRows([]string{"id", "title"}).FromCSVString("5,hello world")
+	mock.ExpectPrepare("SELECT (.+) FROM articles WHERE id = ?").ExpectQuery().
+		WithArgs(5).
+		WillReturnRows(rs)
+
+	stmt, err := db.Prepare("SELECT id, title FROM articles WHERE id = ?")
+	if err != nil {
+		t.Errorf("error '%s' was not expected while creating a prepared statement", err)
+	}
+
+	stmt2, err := db.Prepare("SELECT id, title FROM articles WHERE id = ?")
+	if err != nil {
+		t.Errorf("error '%s' was not expected while creating a prepared statement", err)
+	}
+
+	var id1, id2 int
+	var title1, title2 string
+	err = stmt.QueryRow(5).Scan(&id1, &title1)
+	if err != nil {
+		t.Errorf("error '%s' was not expected querying row from statement and scanning", err)
+	}
+
+	err = stmt2.QueryRow(5).Scan(&id2, &title2)
+	if err != nil {
+		t.Errorf("error '%s' was not expected querying row from statement and scanning", err)
+	}
+
+	if id1 != 5 || id2 != 5 {
+		t.Errorf("expected mocked id to be 5, but got %d instead", id2)
+	}
+
+	if title1 != "hello world" || title2 != "hello world" {
+		t.Errorf("expected mocked title to be 'hello world', but got '%s' instead", title2)
+	}
+
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("there were unfulfilled expections: %s", err)
+	}
+}
+
 func TestUnorderedPreparedQueryExecutions(t *testing.T) {
 	t.Parallel()
 	db, mock, err := New()