Wait for writeStream to open

Author: Brian Chen

packages/firestore/test/unit/specs/spec_test_runner.ts

@@ -174,6 +174,9 @@ class MockConnection implements Connection {
   /** A Deferred that is resolved once watch opens. */
   watchOpen = new Deferred<void>();
 
+  /** A Deferred that is resolved once the write stream opens. */
+  writeOpen = new Deferred<void>();
+
   invokeRPC<Req>(rpcName: string, request: Req): never {
     throw new Error('Not implemented!');
   }
@@ -196,6 +199,10 @@ class MockConnection implements Connection {
     return this.watchOpen.promise;
   }
 
+  waitForWriteOpen(): Promise<void> {
+    return this.writeOpen.promise;
+  }
+
   ackWrite(commitTime?: string, mutationResults?: api.WriteResult[]): void {
     this.writeStream!.callOnMessage({
       // Convert to base64 string so it can later be parsed into ByteString.
@@ -215,6 +222,7 @@ class MockConnection implements Connection {
   private resetAndCloseWriteStream(err?: FirestoreError): void {
     this.writeSendBarriers = [];
     this.earlyWrites = [];
+    this.writeOpen = new Deferred<void>();
     this.writeStream!.callOnClose(err);
     this.writeStream = null;
   }
@@ -280,11 +288,16 @@ class MockConnection implements Connection {
           this.resetAndCloseWriteStream();
         }
       });
-      this.queue.enqueueAndForget(async () => {
-        if (this.writeStream === writeStream) {
-          writeStream.callOnOpen();
-        }
-      });
+      // eslint-disable-next-line @typescript-eslint/no-floating-promises
+      this.queue
+        .enqueue(async () => {
+          if (this.writeStream === writeStream) {
+            writeStream.callOnOpen();
+          }
+        })
+        .then(() => {
+          this.writeOpen.resolve();
+        });
       this.writeStream = writeStream;
       // Replace 'any' with conditional types.
       return writeStream as any; // eslint-disable-line @typescript-eslint/no-explicit-any
@@ -313,12 +326,16 @@ class MockConnection implements Connection {
         }
       });
       // Call on open immediately after returning
-      this.queue.enqueueAndForget(async () => {
-        if (this.watchStream === watchStream) {
-          watchStream.callOnOpen();
+      // eslint-disable-next-line @typescript-eslint/no-floating-promises
+      this.queue
+        .enqueue(async () => {
+          if (this.watchStream === watchStream) {
+            watchStream.callOnOpen();
+          }
+        })
+        .then(() => {
           this.watchOpen.resolve();
-        }
-      });
+        });
       this.watchStream = watchStream;
       // Replace 'any' with conditional types.
       return this.watchStream as any; // eslint-disable-line @typescript-eslint/no-explicit-any
@@ -931,6 +948,9 @@ abstract class TestRunner {
   private async doEnableNetwork(): Promise<void> {
     this.networkEnabled = true;
     await this.syncEngine.enableNetwork();
+    if (this.connection.writeStream !== null) {
+      await this.connection.waitForWriteOpen();
+    }
   }
 
   private async doShutdown(): Promise<void> {