commit 2a64ebfc6d247e9177190f862b0947ac98c5eb43
author Daniel Theophanes <kardianos@gmail.com> Wed Nov 30 09:30:31 2016 -0800
committer Brad Fitzpatrick <bradfitz@golang.org> Wed Nov 30 18:40:03 2016 +0000
tree 3311bb34320bbf491f893e949eed4b96f8d4c6e6
parent f6bff1d587c3523d2d1b98f0737a922ac9b5becc

database/sql: deflake query cancel tests

Rather then using a sleep in the fake DB, go to a channel
select and wait for the context to be done.

Fixes #18115

Change-Id: I6bc3a29db58c568d0a7ea06c2a354c18c9e798b2
Reviewed-on: https://go-review.googlesource.com/33712
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>

src/database/sql/fakedb_test.go

diff --git a/src/database/sql/fakedb_test.go b/src/database/sql/fakedb_test.go
index 9de9289..416b97d 100644
--- a/src/database/sql/fakedb_test.go
+++ b/src/database/sql/fakedb_test.go
@@ -511,6 +511,10 @@
 var hookPrepareBadConn func() bool
 
 func (c *fakeConn) Prepare(query string) (driver.Stmt, error) {
+	panic("use PrepareContext")
+}
+
+func (c *fakeConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
 	c.numPrepare++
 	if c.db == nil {
 		panic("nil c.db; conn = " + fmt.Sprintf("%#v", c))
@@ -549,7 +553,13 @@
 		parts = parts[1:]
 
 		if stmt.wait > 0 {
-			time.Sleep(stmt.wait)
+			wait := time.NewTimer(stmt.wait)
+			select {
+			case <-wait.C:
+			case <-ctx.Done():
+				wait.Stop()
+				return nil, ctx.Err()
+			}
 		}
 
 		c.incrStat(&c.stmtsMade)

src/database/sql/sql_test.go

diff --git a/src/database/sql/sql_test.go b/src/database/sql/sql_test.go
index c0f2cf2..27fb765 100644
--- a/src/database/sql/sql_test.go
+++ b/src/database/sql/sql_test.go
@@ -24,6 +24,17 @@
 		c  *driverConn
 	}
 	freedFrom := make(map[dbConn]string)
+	var mu sync.Mutex
+	getFreedFrom := func(c dbConn) string {
+		mu.Lock()
+		defer mu.Unlock()
+		return freedFrom[c]
+	}
+	setFreedFrom := func(c dbConn, s string) {
+		mu.Lock()
+		defer mu.Unlock()
+		freedFrom[c] = s
+	}
 	putConnHook = func(db *DB, c *driverConn) {
 		idx := -1
 		for i, v := range db.freeConn {
@@ -36,10 +47,10 @@
 			// print before panic, as panic may get lost due to conflicting panic
 			// (all goroutines asleep) elsewhere, since we might not unlock
 			// the mutex in freeConn here.
-			println("double free of conn. conflicts are:\nA) " + freedFrom[dbConn{db, c}] + "\n\nand\nB) " + stack())
+			println("double free of conn. conflicts are:\nA) " + getFreedFrom(dbConn{db, c}) + "\n\nand\nB) " + stack())
 			panic("double free of conn.")
 		}
-		freedFrom[dbConn{db, c}] = stack()
+		setFreedFrom(dbConn{db, c}, stack())
 	}
 }
 
@@ -344,7 +355,7 @@
 	// This will trigger the *fakeConn.Prepare method which will take time
 	// performing the query. The ctxDriverPrepare func will check the context
 	// after this and close the rows and return an error.
-	_, err := db.QueryContext(ctx, "WAIT|30ms|SELECT|people|age,name|")
+	_, err := db.QueryContext(ctx, "WAIT|1s|SELECT|people|age,name|")
 	if err != context.DeadlineExceeded {
 		t.Fatalf("expected QueryContext to error with context deadline exceeded but returned %v", err)
 	}
@@ -372,7 +383,7 @@
 	// This will trigger the *fakeConn.Prepare method which will take time
 	// performing the query. The ctxDriverPrepare func will check the context
 	// after this and close the rows and return an error.
-	_, err = tx.QueryContext(ctx, "WAIT|30ms|SELECT|people|age,name|")
+	_, err = tx.QueryContext(ctx, "WAIT|1s|SELECT|people|age,name|")
 	if err != context.DeadlineExceeded {
 		t.Fatalf("expected QueryContext to error with context deadline exceeded but returned %v", err)
 	}