diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java
index c2ad5ba5a3..b1acfb8229 100644
--- a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java
+++ b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java
@@ -538,7 +538,7 @@ private void waitForDoneCallback(long duration, TimeUnit timeUnit) {
     while (System.nanoTime() <= deadline) {
       this.lock.lock();
       try {
-        if (connectionFinalStatus != null) {
+        if (!this.streamConnectionIsConnected) {
           // Done callback is received, return.
           return;
         }