1 files changed, 57 insertions, 4 deletions

diff --git a/src/pkg/database/sql/fakedb_test.go b/src/pkg/database/sql/fakedb_test.go
index 55597f7de..d900e2ceb 100644
--- a/src/pkg/database/sql/fakedb_test.go
+++ b/src/pkg/database/sql/fakedb_test.go
@@ -13,6 +13,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"testing"
 	"time"
 )
 
@@ -34,9 +35,10 @@ var _ = log.Printf
 // When opening a fakeDriver's database, it starts empty with no
 // tables.  All tables and data are stored in memory only.
 type fakeDriver struct {
-	mu        sync.Mutex
-	openCount int
-	dbs       map[string]*fakeDB
+	mu         sync.Mutex // guards 3 following fields
+	openCount  int        // conn opens
+	closeCount int        // conn closes
+	dbs        map[string]*fakeDB
 }
 
 type fakeDB struct {
@@ -229,7 +231,43 @@ func (c *fakeConn) Begin() (driver.Tx, error) {
 	return c.currTx, nil
 }
 
-func (c *fakeConn) Close() error {
+var hookPostCloseConn struct {
+	sync.Mutex
+	fn func(*fakeConn, error)
+}
+
+func setHookpostCloseConn(fn func(*fakeConn, error)) {
+	hookPostCloseConn.Lock()
+	defer hookPostCloseConn.Unlock()
+	hookPostCloseConn.fn = fn
+}
+
+var testStrictClose *testing.T
+
+// setStrictFakeConnClose sets the t to Errorf on when fakeConn.Close
+// fails to close. If nil, the check is disabled.
+func setStrictFakeConnClose(t *testing.T) {
+	testStrictClose = t
+}
+
+func (c *fakeConn) Close() (err error) {
+	drv := fdriver.(*fakeDriver)
+	defer func() {
+		if err != nil && testStrictClose != nil {
+			testStrictClose.Errorf("failed to close a test fakeConn: %v", err)
+		}
+		hookPostCloseConn.Lock()
+		fn := hookPostCloseConn.fn
+		hookPostCloseConn.Unlock()
+		if fn != nil {
+			fn(c, err)
+		}
+		if err == nil {
+			drv.mu.Lock()
+			drv.closeCount++
+			drv.mu.Unlock()
+		}
+	}()
 	if c.currTx != nil {
 		return errors.New("can't close fakeConn; in a Transaction")
 	}
@@ -424,6 +462,12 @@ func (s *fakeStmt) ColumnConverter(idx int) driver.ValueConverter {
 }
 
 func (s *fakeStmt) Close() error {
+	if s.c == nil {
+		panic("nil conn in fakeStmt.Close")
+	}
+	if s.c.db == nil {
+		panic("in fakeStmt.Close, conn's db is nil (already closed)")
+	}
 	if !s.closed {
 		s.c.incrStat(&s.c.stmtsClosed)
 		s.closed = true
@@ -515,6 +559,15 @@ func (s *fakeStmt) Query(args []driver.Value) (driver.Rows, error) {
 	if !ok {
 		return nil, fmt.Errorf("fakedb: table %q doesn't exist", s.table)
 	}
+
+	if s.table == "magicquery" {
+		if len(s.whereCol) == 2 && s.whereCol[0] == "op" && s.whereCol[1] == "millis" {
+			if args[0] == "sleep" {
+				time.Sleep(time.Duration(args[1].(int64)) * time.Millisecond)
+			}
+		}
+	}
+
 	t.mu.Lock()
 	defer t.mu.Unlock()