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>