Fix broken tests.

// FREEBIE

Example/TSKitiOSTestApp/Podfile.lock

@@ -111,7 +111,7 @@ EXTERNAL SOURCES:
   AxolotlKit:
     :git: https://github.com/WhisperSystems/SignalProtocolKit.git
   SignalServiceKit:
-    :path: "../../SignalServiceKit.podspec"
+    :path: ../../SignalServiceKit.podspec
   SocketRocket:
     :git: https://github.com/facebook/SocketRocket.git
 

src/Messages/Attachments/TSAttachmentStream.h

@@ -46,7 +46,6 @@ NS_ASSUME_NONNULL_BEGIN
 
 + (void)deleteAttachments;
 + (NSString *)attachmentsFolder;
-+ (NSUInteger)numberOfItemsInAttachmentsFolder;
 
 - (CGSize)imageSizeWithTransaction:(YapDatabaseReadWriteTransaction *)transaction;
 - (CGSize)imageSizeWithoutTransaction;

src/Messages/Attachments/TSAttachmentStream.m

@@ -189,19 +189,6 @@ NS_ASSUME_NONNULL_BEGIN
     return attachmentsFolder;
 }
 
-+ (NSUInteger)numberOfItemsInAttachmentsFolder
-{
-    NSError *error;
-    NSUInteger count =
-        [[[NSFileManager defaultManager] contentsOfDirectoryAtPath:[self attachmentsFolder] error:&error] count];
-
-    if (error) {
-        DDLogError(@"Unable to count attachments in attachments folder. Error: %@", error);
-    }
-
-    return count;
-}
-
 - (nullable NSString *)filePath
 {
     if (!self.localRelativeFilePath) {

src/Storage/OWSOrphanedDataCleaner.h

@@ -2,12 +2,29 @@
 //  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
 //
 
+NS_ASSUME_NONNULL_BEGIN
+
+// Notes:
+//
+// * On disk, we only bother cleaning up files, not directories.
+// * For code simplicity, we don't guarantee that everything is
+//   cleaned up in a single pass. If an interaction is cleaned up,
+//   it's attachments might not be cleaned up until the next pass.
+//   If an attachment is cleaned up, it's file on disk might not
+//   be cleaned up until the next pass.
 @interface OWSOrphanedDataCleaner : NSObject
 
 - (instancetype)init NS_UNAVAILABLE;
 
 + (void)auditAsync;
 
-+ (void)auditAndCleanupAsync;
+// completion, if present, will be invoked on the main thread.
++ (void)auditAndCleanupAsync:(void (^_Nullable)())completion;
+
++ (NSSet<NSString *> *)filePathsInAttachmentsFolder;
+
++ (long long)fileSizeOfFilePaths:(NSArray<NSString *> *)filePaths;
 
 @end
+
+NS_ASSUME_NONNULL_END

src/Storage/OWSOrphanedDataCleaner.m

@@ -9,19 +9,29 @@
 #import "TSStorageManager.h"
 #import "TSThread.h"
 
+NS_ASSUME_NONNULL_BEGIN
+
+#ifdef SSK_BUILDING_FOR_TESTS
+#define CleanupLogDebug NSLog
+#define CleanupLogInfo NSLog
+#else
+#define CleanupLogDebug DDLogDebug
+#define CleanupLogInfo DDLogInfo
+#endif
+
 @implementation OWSOrphanedDataCleaner
 
 + (void)auditAsync
 {
     dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
-        [OWSOrphanedDataCleaner auditAndCleanup:NO];
+        [OWSOrphanedDataCleaner auditAndCleanup:NO completion:nil];
     });
 }
 
-+ (void)auditAndCleanupAsync
++ (void)auditAndCleanupAsync:(void (^_Nullable)())completion
 {
     dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
-        [OWSOrphanedDataCleaner auditAndCleanup:YES];
+        [OWSOrphanedDataCleaner auditAndCleanup:YES completion:completion];
     });
 }
 
@@ -36,45 +46,11 @@
 //   They can't be cleaned up - we don't want to delete the TSAttachmentStream or
 //   its corresponding message.  Better that the broken message shows up in the
 //   conversation view.
