Merge branch 'master' of github.com:DATA-DOG/go-sqlmock

Author: Nikita Koryabkin

expectations.go

@@ -338,3 +338,47 @@ type queryBasedExpectation struct {
 	converter driver.ValueConverter
 	args      []driver.Value
 }
+
+func (e *queryBasedExpectation) attemptArgMatch(args []namedValue) (err error) {
+	// catch panic
+	defer func() {
+		if e := recover(); e != nil {
+			_, ok := e.(error)
+			if !ok {
+				err = fmt.Errorf(e.(string))
+			}
+		}
+	}()
+
+	err = e.argsMatches(args)
+	return
+}
+
+// ExpectedPing is used to manage *sql.DB.Ping expectations.
+// Returned by *Sqlmock.ExpectPing.
+type ExpectedPing struct {
+	commonExpectation
+	delay time.Duration
+}
+
+// WillDelayFor allows to specify duration for which it will delay result. May
+// be used together with Context.
+func (e *ExpectedPing) WillDelayFor(duration time.Duration) *ExpectedPing {
+	e.delay = duration
+	return e
+}
+
+// WillReturnError allows to set an error for expected database ping
+func (e *ExpectedPing) WillReturnError(err error) *ExpectedPing {
+	e.err = err
+	return e
+}
+
+// String returns string representation
+func (e *ExpectedPing) String() string {
+	msg := "ExpectedPing => expecting database Ping"
+	if e.err != nil {
+		msg += fmt.Sprintf(", which should return error: %s", e.err)
+	}
+	return msg
+}

options.go

@@ -20,3 +20,19 @@ func QueryMatcherOption(queryMatcher QueryMatcher) func(*sqlmock) error {
 		return nil
 	}
 }
+
+// MonitorPingsOption determines whether calls to Ping on the driver should be
+// observed and mocked.
+//
+// If true is passed, we will check these calls were expected. Expectations can
+// be registered using the ExpectPing() method on the mock.
+//
+// If false is passed or this option is omitted, calls to Ping will not be
+// considered when determining expectations and calls to ExpectPing will have
+// no effect.
+func MonitorPingsOption(monitorPings bool) func(*sqlmock) error {
+	return func(s *sqlmock) error {
+		s.monitorPings = monitorPings
+		return nil
+	}
+}

sqlmock.go

