database/sql: make TestTxContextWaitNoDiscard test more robust

Similar to CL 385934, rely on waiter trigger instead of the WAIT query
prefix and factor out the common test code.

Fixes #53222

Change-Id: I46efc85ca102b350bb4dbe8e514921e016870ffb
Reviewed-on: https://go-review.googlesource.com/c/go/+/416655
Reviewed-by: Michael Knyszek <mknyszek@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
Auto-Submit: Bryan Mills <bcmills@google.com>
Run-TryBot: Dmitri Goutnik <dgoutnik@gmail.com>
Reviewed-by: Bryan Mills <bcmills@google.com>

diff --git a/src/database/sql/sql_test.go b/src/database/sql/sql_test.go
index 6bc869fc864c26d4ac32cba3aadadf3cd9f12fea..8c58723c03011b7d2036e5eba82c26f4836eb7e3 100644 (file)
--- a/src/database/sql/sql_test.go
+++ b/src/database/sql/sql_test.go
@@ -449,6 +449,16 @@ func TestQueryContextWait(t *testing.T) {
 // TestTxContextWait tests the transaction behavior when the tx context is canceled
 // during execution of the query.
 func TestTxContextWait(t *testing.T) {
+       testContextWait(t, false)
+}
+
+// TestTxContextWaitNoDiscard is the same as TestTxContextWait, but should not discard
+// the final connection.
+func TestTxContextWaitNoDiscard(t *testing.T) {
+       testContextWait(t, true)
+}
+
+func testContextWait(t *testing.T, keepConnOnRollback bool) {
        db := newTestDB(t, "people")
        defer closeDB(t, db)
 
@@ -458,7 +468,7 @@ func TestTxContextWait(t *testing.T) {
        if err != nil {
                t.Fatal(err)
        }
-       tx.keepConnOnRollback = false
+       tx.keepConnOnRollback = keepConnOnRollback
 
        tx.dc.ci.(*fakeConn).waiter = func(c context.Context) {
                cancel()
@@ -472,36 +482,11 @@ func TestTxContextWait(t *testing.T) {
                t.Fatalf("expected QueryContext to error with context canceled but returned %v", err)
        }
 
-       waitForFree(t, db, 0)
-}
-
-// TestTxContextWaitNoDiscard is the same as TestTxContextWait, but should not discard
-// the final connection.
-func TestTxContextWaitNoDiscard(t *testing.T) {
-       db := newTestDB(t, "people")
-       defer closeDB(t, db)
-
-       ctx, cancel := context.WithTimeout(context.Background(), 15*time.Millisecond)
-       defer cancel()
-
-       tx, err := db.BeginTx(ctx, nil)
-       if err != nil {
-               // Guard against the context being canceled before BeginTx completes.
-               if err == context.DeadlineExceeded {
-                       t.Skip("tx context canceled prior to first use")
-               }
-               t.Fatal(err)
-       }
-
-       // 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|1s|SELECT|people|age,name|")
-       if err != context.DeadlineExceeded {
-               t.Fatalf("expected QueryContext to error with context deadline exceeded but returned %v", err)
+       if keepConnOnRollback {
+               waitForFree(t, db, 1)
+       } else {
+               waitForFree(t, db, 0)
        }
-
-       waitForFree(t, db, 1)
 }
 
 // TestUnsupportedOptions checks that the database fails when a driver that