-+ (void)auditAndCleanup:(BOOL)shouldCleanup
++ (void)auditAndCleanup:(BOOL)shouldCleanup completion:(void (^_Nullable)())completion
 {
-    NSString *attachmentsFolder = [TSAttachmentStream attachmentsFolder];
+    NSSet<NSString *> *diskFilePaths = [self filePathsInAttachmentsFolder];
-    DDLogDebug(@"attachmentsFolder: %@", attachmentsFolder);
+    long long totalFileSize = [self fileSizeOfFilePaths:diskFilePaths.allObjects];
-
+    NSUInteger fileCount = diskFilePaths.count;
-    __block int fileCount = 0;
-    __block long long totalFileSize = 0;
-    NSMutableSet *diskFilePaths = [NSMutableSet new];
-    __unsafe_unretained __block void (^visitAttachmentFilesRecursable)(NSString *);
-    void (^visitAttachmentFiles)(NSString *);
-    visitAttachmentFiles = ^(NSString *dirPath) {
-        NSError *error;
-        NSArray<NSString *> *fileNames =
-            [[NSFileManager defaultManager] contentsOfDirectoryAtPath:dirPath error:&error];
-        if (error) {
-            OWSFail(@"contentsOfDirectoryAtPath error: %@", error);
-            return;
-        }
-        for (NSString *fileName in fileNames) {
-            NSString *filePath = [dirPath stringByAppendingPathComponent:fileName];
-            BOOL isDirectory;
-            [[NSFileManager defaultManager] fileExistsAtPath:filePath isDirectory:&isDirectory];
-            if (isDirectory) {
-                visitAttachmentFilesRecursable(filePath);
-            } else {
-                NSNumber *fileSize =
-                    [[NSFileManager defaultManager] attributesOfItemAtPath:filePath error:&error][NSFileSize];
-                if (error) {
-                    OWSFail(@"attributesOfItemAtPath: %@ error: %@", filePath, error);
-                    continue;
-                }
-                totalFileSize += fileSize.longLongValue;
-                fileCount++;
-                [diskFilePaths addObject:filePath];
-            }
-        }
-    };
-    visitAttachmentFilesRecursable = visitAttachmentFiles;
-    visitAttachmentFiles(attachmentsFolder);
 
     TSStorageManager *storageManager = [TSStorageManager sharedManager];
     YapDatabaseConnection *databaseConnection = storageManager.newDatabaseConnection;
@@ -98,18 +74,18 @@
                                               }];
     }];
 
-    DDLogDebug(@"fileCount: %d", fileCount);
+    CleanupLogDebug(@"fileCount: %zd", fileCount);
-    DDLogDebug(@"totalFileSize: %lld", totalFileSize);
+    CleanupLogDebug(@"totalFileSize: %lld", totalFileSize);
-    DDLogDebug(@"attachmentStreams: %d", attachmentStreamCount);
+    CleanupLogDebug(@"attachmentStreams: %d", attachmentStreamCount);
-    DDLogDebug(@"attachmentStreams with file paths: %zd", attachmentFilePaths.count);
+    CleanupLogDebug(@"attachmentStreams with file paths: %zd", attachmentFilePaths.count);
 
     NSMutableSet<NSString *> *orphanDiskFilePaths = [diskFilePaths mutableCopy];
     [orphanDiskFilePaths minusSet:attachmentFilePaths];
     NSMutableSet<NSString *> *missingAttachmentFilePaths = [attachmentFilePaths mutableCopy];
     [missingAttachmentFilePaths minusSet:diskFilePaths];
 
-    DDLogDebug(@"orphan disk file paths: %zd", orphanDiskFilePaths.count);
+    CleanupLogDebug(@"orphan disk file paths: %zd", orphanDiskFilePaths.count);
-    DDLogDebug(@"missing attachment file paths: %zd", missingAttachmentFilePaths.count);
+    CleanupLogDebug(@"missing attachment file paths: %zd", missingAttachmentFilePaths.count);
 
     [self printPaths:orphanDiskFilePaths.allObjects label:@"orphan disk file paths"];
     [self printPaths:missingAttachmentFilePaths.allObjects label:@"missing attachment file paths"];
@@ -141,21 +117,25 @@
                                               }];
     }];
 
-    DDLogDebug(@"attachmentIds: %zd", attachmentIds.count);
+    CleanupLogDebug(@"attachmentIds: %zd", attachmentIds.count);
-    DDLogDebug(@"messageAttachmentIds: %zd", messageAttachmentIds.count);
+    CleanupLogDebug(@"messageAttachmentIds: %zd", messageAttachmentIds.count);
 
     NSMutableSet<NSString *> *orphanAttachmentIds = [attachmentIds mutableCopy];
     [orphanAttachmentIds minusSet:messageAttachmentIds];
     NSMutableSet<NSString *> *missingAttachmentIds = [messageAttachmentIds mutableCopy];
     [missingAttachmentIds minusSet:attachmentIds];
 
