1 files changed, 9 insertions, 3 deletions
diff --git a/Firestore/Example/Tests/Integration/FSTStreamTests.mm b/Firestore/Example/Tests/Integration/FSTStreamTests.mm
index a36361a..cb967b1 100644
--- a/Firestore/Example/Tests/Integration/FSTStreamTests.mm
+++ b/Firestore/Example/Tests/Integration/FSTStreamTests.mm
@@ -16,6 +16,8 @@
 
 #import <XCTest/XCTest.h>
 
+#import <GRPCClient/GRPCCall.h>
+
 #import <FirebaseFirestore/FIRFirestoreSettings.h>
 
 #import "Firestore/Example/Tests/Util/FSTHelpers.h"
@@ -35,7 +37,7 @@ using firebase::firestore::model::DatabaseId;
 
 /** Exposes otherwise private methods for testing. */
 @interface FSTStream (Testing)
-- (void)writesFinishedWithError:(NSError *_Nullable)error;
+@property(nonatomic, strong, readwrite) id<GRXWriteable> callbackFilter;
 @end
 
 /**
@@ -202,7 +204,9 @@ using firebase::firestore::model::DatabaseId;
   }];
 
   // Simulate a final callback from GRPC
-  [watchStream writesFinishedWithError:nil];
+  [_workerDispatchQueue dispatchAsync:^{
+    [watchStream.callbackFilter writesFinishedWithError:nil];
+  }];
 
   [self verifyDelegateObservedStates:@[ @"watchStreamDidOpen" ]];
 }
@@ -224,7 +228,9 @@ using firebase::firestore::model::DatabaseId;
   }];
 
   // Simulate a final callback from GRPC
-  [writeStream writesFinishedWithError:nil];
+  [_workerDispatchQueue dispatchAsync:^{
+    [writeStream.callbackFilter writesFinishedWithError:nil];
+  }];
 
   [self verifyDelegateObservedStates:@[ @"writeStreamDidOpen" ]];
 }