closes #114 allow expecting rows to be closed

Author: l3pp4rd

README.md

@@ -213,6 +213,7 @@ It only asserts that argument is of `time.Time` type.
 
 ## Change Log
 
+- **2018-12-11** - added expectation of Rows to be closed, while mocking expected query.
 - **2018-12-11** - introduced an option to provide **QueryMatcher** in order to customize SQL query matching.
 - **2017-09-01** - it is now possible to expect that prepared statement will be closed,
   using **ExpectedPrepare.WillBeClosed**.

expectations.go

@@ -125,8 +125,10 @@ func (e *ExpectedRollback) String() string {
 // Returned by *Sqlmock.ExpectQuery.
 type ExpectedQuery struct {
 	queryBasedExpectation
-	rows  driver.Rows
-	delay time.Duration
+	rows             driver.Rows
+	delay            time.Duration
+	rowsMustBeClosed bool
+	rowsWereClosed   bool
 }
 
 // WithArgs will match given expected args to actual database query arguments.
@@ -137,6 +139,12 @@ func (e *ExpectedQuery) WithArgs(args ...driver.Value) *ExpectedQuery {
 	return e
 }
 
+// RowsWillBeClosed expects this query rows to be closed.
+func (e *ExpectedQuery) RowsWillBeClosed() *ExpectedQuery {
+	e.rowsMustBeClosed = true
+	return e
+}
+
 // WillReturnError allows to set an error for expected database query
 func (e *ExpectedQuery) WillReturnError(err error) *ExpectedQuery {
 	e.err = err

expectations_before_go18.go

@@ -11,7 +11,7 @@ import (
 // WillReturnRows specifies the set of resulting rows that will be returned
 // by the triggered query
 func (e *ExpectedQuery) WillReturnRows(rows *Rows) *ExpectedQuery {
-	e.rows = &rowSets{sets: []*Rows{rows}}
+	e.rows = &rowSets{sets: []*Rows{rows}, ex: e}
 	return e
 }
 

expectations_go18.go

@@ -16,7 +16,7 @@ func (e *ExpectedQuery) WillReturnRows(rows ...*Rows) *ExpectedQuery {
 	for i, r := range rows {
 		sets[i] = r
 	}
-	e.rows = &rowSets{sets: sets}
+	e.rows = &rowSets{sets: sets, ex: e}
 	return e
 }
 

rows.go

@@ -22,13 +22,15 @@ var CSVColumnParser = func(s string) []byte {
 type rowSets struct {
 	sets []*Rows
 	pos  int
+	ex   *ExpectedQuery
 }
 
 func (rs *rowSets) Columns() []string {
 	return rs.sets[rs.pos].cols
 }
 
 func (rs *rowSets) Close() error {
+	rs.ex.rowsWereClosed = true
 	return rs.sets[rs.pos].closeErr
 }
 

rows_test.go

@@ -88,6 +88,29 @@ func ExampleRows_closeError() {
 	// Output: got error: close error
 }
 
+func ExampleRows_expectToBeClosed() {
+	db, mock, err := New()
+	if err != nil {
+		fmt.Println("failed to open sqlmock database:", err)
+	}
+	defer db.Close()
+
+	rows := NewRows([]string{"id", "title"}).AddRow(1, "john")
+	mock.ExpectQuery("SELECT").WillReturnRows(rows).RowsWillBeClosed()
+
+	db.Query("SELECT")
+
+	if err := mock.ExpectationsWereMet(); err != nil {
+		fmt.Println("got error:", err)
+	}
+
+	// Output: got error: expected query rows to be closed, but it was not: ExpectedQuery => expecting Query, QueryContext or QueryRow which:
+	//   - matches sql: 'SELECT'
+	//   - is without arguments
+	//   - should return rows:
+	//     row 0 - [1 john]
+}
+
 func ExampleRows_customDriverValue() {
 	db, mock, err := New()
 	if err != nil {
@@ -184,6 +207,31 @@ func TestRowsCloseError(t *testing.T) {
 	}
 }
 
+func TestRowsClosed(t *testing.T) {
+	t.Parallel()
+	db, mock, err := New()
+	if err != nil {
+		t.Fatalf("an error '%s' was not expected when opening a stub database connection", err)
+	}
+	defer db.Close()
+
+	rows := NewRows([]string{"id"}).AddRow(1)
+	mock.ExpectQuery("SELECT").WillReturnRows(rows).RowsWillBeClosed()
+
+	rs, err := db.Query("SELECT")
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+
+	if err := rs.Close(); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Fatal(err)
+	}
+}
+
 func TestQuerySingleRow(t *testing.T) {
 	t.Parallel()
 	db, mock, err := New()

sqlmock.go

@@ -176,6 +176,13 @@ func (c *sqlmock) ExpectationsWereMet() error {
 				return fmt.Errorf("expected prepared statement to be closed, but it was not: %s", prep)
 			}
 		}
+
+		// must check whether all expected queried rows are closed
+		if query, ok := e.(*ExpectedQuery); ok {
+			if query.rowsMustBeClosed && !query.rowsWereClosed {
+				return fmt.Errorf("expected query rows to be closed, but it was not: %s", query)
+			}
+		}
 	}
 	return nil
 }