-    DDLogDebug(@"orphan attachmentIds: %zd", orphanAttachmentIds.count);
+    CleanupLogDebug(@"orphan attachmentIds: %zd", orphanAttachmentIds.count);
-    DDLogDebug(@"missing attachmentIds: %zd", missingAttachmentIds.count);
+    CleanupLogDebug(@"missing attachmentIds: %zd", missingAttachmentIds.count);
-    DDLogDebug(@"orphan interactions: %zd", orphanInteractionIds.count);
+    CleanupLogDebug(@"orphan interactions: %zd", orphanInteractionIds.count);
 
     // We need to avoid cleaning up new attachments and files that are still in the process of
     // being created/written, so we don't clean up anything recent.
+#ifdef SSK_BUILDING_FOR_TESTS
+    const NSTimeInterval kMinimumOrphanAge = 0.f;
+#else
     const NSTimeInterval kMinimumOrphanAge = 15 * 60.f;
+#endif
 
     if (!shouldCleanup) {
         return;
@@ -169,7 +149,7 @@
                 OWSFail(@"Could not load interaction: %@", interactionId);
                 continue;
             }
-            DDLogInfo(@"Removing orphan message: %@", interaction.uniqueId);
+            CleanupLogInfo(@"Removing orphan message: %@", interaction.uniqueId);
             [interaction removeWithTransaction:transaction];
         }
         for (NSString *attachmentId in orphanAttachmentIds) {
@@ -185,11 +165,11 @@
             TSAttachmentStream *attachmentStream = (TSAttachmentStream *)attachment;
             // Don't delete attachments which were created in the last N minutes.
             if (fabs([attachmentStream.creationTimestamp timeIntervalSinceNow]) < kMinimumOrphanAge) {
-                DDLogInfo(@"Skipping orphan attachment due to age: %f",
+                CleanupLogInfo(@"Skipping orphan attachment due to age: %f",
                     fabs([attachmentStream.creationTimestamp timeIntervalSinceNow]));
                 continue;
             }
-            DDLogInfo(@"Removing orphan attachment: %@", attachmentStream.uniqueId);
+            CleanupLogInfo(@"Removing orphan attachment: %@", attachmentStream.uniqueId);
             [attachmentStream removeWithTransaction:transaction];
         }
     }];
@@ -203,24 +183,82 @@
         }
         // Don't delete files which were created in the last N minutes.
         if (fabs([attributes.fileModificationDate timeIntervalSinceNow]) < kMinimumOrphanAge) {
-            DDLogInfo(@"Skipping orphan attachment file due to age: %f",
+            CleanupLogInfo(@"Skipping orphan attachment file due to age: %f",
                 fabs([attributes.fileModificationDate timeIntervalSinceNow]));
             continue;
         }
 
-        DDLogInfo(@"Removing orphan attachment file: %@", filePath);
+        CleanupLogInfo(@"Removing orphan attachment file: %@", filePath);
         [[NSFileManager defaultManager] removeItemAtPath:filePath error:&error];
         if (error) {
             OWSFail(@"Could not remove orphan file at: %@", filePath);
         }
     }
