2 files changed, 164 insertions, 9 deletions
diff --git a/Firestore/Example/Tests/Integration/FSTStreamTests.mm b/Firestore/Example/Tests/Integration/FSTStreamTests.mm
index 6550368..7e37913 100644
--- a/Firestore/Example/Tests/Integration/FSTStreamTests.mm
+++ b/Firestore/Example/Tests/Integration/FSTStreamTests.mm
@@ -243,9 +243,9 @@ using firebase::firestore::model::DatabaseId;
   }];
 
   // Writing before the handshake should throw
-  dispatch_sync(_testQueue, ^{
+  [_workerDispatchQueue dispatchSync:^{
     XCTAssertThrows([writeStream writeMutations:_mutations]);
-  });
+  }];
 
   [_delegate awaitNotificationFromBlock:^{
     [writeStream writeHandshake];
@@ -285,9 +285,9 @@ using firebase::firestore::model::DatabaseId;
 
   [_workerDispatchQueue runDelayedCallbacksUntil:FSTTimerIDWriteStreamIdle];
 
-  dispatch_sync(_testQueue, ^{
+  [_workerDispatchQueue dispatchSync:^{
     XCTAssertFalse([writeStream isOpen]);
-  });
+  }];
 
   [self verifyDelegateObservedStates:@[
     @"writeStreamDidOpen", @"writeStreamDidCompleteHandshake", @"writeStreamWasInterrupted"
@@ -315,9 +315,9 @@ using firebase::firestore::model::DatabaseId;
         [_workerDispatchQueue containsDelayedCallbackWithTimerID:FSTTimerIDWriteStreamIdle]);
   }];
 
-  dispatch_sync(_testQueue, ^{
+  [_workerDispatchQueue dispatchSync:^{
     XCTAssertTrue([writeStream isOpen]);
-  });
+  }];
 
   [self verifyDelegateObservedStates:@[
     @"writeStreamDidOpen", @"writeStreamDidCompleteHandshake",
diff --git a/Firestore/Example/Tests/Util/FSTDispatchQueueTests.mm b/Firestore/Example/Tests/Util/FSTDispatchQueueTests.mm
index 5ef860c..60b1705 100644
--- a/Firestore/Example/Tests/Util/FSTDispatchQueueTests.mm
+++ b/Firestore/Example/Tests/Util/FSTDispatchQueueTests.mm
@@ -29,6 +29,7 @@ static const FSTTimerID timerID3 = FSTTimerIDWriteStreamConnectionBackoff;
 @end
 
 @implementation FSTDispatchQueueTests {
+  dispatch_queue_t _underlyingQueue;
   FSTDispatchQueue *_queue;
   NSMutableArray *_completedSteps;
   NSArray *_expectedSteps;
@@ -37,13 +38,167 @@ static const FSTTimerID timerID3 = FSTTimerIDWriteStreamConnectionBackoff;
 
 - (void)setUp {
   [super setUp];
-  dispatch_queue_t dispatch_queue =
-      dispatch_queue_create("FSTDispatchQueueTests", DISPATCH_QUEUE_SERIAL);
-  _queue = [[FSTDispatchQueue alloc] initWithQueue:dispatch_queue];
+  _underlyingQueue = dispatch_queue_create("FSTDispatchQueueTests", DISPATCH_QUEUE_SERIAL);
+  _queue = [[FSTDispatchQueue alloc] initWithQueue:_underlyingQueue];
   _completedSteps = [NSMutableArray array];
   _expectedSteps = nil;
 }
 
+- (void)testDispatchAsyncBlocksSubmissionFromTasksOnTheQueue {
+  XCTestExpectation *expectation = [self expectationWithDescription:@"completion"];
+  __block NSException *caught = nil;
+  __block NSString *problem = nil;
+
+  [_queue dispatchAsync:^{
+    @try {
+      [self->_queue dispatchAsync:^{
+      }];
+      problem = @"Should have disallowed submission into the queue while running";
+      [expectation fulfill];
+    } @catch (NSException *ex) {
+      caught = ex;
+      [expectation fulfill];
+    }
+  }];
+
+  [self awaitExpectations];
+  XCTAssertNil(problem);
+  XCTAssertNotNil(caught);
+
+  XCTAssertEqualObjects(caught.name, NSInternalInconsistencyException);
+  XCTAssertTrue(
+      [caught.reason hasPrefix:@"FIRESTORE INTERNAL ASSERTION FAILED: "
+                               @"dispatchAsync called when we are already running on target"]);
+}
+
+- (void)testDispatchAsyncAllowingSameQueueActuallyAllowsSameQueue {
+  XCTestExpectation *expectation = [self expectationWithDescription:@"completion"];
+  __block NSException *caught = nil;
+
+  [_queue dispatchAsync:^{
+    @try {
+      [self->_queue dispatchAsyncAllowingSameQueue:^{
+        [expectation fulfill];
+      }];
+    } @catch (NSException *ex) {
+      caught = ex;
+      [expectation fulfill];
+    }
+  }];
+
+  [self awaitExpectations];
+  XCTAssertNil(caught);
+}
+
+- (void)testDispatchAsyncAllowsSameQueueForUnownedActions {
+  XCTestExpectation *expectation = [self expectationWithDescription:@"completion"];
+  __block NSException *caught = nil;
+
+  // Simulate the case of an action that runs on our queue because e.g. it's run by a user-owned
+  // deinitializer that happened to be last held in one of our API methods.
+  dispatch_async(_underlyingQueue, ^{
+    @try {
+      [self->_queue dispatchAsync:^{
+        [expectation fulfill];
+      }];
+    } @catch (NSException *ex) {
+      caught = ex;
+      [expectation fulfill];
+    }
+  });
+
+  [self awaitExpectations];
+  XCTAssertNil(caught);
+}
+
+- (void)testDispatchSyncBlocksSubmissionFromTasksOnTheQueue {
+  XCTestExpectation *expectation = [self expectationWithDescription:@"completion"];
+  __block NSException *caught = nil;
+  __block NSString *problem = nil;
+
+  [_queue dispatchSync:^{
+    @try {
+      [self->_queue dispatchSync:^{
+      }];
+      problem = @"Should have disallowed submission into the queue while running";
+      [expectation fulfill];
+    } @catch (NSException *ex) {
+      caught = ex;
+      [expectation fulfill];
+    }
+  }];
+
+  [self awaitExpectations];
+  XCTAssertNil(problem);
+  XCTAssertNotNil(caught);
+
+  XCTAssertEqualObjects(caught.name, NSInternalInconsistencyException);
+  XCTAssertTrue(
+      [caught.reason hasPrefix:@"FIRESTORE INTERNAL ASSERTION FAILED: "
+                               @"dispatchSync called when we are already running on target"]);
+}
+
+- (void)testVerifyIsCurrentQueueActuallyRequiresCurrentQueue {
+  XCTAssertNotEqualObjects(_underlyingQueue, dispatch_get_main_queue());
+
+  __block NSException *caught = nil;
+  @try {
+    // Run on the main queue not the FSTDispatchQueue's queue
+    [_queue verifyIsCurrentQueue];
+  } @catch (NSException *ex) {
+    caught = ex;
+  }
+  XCTAssertNotNil(caught);
+  XCTAssertTrue([caught.reason hasPrefix:@"FIRESTORE INTERNAL ASSERTION FAILED: "
+                                         @"We are running on the wrong dispatch queue"]);
+}
+
+- (void)testVerifyIsCurrentQueueRequiresOperationIsInProgress {
+  __block NSException *caught = nil;
+  dispatch_sync(_underlyingQueue, ^{
+    @try {
+      [_queue verifyIsCurrentQueue];
+    } @catch (NSException *ex) {
+      caught = ex;
+    }
+  });
+  XCTAssertNotNil(caught);
+  XCTAssertTrue(
+      [caught.reason hasPrefix:@"FIRESTORE INTERNAL ASSERTION FAILED: "
+                               @"verifyIsCurrentQueue called outside enterCheckedOperation"]);
+}
+
+- (void)testVerifyIsCurrentQueueWorksWithOperationIsInProgress {
+  __block NSException *caught = nil;
+  [_queue dispatchSync:^{
+    @try {
+      [_queue verifyIsCurrentQueue];
+    } @catch (NSException *ex) {
+      caught = ex;
+    }
+  }];
+  XCTAssertNil(caught);
+}
+
+- (void)testEnterCheckedOperationDisallowsNesting {
+  __block NSException *caught = nil;
+  __block NSString *problem = nil;
+  [_queue dispatchSync:^{
+    @try {
+      [_queue enterCheckedOperation:^{
+      }];
+      problem = @"Should not have been able to enter nested enterCheckedOperation";
+    } @catch (NSException *ex) {
+      caught = ex;
+    }
+  }];
+  XCTAssertNil(problem);
+  XCTAssertNotNil(caught);
+  XCTAssertTrue([caught.reason
+      hasPrefix:@"FIRESTORE INTERNAL ASSERTION FAILED: "
+                @"enterCheckedOperation may not be called when an operation is in progress"]);
+}
+
 /**
  * Helper to return a block that adds @(n) to _completedSteps when run and fulfils _expectation if
  * the _completedSteps match the _expectedSteps.