FCM: async remove database in FIRMessagingRmqManager (#4771)

* Add tests for removeDatabase deadlock

* Cleanup

* Cleanup and comments.

* FIRMessagingRmqManager: async remove database

* FIRMessagingRmqManager: async remove DB test fixes.

* Fix trailing spaces

* Cleanup

* FIR_MESSAGING_ASSERTIONS_BLOCKED handling added. FIR_MESSAGING_ASSERTIONS_BLOCKED is set in google3 tests

* Fix blaze tests

Example/Messaging/Tests/FIRInstanceIDWithFCMTest.m

@@ -54,7 +54,7 @@
 }
 
 - (void)tearDown {
-  [_testUtil cleanupAfterTest];
+  [_testUtil cleanupAfterTest:self];
   _instanceID = nil;
   _messaging = nil;
   [_mockFirebaseApp stopMocking];

Example/Messaging/Tests/FIRMessagingDataMessageManagerTest.m

@@ -16,6 +16,7 @@
 
 #import <OCMock/OCMock.h>
 #import <XCTest/XCTest.h>
+#import "XCTestCase+FIRMessagingRmqManagerTests.h"
 
 #import <FirebaseMessaging/FIRMessaging.h>
 
@@ -90,6 +91,7 @@ static NSString *const kRmqDatabaseName = @"gcm-dmm-test";
 -(void)tearDown {
     if (_dataMessageManager.rmq2Manager) {
         [_dataMessageManager.rmq2Manager removeDatabase];
+        [self waitForDrainDatabaseQueueForRmqManager:_dataMessageManager.rmq2Manager];
     }
     [super tearDown];
 }

Example/Messaging/Tests/FIRMessagingHandlingTest.m

@@ -76,7 +76,7 @@ static NSString *const kFIRMessagingDefaultsTestDomain = @"com.messaging.tests";
 }
 
 - (void)tearDown {
-  [_testUtil cleanupAfterTest];
+  [_testUtil cleanupAfterTest:self];
   [_mockMessagingAnalytics stopMocking];
   [_mockFirebaseApp stopMocking];
   [super tearDown];

Example/Messaging/Tests/FIRMessagingLinkHandlingTest.m

@@ -53,7 +53,7 @@ NSString *const kFIRMessagingTestsLinkHandlingSuiteName = @"com.messaging.test_l
 }
 
 - (void)tearDown {
-  [_testUtil cleanupAfterTest];
+  [_testUtil cleanupAfterTest:self];
   _messaging = nil;
   [[[NSUserDefaults alloc] initWithSuiteName:kFIRMessagingTestsLinkHandlingSuiteName] removePersistentDomainForName:kFIRMessagingTestsLinkHandlingSuiteName];
   [super tearDown];

Example/Messaging/Tests/FIRMessagingRmqManagerTest.m

@@ -16,6 +16,11 @@
 
 #import <XCTest/XCTest.h>
 
+#import <OCMock/OCMock.h>
+
+#import "FIRTestsAssertionHandler.h"
+#import "XCTestCase+FIRMessagingRmqManagerTests.h"
+
 #import "Firebase/Messaging/FIRMessagingPersistentSyncMessage.h"
 #import "Firebase/Messaging/FIRMessagingRmqManager.h"
 #import "Firebase/Messaging/FIRMessagingUtilities.h"
@@ -28,12 +33,15 @@ static const NSTimeInterval kAsyncTestTimout = 0.5;
 @interface FIRMessagingRmqManager (ExposedForTest)
 
 - (void)removeDatabase;
+- (dispatch_queue_t)databaseOperationQueue;
 
 @end
 
 @interface FIRMessagingRmqManagerTest : XCTestCase
 
 @property(nonatomic, readwrite, strong) FIRMessagingRmqManager *rmqManager;
+@property(nonatomic, strong) id assertionHandlerMock;
+@property(nonatomic, strong) FIRTestsAssertionHandler *testAssertionHandler;
 
 @end
 
@@ -41,13 +49,23 @@ static const NSTimeInterval kAsyncTestTimout = 0.5;
 
 - (void)setUp {
   [super setUp];
+
+  self.testAssertionHandler = [[FIRTestsAssertionHandler alloc] init];
+  self.assertionHandlerMock = OCMClassMock([NSAssertionHandler class]);
+  OCMStub([self.assertionHandlerMock currentHandler]).andReturn(self.testAssertionHandler);
+
   // Make sure we start off with a clean state each time
   _rmqManager = [[FIRMessagingRmqManager alloc] initWithDatabaseName:kRmqDatabaseName];
-
 }
 
 - (void)tearDown {
   [self.rmqManager removeDatabase];
+  [self waitForDrainDatabaseQueueForRmqManager:self.rmqManager];
+
+  [self.assertionHandlerMock stopMocking];
+  self.assertionHandlerMock = nil;
+  self.testAssertionHandler = nil;
+
   [super tearDown];
 }
 
@@ -302,6 +320,40 @@ static const NSTimeInterval kAsyncTestTimout = 0.5;
   XCTAssertNil([self.rmqManager querySyncMessageWithRmqID:rmqID]);
 }
 