@@ -21,7 +21,6 @@ import (
 // for any kind of database action in order to mock
 // and test real database behavior.
 type Sqlmock interface {
-
 	// ExpectClose queues an expectation for this database
 	// action to be triggered. the *ExpectedClose allows
 	// to mock database response
@@ -57,6 +56,17 @@ type Sqlmock interface {
 	// the *ExpectedRollback allows to mock database response
 	ExpectRollback() *ExpectedRollback
 
+	// ExpectPing expected *sql.DB.Ping to be called.
+	// the *ExpectedPing allows to mock database response
+	//
+	// Ping support only exists in the SQL library in Go 1.8 and above.
+	// ExpectPing in Go <=1.7 will return an ExpectedPing but not register
+	// any expectations.
+	//
+	// You must enable pings using MonitorPingsOption for this to register
+	// any expectations.
+	ExpectPing() *ExpectedPing
+
 	// MatchExpectationsInOrder gives an option whether to match all
 	// expectations in the order they were set or not.
 	//
@@ -83,6 +93,7 @@ type sqlmock struct {
 	drv          *mockDriver
 	converter    driver.ValueConverter
 	queryMatcher QueryMatcher
+	monitorPings bool
 
 	expected []expectation
 }
@@ -104,6 +115,15 @@ func (c *sqlmock) open(options []func(*sqlmock) error) (*sql.DB, Sqlmock, error)
 	if c.queryMatcher == nil {
 		c.queryMatcher = QueryMatcherRegexp
 	}
+
+	if c.monitorPings {
+		// We call Ping on the driver shortly to verify startup assertions by
+		// driving internal behaviour of the sql standard library. We don't
+		// want this call to ping to be monitored for expectation purposes so
+		// temporarily disable.
+		c.monitorPings = false
+		defer func() { c.monitorPings = true }()
+	}
 	return db, c, db.Ping()
 }
 

sqlmock_before_go18.go

@@ -1,16 +1,13 @@
 // +build !go1.8
 
-/*
-Package sqlmock is a mock library implementing sql driver. Which has one and only
-purpose - to simulate any sql driver behavior in tests, without needing a real
-database connection. It helps to maintain correct **TDD** workflow.
+package sqlmock
 
-It does not require any modifications to your source code in order to test
-and mock database operations. Supports concurrency and multiple database mocking.
+import "log"
 
-The driver allows to mock any sql driver method behavior.
-*/
-package sqlmock
+func (c *sqlmock) ExpectPing() *ExpectedPing {
+	log.Println("ExpectPing has no effect on Go 1.7 or below")
+	return &ExpectedPing{}
+}
 
 func (c *sqlmock) exec(query string, args []namedValue) (*ExpectedExec, error) {
 	var expected *ExpectedExec

sqlmock_before_go18_test.go

@@ -0,0 +1,26 @@
+// +build !go1.8
+
+package sqlmock
+
+import (
+	"fmt"
+	"testing"
+	"time"
+)
+
+func TestSqlmockExpectPingHasNoEffect(t *testing.T) {
+	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()
+
+	e := mock.ExpectPing()
+
+	// Methods on the expectation can be called
+	e.WillDelayFor(time.Hour).WillReturnError(fmt.Errorf("an error"))
+
+	if err = mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("expected no error to be returned, but got '%s'", err)
+	}
+}

sqlmock_go18.go

@@ -6,6 +6,8 @@ import (
 	"context"
 	"database/sql/driver"
 	"errors"
+	"fmt"
+	"log"
 	"time"
 )
 
@@ -85,11 +87,57 @@ func (c *sqlmock) PrepareContext(ctx context.Context, query string) (driver.Stmt
 	return nil, err
 }
 
-// Implement the "Pinger" interface
-// for now we do not have a Ping expectation
-// may be something for the future
+// Implement the "Pinger" interface - the explicit DB driver ping was only added to database/sql in Go 1.8
 func (c *sqlmock) Ping(ctx context.Context) error {
-	return nil
+	if !c.monitorPings {
+		return nil
+	}
+
+	ex, err := c.ping()
+	if ex != nil {
+		select {
+		case <-ctx.Done():
+			return ErrCancelled
+		case <-time.After(ex.delay):
+		}
+	}
+
+	return err
+}
+
+func (c *sqlmock) ping() (*ExpectedPing, error) {
+	var expected *ExpectedPing
+	var fulfilled int
+	var ok bool
+	for _, next := range c.expected {
+		next.Lock()
+		if next.fulfilled() {
+			next.Unlock()
+			fulfilled++
+			continue
+		}
+
+		if expected, ok = next.(*ExpectedPing); ok {
+			break
+		}
+
+		next.Unlock()
+		if c.ordered {
+			return nil, fmt.Errorf("call to database Ping, was not expected, next expectation is: %s", next)
+		}
+	}
+
+	if expected == nil {
+		msg := "call to database Ping was not expected"
+		if fulfilled == len(c.expected) {
+			msg = "all expectations were already fulfilled, " + msg
+		}
+		return nil, fmt.Errorf(msg)
+	}
+
+	expected.triggered = true
+	expected.Unlock()
+	return expected, expected.err
 }
 
 // Implement the "StmtExecContext" interface
@@ -102,4 +150,14 @@ func (stmt *statement) QueryContext(ctx context.Context, args []driver.NamedValu
 	return stmt.conn.QueryContext(ctx, stmt.query, args)
 }
 
+func (c *sqlmock) ExpectPing() *ExpectedPing {
+	if !c.monitorPings {
+		log.Println("ExpectPing will have no effect as monitoring pings is disabled. Use MonitorPingsOption to enable.")
+		return nil
+	}
+	e := &ExpectedPing{}
+	c.expected = append(c.expected, e)
+	return e
+}
+
 // @TODO maybe add ExpectedBegin.WithOptions(driver.TxOptions)