1 files changed, 30 insertions, 0 deletions
diff --git a/src/database/sql/sql_test.go b/src/database/sql/sql_test.go
index 41265caed8..00a67c79ee 100644
--- a/src/database/sql/sql_test.go
+++ b/src/database/sql/sql_test.go
@@ -442,6 +442,7 @@ func TestTxContextWait(t *testing.T) {
 		}
 		t.Fatal(err)
 	}
+	tx.keepConnOnRollback = false
 
 	// This will trigger the *fakeConn.Prepare method which will take time
 	// performing the query. The ctxDriverPrepare func will check the context
@@ -454,6 +455,35 @@ func TestTxContextWait(t *testing.T) {
 	waitForFree(t, db, 5*time.Second, 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)
+	}
+
+	waitForFree(t, db, 5*time.Second, 1)
+}
+
 // TestUnsupportedOptions checks that the database fails when a driver that
 // doesn't implement ConnBeginTx is used with non-default options and an
 // un-cancellable context.