+- (void)testInitWhenDatabaseIsBrokenThenDatabaseIsDeleted {
+  NSString *databaseName = @"invalid-database-file";
+  NSString *databasePath = [self createBrokenDatabaseWithName:databaseName];
+  XCTAssert([[NSFileManager defaultManager] fileExistsAtPath:databasePath]);
+
+  // Expect for at least one assertion.
+  XCTestExpectation *assertionFailureExpectation = [self expectationWithDescription:@"assertionFailureExpectation"];
+  assertionFailureExpectation.assertForOverFulfill = NO;
+
+// The flag FIR_MESSAGING_ASSERTIONS_BLOCKED can be set by blaze when running tests from google3.
+#ifndef FIR_MESSAGING_ASSERTIONS_BLOCKED
+  [self.testAssertionHandler setMethodFailureHandlerForClass:[FIRMessagingRmqManager class]
+                                                     handler:^(id object, NSString *fileName, NSInteger lineNumber) {
+    [assertionFailureExpectation fulfill];
+  }];
+#else
+  // If FIR_MESSAGING_ASSERTIONS_BLOCKED is defined, then no assertion handlers will be called,
+  // so don't wait for it.
+  dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(0.5 * NSEC_PER_SEC)), dispatch_get_main_queue(), ^{
+    [assertionFailureExpectation fulfill];
+  });
+#endif // FIR_MESSAGING_ASSERTIONS_BLOCKED
+
+  // Create `FIRMessagingRmqManager` instance with a broken database.
+  FIRMessagingRmqManager *manager = [[FIRMessagingRmqManager alloc] initWithDatabaseName:databaseName];
+
+  [self waitForExpectations:@[ assertionFailureExpectation ] timeout:0.5];
+
+  [self waitForDrainDatabaseQueueForRmqManager:manager];
+
+  // Check that the file was deleted.
+  XCTAssertFalse([[NSFileManager defaultManager] fileExistsAtPath:databasePath]);
+}
+
 #pragma mark - Private Helpers
 
 - (GtalkDataMessageStanza *)dataMessageWithMessageID:(NSString *)messageID
@@ -323,4 +375,24 @@ static const NSTimeInterval kAsyncTestTimout = 0.5;
   return stanza;
 }
 
+- (NSString *)createBrokenDatabaseWithName:(NSString *)name {
+  NSString *databasePath = [FIRMessagingRmqManager pathForDatabaseWithName:name];
+  NSMutableArray *pathComponents = [[databasePath pathComponents] mutableCopy];
+  [pathComponents removeLastObject];
+  NSString *directoryPath = [NSString pathWithComponents:pathComponents];
+
+  // Create directory if doesn't exist.
+  [[NSFileManager defaultManager] createDirectoryAtPath:directoryPath withIntermediateDirectories:YES attributes:nil error:NULL];
+  // Remove the file if exists.
+  [[NSFileManager defaultManager] removeItemAtPath:databasePath error:NULL];
+
+  NSData *brokenDBFileContent = [@"not a valid DB" dataUsingEncoding:NSUTF8StringEncoding];
+  [brokenDBFileContent writeToFile:databasePath atomically:YES];
+
+  XCTAssertEqualObjects([NSData dataWithContentsOfFile:databasePath], brokenDBFileContent);
+  return databasePath;
+}
+
+
+
 @end

Example/Messaging/Tests/FIRMessagingServiceTest.m

