Cleaning up test

Author: schmidt-sebastian

Firestore/Example/Tests/Integration/FSTStreamTests.m

@@ -38,17 +38,25 @@ - (void)writesFinishedWithError:(NSError *_Nullable)error;
  */
 @interface FSTStreamStatusDelegate : NSObject <FSTWatchStreamDelegate, FSTWriteStreamDelegate>
 
+- (instancetype)initFrom:(XCTestCase *)testCase
+              usingQueue:(FSTDispatchQueue *)dispatchQueue NS_DESIGNATED_INITIALIZER;
+- (instancetype)init NS_UNAVAILABLE;
+
 @property(nonatomic, readonly) NSMutableArray<NSString *> *states;
 @property(atomic, readwrite) BOOL invokeCallbacks;
 @property(nonatomic, weak) XCTestExpectation *expectation;
 @property(nonatomic, weak, readonly) FSTStream *stream;
+@property(nonatomic, weak, readonly) XCTestCase *testCase;
+@property(nonatomic, weak, readonly) FSTDispatchQueue *dispatchQueue;
 
 @end
 
 @implementation FSTStreamStatusDelegate
 
-- (instancetype)init {
+- (instancetype)initFrom:(XCTestCase *)testCase usingQueue:(FSTDispatchQueue *)dispatchQueue {
   if (self = [super init]) {
+    _testCase = testCase;
+    _dispatchQueue = dispatchQueue;
     _states = [NSMutableArray new];
   }
 
@@ -99,9 +107,15 @@ - (void)writeStreamDidReceiveResponseWithVersion:(FSTSnapshotVersion *)commitVer
   _expectation = nil;
 }
 
-- (void)fulfillOnCallback:(XCTestExpectation *)expectation {
+/** Executes 'block' using the provided FSTDispatchQueue and waits for any callback on this delegate
+ * to be called. */
+- (void)awaitNotificationFromBlock:(void (^)(void))block {
   FSTAssert(_expectation == nil, @"Previous expectation still active");
+  XCTestExpectation *expectation =
+      [self.testCase expectationWithDescription:@"awaitCallbackInBlock"];
   _expectation = expectation;
+  [self.dispatchQueue dispatchAsync:block];
+  [self.testCase awaitExpectations];
 }
 
 @end
@@ -140,16 +154,12 @@ - (void)setUp {
   _credentials = [[FSTEmptyCredentialsProvider alloc] init];
 }
 
-- (void)tearDown {
-  [super tearDown];
-}
-
 - (FSTWriteStream *)setUpWriteStream {
   FSTDatastore *datastore = [[FSTDatastore alloc] initWithDatabaseInfo:_databaseInfo
                                                    workerDispatchQueue:_workerDispatchQueue
                                                            credentials:_credentials];
 
-  _delegate = [FSTStreamStatusDelegate new];
+  _delegate = [[FSTStreamStatusDelegate alloc] initFrom:self usingQueue:_workerDispatchQueue];
   return [datastore createWriteStreamWithDelegate:_delegate];
 }
 
@@ -158,28 +168,30 @@ - (FSTWatchStream *)setUpWatchStream {
                                                    workerDispatchQueue:_workerDispatchQueue
                                                            credentials:_credentials];
 
-  _delegate = [FSTStreamStatusDelegate new];
+  _delegate = [[FSTStreamStatusDelegate alloc] initFrom:self usingQueue:_workerDispatchQueue];
   return [datastore createWatchStreamWithDelegate:_delegate];
 }
 
-- (void)verifyDelegate:(NSArray<NSString *> *)expectedStates {
+/**
+ * Drains the test queue and asserts that all the observed callbacks (up to this point) match
+ * 'expectedStates'. Clears the list of observed callbacks on completion.
+ */
+- (void)verifyDelegateObservedStates:(NSArray<NSString *> *)expectedStates {
   // Drain queue
   dispatch_sync(_testQueue, ^{
                 });
 
   XCTAssertEqualObjects(_delegate.states, expectedStates);
+  [_delegate.states removeAllObjects];
 }
 
 /** Verifies that the watch stream does not issue an onClose callback after a call to stop(). */
 - (void)testWatchStreamStopBeforeHandshake {
   FSTWatchStream *watchStream = [self setUpWatchStream];
 
-  XCTestExpectation *openExpectation = [self expectationWithDescription:@"open"];
-  [_delegate fulfillOnCallback:openExpectation];
-  [_workerDispatchQueue dispatchAsync:^{
+  [_delegate awaitNotificationFromBlock:^() {
     [watchStream start];
   }];
-  [self awaitExpectations];
 
   // Stop must not call watchStreamDidClose because the full implementation of the delegate could
   // attempt to restart the stream in the event it had pending watches.
@@ -190,19 +202,16 @@ - (void)testWatchStreamStopBeforeHandshake {
   // Simulate a final callback from GRPC
   [watchStream writesFinishedWithError:nil];
 
-  [self verifyDelegate:@[ @"watchStreamDidOpen" ]];
+  [self verifyDelegateObservedStates:@[ @"watchStreamDidOpen" ]];
 }
 
 /** Verifies that the write stream does not issue an onClose callback after a call to stop(). */
 - (void)testWriteStreamStopBeforeHandshake {
   FSTWriteStream *writeStream = [self setUpWriteStream];
 
-  XCTestExpectation *openExpectation = [self expectationWithDescription:@"open"];
-  [_delegate fulfillOnCallback:openExpectation];
-  [_workerDispatchQueue dispatchAsync:^{
+  [_delegate awaitNotificationFromBlock:^() {
     [writeStream start];
   }];
-  [self awaitExpectations];
 
   // Don't start the handshake.
 
@@ -215,44 +224,35 @@ - (void)testWriteStreamStopBeforeHandshake {
   // Simulate a final callback from GRPC
   [writeStream writesFinishedWithError:nil];
 
-  [self verifyDelegate:@[ @"writeStreamDidOpen" ]];
+  [self verifyDelegateObservedStates:@[ @"writeStreamDidOpen" ]];
 }
 
 - (void)testWriteStreamStopAfterHandshake {
   FSTWriteStream *writeStream = [self setUpWriteStream];
 
-  XCTestExpectation *openExpectation = [self expectationWithDescription:@"open"];
-  [_delegate fulfillOnCallback:openExpectation];
-  [_workerDispatchQueue dispatchAsync:^{
+  [_delegate awaitNotificationFromBlock:^() {
     [writeStream start];
   }];
-  [self awaitExpectations];
 
   // Writing before the handshake should throw
   dispatch_sync(_testQueue, ^{
     XCTAssertThrows([writeStream writeMutations:_mutations]);
   });
 
-  XCTestExpectation *handshakeExpectation = [self expectationWithDescription:@"handshake"];
-  [_delegate fulfillOnCallback:handshakeExpectation];
-  [_workerDispatchQueue dispatchAsync:^{
+  [_delegate awaitNotificationFromBlock:^() {
     [writeStream writeHandshake];
   }];
-  [self awaitExpectations];
 
   // Now writes should succeed
-  XCTestExpectation *writeExpectation = [self expectationWithDescription:@"write"];
-  [_delegate fulfillOnCallback:writeExpectation];
-  [_workerDispatchQueue dispatchAsync:^{
+  [_delegate awaitNotificationFromBlock:^() {
     [writeStream writeMutations:_mutations];
   }];
-  [self awaitExpectations];
 
   [_workerDispatchQueue dispatchAsync:^{
     [writeStream stop];
   }];
 
-  [self verifyDelegate:@[
+  [self verifyDelegateObservedStates:@[
     @"writeStreamDidOpen", @"writeStreamDidCompleteHandshake",
     @"writeStreamDidReceiveResponseWithVersion"
   ]];