From 3b2c5bfc7b9df5a4d7196825ea9eaca7c62655bd Mon Sep 17 00:00:00 2001
From: Matthew Chen <matthew@signal.org>
Date: Wed, 19 Sep 2018 16:29:15 -0400
Subject: [PATCH] Modify mock environments to register all db views.

---
 SignalMessaging/environment/AppSetup.m        | 16 +++++++++-------
 .../src/Storage/OWSPrimaryStorage.m           | 19 ++++++++++---------
 .../src/Storage/OWSStorage+Subclass.h         |  1 +
 SignalServiceKit/src/Storage/OWSStorage.h     |  1 +
 .../src/Tests/MockSSKEnvironment.m            | 19 ++++++++++++-------
 5 files changed, 33 insertions(+), 23 deletions(-)

diff --git a/SignalMessaging/environment/AppSetup.m b/SignalMessaging/environment/AppSetup.m
index c07dd3198..96ed299ea 100644
--- a/SignalMessaging/environment/AppSetup.m
+++ b/SignalMessaging/environment/AppSetup.m
@@ -68,15 +68,17 @@ NS_ASSUME_NONNULL_BEGIN
         [NSKeyedUnarchiver setClass:[OWSDatabaseMigration class] forClassName:[OWSDatabaseMigration collection]];
 
         [OWSStorage registerExtensionsWithMigrationBlock:^() {
-            // Don't start database migrations until storage is ready.
-            [VersionMigrations performUpdateCheckWithCompletion:^() {
-                OWSAssertIsOnMainThread();
+            dispatch_async(dispatch_get_main_queue(), ^{
+                // Don't start database migrations until storage is ready.
+                [VersionMigrations performUpdateCheckWithCompletion:^() {
+                    OWSAssertIsOnMainThread();
 
-                migrationCompletion();
+                    migrationCompletion();
 
-                OWSAssertDebug(backgroundTask);
-                backgroundTask = nil;
-            }];
+                    OWSAssertDebug(backgroundTask);
+                    backgroundTask = nil;
+                }];
+            });
         }];
     });
 }
diff --git a/SignalServiceKit/src/Storage/OWSPrimaryStorage.m b/SignalServiceKit/src/Storage/OWSPrimaryStorage.m
index 3cf404328..b7000f0f7 100644
--- a/SignalServiceKit/src/Storage/OWSPrimaryStorage.m
+++ b/SignalServiceKit/src/Storage/OWSPrimaryStorage.m
@@ -213,17 +213,18 @@ void VerifyRegistrationsForPrimaryStorage(OWSStorage *storage)
     [OWSMediaGalleryFinder asyncRegisterDatabaseExtensionsWithPrimaryStorage:self];
     [TSDatabaseView asyncRegisterLazyRestoreAttachmentsDatabaseView:self];
 
-    [self.database flushExtensionRequestsWithCompletionQueue:nil
-                                             completionBlock:^{
-                                                 OWSAssertIsOnMainThread();
-                                                 OWSAssertDebug(!self.areAsyncRegistrationsComplete);
-                                                 OWSLogVerbose(@"async registrations complete.");
-                                                 self.areAsyncRegistrationsComplete = YES;
+    [self.database
+        flushExtensionRequestsWithCompletionQueue:dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0)
+                                  completionBlock:^{
+                                      OWSAssertDebug(!self.areAsyncRegistrationsComplete);
+                                      OWSLogVerbose(@"async registrations complete.");
 
-                                                 completion();
+                                      self.areAsyncRegistrationsComplete = YES;
 
-                                                 [self verifyDatabaseViews];
-                                             }];
+                                      completion();
+
+                                      [self verifyDatabaseViews];
+                                  }];
 }
 
 - (void)verifyDatabaseViews
diff --git a/SignalServiceKit/src/Storage/OWSStorage+Subclass.h b/SignalServiceKit/src/Storage/OWSStorage+Subclass.h
index 4a4232c29..5808c2863 100644
--- a/SignalServiceKit/src/Storage/OWSStorage+Subclass.h
+++ b/SignalServiceKit/src/Storage/OWSStorage+Subclass.h
@@ -15,6 +15,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)loadDatabase;
 
 - (void)runSyncRegistrations;
+// completion will be invoked _off_ the main thread.
 - (void)runAsyncRegistrationsWithCompletion:(void (^_Nonnull)(void))completion;
 
 - (BOOL)areAsyncRegistrationsComplete;
diff --git a/SignalServiceKit/src/Storage/OWSStorage.h b/SignalServiceKit/src/Storage/OWSStorage.h
index ead339ca9..199421c05 100644
--- a/SignalServiceKit/src/Storage/OWSStorage.h
+++ b/SignalServiceKit/src/Storage/OWSStorage.h
@@ -58,6 +58,7 @@ typedef void (^OWSStorageMigrationBlock)(void);
 // This object can be used to filter database notifications.
 @property (nonatomic, readonly, nullable) id dbNotificationObject;
 
+// migrationBlock will be invoked _off_ the main thread.
 + (void)registerExtensionsWithMigrationBlock:(OWSStorageMigrationBlock)migrationBlock;
 
 #ifdef DEBUG
diff --git a/SignalServiceKit/src/Tests/MockSSKEnvironment.m b/SignalServiceKit/src/Tests/MockSSKEnvironment.m
index 93c668c82..5f7fac71b 100644
--- a/SignalServiceKit/src/Tests/MockSSKEnvironment.m
+++ b/SignalServiceKit/src/Tests/MockSSKEnvironment.m
@@ -27,7 +27,9 @@ NS_ASSUME_NONNULL_BEGIN
 
 + (void)activate
 {
-    [self setShared:[self new]];
+    MockSSKEnvironment *instance = [self new];
+    [self setShared:instance];
+    [instance configure];
 }
 
 - (instancetype)init
@@ -51,19 +53,22 @@ NS_ASSUME_NONNULL_BEGIN
     return self;
 }
 
-
 + (OWSPrimaryStorage *)createPrimaryStorageForTests
 {
     OWSPrimaryStorage *primaryStorage = [[OWSPrimaryStorage alloc] initStorage];
     [OWSPrimaryStorage protectFiles];
-
-    // TODO: Should we inject a block to do view registrations?
-    primaryStorage.areAsyncRegistrationsComplete = YES;
-    primaryStorage.areSyncRegistrationsComplete = YES;
-
     return primaryStorage;
 }
 
+- (void)configure
+{
+    __block dispatch_semaphore_t semaphore = dispatch_semaphore_create(0);
+    [OWSStorage registerExtensionsWithMigrationBlock:^() {
+        dispatch_semaphore_signal(semaphore);
+    }];
+    dispatch_semaphore_wait(semaphore, DISPATCH_TIME_FOREVER);
+}
+
 @end
 
 NS_ASSUME_NONNULL_END