@@ -85,7 +85,7 @@ static NSString *const kFIRMessagingTestsServiceSuiteName = @"com.messaging.test
 }
 
 - (void)tearDown {
-  [_testUtil cleanupAfterTest];
+  [_testUtil cleanupAfterTest:self];
   _messaging = nil;
   [_mockPubSub stopMocking];
   [[[NSUserDefaults alloc] initWithSuiteName:kFIRMessagingTestsServiceSuiteName] removePersistentDomainForName:kFIRMessagingTestsServiceSuiteName];

Example/Messaging/Tests/FIRMessagingSyncMessageManagerTest.m

@@ -16,6 +16,8 @@
 
 #import <XCTest/XCTest.h>
 
+#import "XCTestCase+FIRMessagingRmqManagerTests.h"
+
 #import "Firebase/Messaging/FIRMessagingPersistentSyncMessage.h"
 #import "Firebase/Messaging/FIRMessagingRmqManager.h"
 #import "Firebase/Messaging/FIRMessagingSyncMessageManager.h"
@@ -48,6 +50,7 @@ static NSString *const kRmqSqliteFilename = @"rmq-sync-manager-test";
 
 - (void)tearDown {
   [_rmqManager removeDatabase];
+  [self waitForDrainDatabaseQueueForRmqManager:_rmqManager];
   [super tearDown];
 }
 

Example/Messaging/Tests/FIRMessagingTest.m

@@ -82,7 +82,7 @@ static NSString *const kFIRMessagingDefaultsTestDomain = @"com.messaging.tests";
 }
 
 - (void)tearDown {
-  [_testUtil cleanupAfterTest];
+  [_testUtil cleanupAfterTest:self];
   [_mockFirebaseApp stopMocking];
   _messaging = nil;
   [[[NSUserDefaults alloc] initWithSuiteName:kFIRMessagingDefaultsTestDomain]

Example/Messaging/Tests/FIRMessagingTestUtilities.h

@@ -41,7 +41,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (instancetype)initWithUserDefaults:(NSUserDefaults *)userDefaults withRMQManager:(BOOL)withRMQManager;
 
-- (void)cleanupAfterTest;
+- (void)cleanupAfterTest:(XCTestCase *)testCase;
 
 @end
 

Example/Messaging/Tests/FIRMessagingTestUtilities.m

@@ -16,6 +16,8 @@
 
 #import <OCMock/OCMock.h>
 
+#import "XCTestCase+FIRMessagingRmqManagerTests.h"
+
 #import "Example/Messaging/Tests/FIRMessagingTestUtilities.h"
 
 #import <FirebaseAnalyticsInterop/FIRAnalyticsInterop.h>
@@ -95,8 +97,9 @@ static NSString *const kFIRMessagingDefaultsTestDomain = @"com.messaging.tests";
   return self;
 }
 
-- (void)cleanupAfterTest {
+- (void)cleanupAfterTest:(XCTestCase *)testCase {
   [_messaging.rmq2Manager removeDatabase];
+  [testCase waitForDrainDatabaseQueueForRmqManager:_messaging.rmq2Manager];
   [_messaging.messagingUserDefaults removePersistentDomainForName:kFIRMessagingDefaultsTestDomain];
   _messaging.shouldEstablishDirectChannel = NO;
   [_mockPubsub stopMocking];

Example/Messaging/Tests/FIRTestsAssertionHandler.h

@@ -0,0 +1,37 @@
+/*
+* Copyright 2020 Google
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+*      http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+
+#import <Foundation/Foundation.h>
+
+NS_ASSUME_NONNULL_BEGIN
+
+typedef void(^FIRTestsAssertionHandlerBlock)(id object, NSString *fileName, NSInteger lineNumber);
+
+@interface FIRTestsAssertionHandler : NSAssertionHandler
+
+/**
+ * Sets a handler for assertions in objects of the specified class.
+ * @param aClass A class to set an assertion method failure handler.
+ * @param handler A custom handler that will be called each time when
+ * `-[FIRTestsAssertionHandler handleFailureInMethod:object:file:lineNumber:description:]` is called. If `nil` then
+ * `-[super handleFailureInMethod:object:file:lineNumber:description:]` (default implementation) will be called.
+ */
+- (void)setMethodFailureHandlerForClass:(Class)aClass
+                                handler:(nullable FIRTestsAssertionHandlerBlock)handler;
+
+@end
+
+NS_ASSUME_NONNULL_END

Example/Messaging/Tests/FIRTestsAssertionHandler.m

@@ -0,0 +1,59 @@
+/*
+* Copyright 2020 Google
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+*      http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+
+#import "FIRTestsAssertionHandler.h"
+
+@interface FIRTestsAssertionHandler ()
+
+@property(nonatomic) NSMutableDictionary<NSString *, FIRTestsAssertionHandlerBlock> *methodFailureHandlers;
+
+@end
+
+@implementation FIRTestsAssertionHandler
+
+- (instancetype)init
+{
+  self = [super init];
+  if (self) {
+    _methodFailureHandlers = [NSMutableDictionary dictionary];
+  }
+  return self;
+}
+
+- (void)handleFailureInMethod:(SEL)selector
+                       object:(id<NSObject>)object
+                         file:(NSString *)fileName
+                   lineNumber:(NSInteger)line
+                  description:(NSString *)format, ... {
+  FIRTestsAssertionHandlerBlock handler;
+  @synchronized (self) {
+    handler = self.methodFailureHandlers[NSStringFromClass([object class])];
+  }
+
+  if (handler) {
+    handler(object, fileName, line);
+  } else {
+    [super handleFailureInMethod:selector object:object file:fileName lineNumber:line description:@"%@", format];
+  }
+}
+
+- (void)setMethodFailureHandlerForClass:(Class)aClass handler:(FIRTestsAssertionHandlerBlock)handler {
+  @synchronized (self) {
+    self.methodFailureHandlers[NSStringFromClass(aClass)]  = [handler copy];
+  }
+}
+
+@end

Example/Messaging/Tests/XCTestCase+FIRMessagingRmqManagerTests.h

@@ -0,0 +1,29 @@
+/*
+* Copyright 2020 Google
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+*      http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+
+#import <XCTest/XCTest.h>
+
+@class FIRMessagingRmqManager;
+
+NS_ASSUME_NONNULL_BEGIN
+
+@interface XCTestCase (FIRMessagingRmqManagerTests)
+
+- (void)waitForDrainDatabaseQueueForRmqManager:(FIRMessagingRmqManager *)manager;
+
+@end
+
+NS_ASSUME_NONNULL_END

Example/Messaging/Tests/XCTestCase+FIRMessagingRmqManagerTests.m

@@ -0,0 +1,42 @@
+/*
+* Copyright 2020 Google
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+*      http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+
+#import "XCTestCase+FIRMessagingRmqManagerTests.h"
+
+#import <Foundation/Foundation.h>
+
+#import "Firebase/Messaging/FIRMessagingRmqManager.h"
+
+@interface FIRMessagingRmqManager (FIRMessagingRmqManagerTests)
+- (dispatch_queue_t)databaseOperationQueue;
+@end
+
+@implementation XCTestCase (FIRMessagingRmqManagerTests)
+
+- (void)waitForDrainDatabaseQueueForRmqManager:(FIRMessagingRmqManager *)manager {
+  dispatch_queue_t databaseQueue = [manager databaseOperationQueue];
+  if (databaseQueue == nil) {
+    return;
+  }
+
+  XCTestExpectation *drainDatabaseQueueExpectation = [self expectationWithDescription:@"drainDatabaseQueue"];
+  dispatch_async([manager databaseOperationQueue], ^{
+    [drainDatabaseQueueExpectation fulfill];
+  });
+  [self waitForExpectations:@[drainDatabaseQueueExpectation] timeout:1.5];
+}
+
+@end

Firebase/Messaging/FIRMessagingRmqManager.h

@@ -154,4 +154,11 @@ typedef void(^FIRMessagingRmqMessageHandler)(NSDictionary<NSString *, GPBMessage
  */
 - (void)updateSyncMessageViaMCSWithRmqID:(NSString *)rmqID;
 
+/**
+ * Returns path for database with specified name.
+ * @param databaseName The database name without extension: "<databaseName>.sqlite".
+ * @returns Path to the database with the specified name.
+ */
++ (NSString *)pathForDatabaseWithName:(NSString *)databaseName;
+
 @end

Firebase/Messaging/FIRMessagingRmqManager.m

@@ -544,7 +544,11 @@ NSString * _Nonnull FIRMessagingStringFromSQLiteResult(int result) {
 # pragma mark - Database
 
 - (NSString *)pathForDatabase {
-  NSString *dbNameWithExtension = [NSString stringWithFormat:@"%@.sqlite", _databaseName];
+  return [[self class] pathForDatabaseWithName:_databaseName];
+}
+
++ (NSString *)pathForDatabaseWithName:(NSString *)databaseName {
+  NSString *dbNameWithExtension = [NSString stringWithFormat:@"%@.sqlite", databaseName];
   NSArray *paths = NSSearchPathForDirectoriesInDomains(FIRMessagingSupportedDirectory(),
                                                   NSUserDomainMask,
                                                   YES);
@@ -585,7 +589,7 @@ NSString * _Nonnull FIRMessagingStringFromSQLiteResult(int result) {
 
 - (void)removeDatabase {
   // Ensure database is removed in a sync queue as this sometimes makes test have race conditions.
-  dispatch_sync(_databaseOperationQueue, ^{
+  dispatch_async(_databaseOperationQueue, ^{
     NSString *path = [self pathForDatabase];
     [[NSFileManager defaultManager] removeItemAtPath:path error:nil];
   });
@@ -853,4 +857,9 @@ NSString * _Nonnull FIRMessagingStringFromSQLiteResult(int result) {
   [self logError];
   sqlite3_finalize(stmt);
 }
+
+- (dispatch_queue_t)databaseOperationQueue {
+  return _databaseOperationQueue;
+}
+
 @end