+
+    if (completion) {
+        dispatch_async(dispatch_get_main_queue(), ^{
+            completion();
+        });
+    }
 }
 
 + (void)printPaths:(NSArray<NSString *> *)paths label:(NSString *)label
 {
     for (NSString *path in [paths sortedArrayUsingSelector:@selector(compare:)]) {
-        DDLogDebug(@"%@: %@", label, path);
+        CleanupLogDebug(@"%@: %@", label, path);
+    }
+}
+
++ (NSSet<NSString *> *)filePathsInAttachmentsFolder
+{
+    NSString *attachmentsFolder = [TSAttachmentStream attachmentsFolder];
+    CleanupLogDebug(@"attachmentsFolder: %@", attachmentsFolder);
+
+    return [self filePathsInDirectory:attachmentsFolder];
+}
+
++ (NSSet<NSString *> *)filePathsInDirectory:(NSString *)dirPath
+{
+    NSMutableSet *filePaths = [NSMutableSet new];
+    NSError *error;
+    NSArray<NSString *> *fileNames = [[NSFileManager defaultManager] contentsOfDirectoryAtPath:dirPath error:&error];
+    if (error) {
+        OWSFail(@"contentsOfDirectoryAtPath error: %@", error);
+        return [NSSet new];
     }
+    for (NSString *fileName in fileNames) {
+        NSString *filePath = [dirPath stringByAppendingPathComponent:fileName];
+        BOOL isDirectory;
+        [[NSFileManager defaultManager] fileExistsAtPath:filePath isDirectory:&isDirectory];
+        if (isDirectory) {
+            [filePaths addObjectsFromArray:[self filePathsInDirectory:filePath].allObjects];
+        } else {
+            [filePaths addObject:filePath];
+        }
+    }
+    return filePaths;
+}
+
++ (long long)fileSizeOfFilePath:(NSString *)filePath
+{
+    NSError *error;
+    NSNumber *fileSize = [[NSFileManager defaultManager] attributesOfItemAtPath:filePath error:&error][NSFileSize];
+    if (error) {
+        OWSFail(@"attributesOfItemAtPath: %@ error: %@", filePath, error);
+        return 0;
+    }
+    return fileSize.longLongValue;
+}
+
++ (long long)fileSizeOfFilePaths:(NSArray<NSString *> *)filePaths
+{
+    long long result = 0;
+    for (NSString *filePath in filePaths) {
+        result += [self fileSizeOfFilePath:filePath];
+    }
+    return result;
 }
 
 @end
+
+NS_ASSUME_NONNULL_END

tests/Storage/OWSOrphanedDataCleanerTest.m

@@ -14,6 +14,8 @@
 
 @end
 
+#pragma mark -
+
 @implementation OWSOrphanedDataCleanerTest
 
 - (void)setUp
@@ -24,7 +26,7 @@
 
     // Set up initial conditions & Sanity check
     [TSAttachmentStream deleteAttachments];
-    XCTAssertEqual(0, [TSAttachmentStream numberOfItemsInAttachmentsFolder]);
+    XCTAssertEqual(0, [self numberOfItemsInAttachmentsFolder]);
     [TSAttachmentStream removeAllObjectsInCollection];
     XCTAssertEqual(0, [TSAttachmentStream numberOfKeysInCollection]);
     [TSIncomingMessage removeAllObjectsInCollection];
@@ -38,6 +40,11 @@
     [super tearDown];
 }
 
+- (NSUInteger)numberOfItemsInAttachmentsFolder
+{
+    return [OWSOrphanedDataCleaner filePathsInAttachmentsFolder].count;
+}
+
 - (void)testInteractionsWithoutThreadAreDeleted
 {
     // This thread is intentionally not saved. It's meant to recreate a situation we've seen where interactions exist
@@ -53,7 +60,17 @@
     [incomingMessage save];
     XCTAssertEqual(1, [TSIncomingMessage numberOfKeysInCollection]);
 
-    [[OWSOrphanedDataCleaner new] removeOrphanedData];
+    XCTestExpectation *expectation = [self expectationWithDescription:@"Cleanup"];
+    [OWSOrphanedDataCleaner auditAndCleanupAsync:^{
+        [expectation fulfill];
+    }];
+    [self waitForExpectationsWithTimeout:5.0
+                                 handler:^(NSError *error) {
+                                     if (error) {
+                                         XCTFail(@"Expectation Failed with error: %@", error);
+                                     }
+                                 }];
+
     XCTAssertEqual(0, [TSIncomingMessage numberOfKeysInCollection]);
 }
 
@@ -70,14 +87,24 @@
     [incomingMessage save];
     XCTAssertEqual(1, [TSIncomingMessage numberOfKeysInCollection]);
 
