database/sql: let tests wait for db pool to come to expected state

Slower builders were failing TestQueryContext because the cancel
and return to conn pool happens async. TestQueryContext already
uses a wait method for this reason. Use the same method for
other context tests.

Fixes #18759

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

diff --git a/src/database/sql/sql_test.go b/src/database/sql/sql_test.go
index 63e1292cb1f2c50591f2e1f63f2606a879de5e95..3f8e03ce135426b241b089e9623a1b408b874ac2 100644 (file)
--- a/src/database/sql/sql_test.go
+++ b/src/database/sql/sql_test.go
@@ -326,9 +326,7 @@ func TestQueryContext(t *testing.T) {
 
        // And verify that the final rows.Next() call, which hit EOF,
        // also closed the rows connection.
-       if n := db.numFreeConns(); n != 1 {
-               t.Fatalf("free conns after query hitting EOF = %d; want 1", n)
-       }
+       waitForFree(t, db, 5*time.Second, 1)
        if prepares := numPrepares(t, db) - prepares0; prepares != 1 {
                t.Errorf("executed %d Prepare statements; want 1", prepares)
        }
@@ -345,6 +343,18 @@ func waitCondition(waitFor, checkEvery time.Duration, fn func() bool) bool {
        return false
 }
 
+// waitForFree checks db.numFreeConns until either it equals want or
+// the maxWait time elapses.
+func waitForFree(t *testing.T, db *DB, maxWait time.Duration, want int) {
+       var numFree int
+       if !waitCondition(maxWait, 5*time.Millisecond, func() bool {
+               numFree = db.numFreeConns()
+               return numFree == want
+       }) {
+               t.Fatalf("free conns after hitting EOF = %d; want %d", numFree, want)
+       }
+}
+
 func TestQueryContextWait(t *testing.T) {
        db := newTestDB(t, "people")
        defer closeDB(t, db)
@@ -361,9 +371,7 @@ func TestQueryContextWait(t *testing.T) {
        }
 
        // Verify closed rows connection after error condition.
-       if n := db.numFreeConns(); n != 1 {
-               t.Fatalf("free conns after query hitting EOF = %d; want 1", n)
-       }
+       waitForFree(t, db, 5*time.Second, 1)
        if prepares := numPrepares(t, db) - prepares0; prepares != 1 {
                t.Errorf("executed %d Prepare statements; want 1", prepares)
        }
@@ -388,13 +396,7 @@ func TestTxContextWait(t *testing.T) {
                t.Fatalf("expected QueryContext to error with context deadline exceeded but returned %v", err)
        }
 
-       var numFree int
-       if !waitCondition(5*time.Second, 5*time.Millisecond, func() bool {
-               numFree = db.numFreeConns()
-               return numFree == 0
-       }) {
-               t.Fatalf("free conns after hitting EOF = %d; want 0", numFree)
-       }
+       waitForFree(t, db, 5*time.Second, 0)
 
        // Ensure the dropped connection allows more connections to be made.
        // Checked on DB Close.
@@ -471,9 +473,7 @@ func TestMultiResultSetQuery(t *testing.T) {
 
        // And verify that the final rows.Next() call, which hit EOF,
        // also closed the rows connection.
-       if n := db.numFreeConns(); n != 1 {
-               t.Fatalf("free conns after query hitting EOF = %d; want 1", n)
-       }
+       waitForFree(t, db, 5*time.Second, 1)
        if prepares := numPrepares(t, db) - prepares0; prepares != 1 {
                t.Errorf("executed %d Prepare statements; want 1", prepares)
        }