-    [[OWSOrphanedDataCleaner new] removeOrphanedData];
+    XCTestExpectation *expectation = [self expectationWithDescription:@"Cleanup"];
+    [OWSOrphanedDataCleaner auditAndCleanupAsync:^{
+        [expectation fulfill];
+    }];
+    [self waitForExpectationsWithTimeout:5.0
+                                 handler:^(NSError *error) {
+                                     if (error) {
+                                         XCTFail(@"Expectation Failed with error: %@", error);
+                                     }
+                                 }];
+
     XCTAssertEqual(1, [TSIncomingMessage numberOfKeysInCollection]);
 }
 
 - (void)testFilesWithoutInteractionsAreDeleted
 {
     // sanity check
-    XCTAssertEqual(0, [TSAttachmentStream numberOfItemsInAttachmentsFolder]);
+    XCTAssertEqual(0, [self numberOfItemsInAttachmentsFolder]);
 
     NSError *error;
     TSAttachmentStream *attachmentStream = [[TSAttachmentStream alloc] initWithContentType:@"image/jpeg" sourceFilename:nil];
@@ -86,12 +113,25 @@
     NSString *orphanedFilePath = [attachmentStream filePath];
     BOOL fileExists = [[NSFileManager defaultManager] fileExistsAtPath:orphanedFilePath];
     XCTAssert(fileExists);
-    XCTAssertEqual(1, [TSAttachmentStream numberOfItemsInAttachmentsFolder]);
+    XCTAssertEqual(1, [self numberOfItemsInAttachmentsFolder]);
+
+    // Do multiple cleanup passes.
+    for (int i = 0; i < 2; i++) {
+        XCTestExpectation *expectation = [self expectationWithDescription:@"Cleanup"];
+        [OWSOrphanedDataCleaner auditAndCleanupAsync:^{
+            [expectation fulfill];
+        }];
+        [self waitForExpectationsWithTimeout:5.0
+                                     handler:^(NSError *error) {
+                                         if (error) {
+                                             XCTFail(@"Expectation Failed with error: %@", error);
+                                         }
+                                     }];
+    }
 
-    [[OWSOrphanedDataCleaner new] removeOrphanedData];
     fileExists = [[NSFileManager defaultManager] fileExistsAtPath:orphanedFilePath];
     XCTAssertFalse(fileExists);
-    XCTAssertEqual(0, [TSAttachmentStream numberOfItemsInAttachmentsFolder]);
+    XCTAssertEqual(0, [self numberOfItemsInAttachmentsFolder]);
 }
 
 - (void)testFilesWithInteractionsAreNotDeleted
@@ -116,13 +156,22 @@
     NSString *attachmentFilePath = [attachmentStream filePath];
     BOOL fileExists = [[NSFileManager defaultManager] fileExistsAtPath:attachmentFilePath];
     XCTAssert(fileExists);
-    XCTAssertEqual(1, [TSAttachmentStream numberOfItemsInAttachmentsFolder]);
+    XCTAssertEqual(1, [self numberOfItemsInAttachmentsFolder]);
-
+
-    [[OWSOrphanedDataCleaner new] removeOrphanedData];
+    XCTestExpectation *expectation = [self expectationWithDescription:@"Cleanup"];
+    [OWSOrphanedDataCleaner auditAndCleanupAsync:^{
+        [expectation fulfill];
+    }];
+    [self waitForExpectationsWithTimeout:5.0
+                                 handler:^(NSError *error) {
+                                     if (error) {
+                                         XCTFail(@"Expectation Failed with error: %@", error);
+                                     }
+                                 }];
 
     fileExists = [[NSFileManager defaultManager] fileExistsAtPath:attachmentFilePath];
     XCTAssert(fileExists);
-    XCTAssertEqual(1, [TSAttachmentStream numberOfItemsInAttachmentsFolder]);
+    XCTAssertEqual(1, [self numberOfItemsInAttachmentsFolder]);
 }
 
 - (void)testFilesWithoutAttachmentStreamsAreDeleted
@@ -135,12 +184,22 @@
     NSString *orphanedFilePath = [attachmentStream filePath];
     BOOL fileExists = [[NSFileManager defaultManager] fileExistsAtPath:orphanedFilePath];
     XCTAssert(fileExists);
-    XCTAssertEqual(1, [TSAttachmentStream numberOfItemsInAttachmentsFolder]);
+    XCTAssertEqual(1, [self numberOfItemsInAttachmentsFolder]);
+
+    XCTestExpectation *expectation = [self expectationWithDescription:@"Cleanup"];
+    [OWSOrphanedDataCleaner auditAndCleanupAsync:^{
+        [expectation fulfill];
+    }];
+    [self waitForExpectationsWithTimeout:5.0
+                                 handler:^(NSError *error) {
+                                     if (error) {
+                                         XCTFail(@"Expectation Failed with error: %@", error);
+                                     }
+                                 }];
 
-    [[OWSOrphanedDataCleaner new] removeOrphanedData];
     fileExists = [[NSFileManager defaultManager] fileExistsAtPath:orphanedFilePath];
     XCTAssertFalse(fileExists);
-    XCTAssertEqual(0, [TSAttachmentStream numberOfItemsInAttachmentsFolder]);
+    XCTAssertEqual(0, [self numberOfItemsInAttachmentsFolder]);
 }
 
 @end