Add PerMessageExpirationTest.

SignalServiceKit/src/TestUtils/Factories.swift

@@ -142,7 +142,7 @@ public class OutgoingMessageFactory: NSObject, Factory {
                                      contactShare: contactShareBuilder(),
                                      linkPreview: linkPreviewBuilder(),
                                      messageSticker: messageStickerBuilder(),
-                                     perMessageExpirationDurationSeconds: 0)
+                                     perMessageExpirationDurationSeconds: perMessageExpirationDurationSecondsBuilder())
 
         return item
     }
@@ -219,6 +219,11 @@ public class OutgoingMessageFactory: NSObject, Factory {
         return nil
     }
 
+    @objc
+    public var perMessageExpirationDurationSecondsBuilder: () -> UInt32 = {
+        return 0
+    }
+
     // MARK: Delivery Receipts
 
     @objc
@@ -266,7 +271,7 @@ public class IncomingMessageFactory: NSObject, Factory {
                                      messageSticker: messageStickerBuilder(),
                                      serverTimestamp: serverTimestampBuilder(),
                                      wasReceivedByUD: wasReceivedByUDBuilder(),
-                                     perMessageExpirationDurationSeconds: 0)
+                                     perMessageExpirationDurationSeconds: perMessageExpirationDurationSecondsBuilder())
 
         item.anyInsert(transaction: transaction)
 
@@ -349,6 +354,11 @@ public class IncomingMessageFactory: NSObject, Factory {
     public var wasReceivedByUDBuilder: () -> Bool = {
         return false
     }
+
+    @objc
+    public var perMessageExpirationDurationSecondsBuilder: () -> UInt32 = {
+        return 0
+    }
 }
 
 @objc

SignalServiceKit/src/Util/PerMessageExpiration.swift

@@ -75,32 +75,40 @@ public class PerMessageExpiration: NSObject {
         let perMessageExpiresAtMS = message.perMessageExpiresAt
 
         guard perMessageExpiresAtMS > nowMs else {
+            Logger.verbose("Expiring immediately.")
             // Message has expired; remove it immediately.
             completePerMessageExpiration(forMessage: message,
                                          transaction: transaction)
             return
         }
 
+        Logger.verbose("Scheduling expiration.")
         let delaySeconds: TimeInterval = Double(perMessageExpiresAtMS - nowMs) / 1000
         DispatchQueue.global().asyncAfter(wallDeadline: .now() + delaySeconds) {
-            self.completePerMessageExpiration(forMessage: message)
-        }
-    }
-
-    private class func completePerMessageExpiration(forMessage message: TSMessage) {
-        databaseStorage.write { (transaction) in
-            self.completePerMessageExpiration(forMessage: message,
-                                              transaction: transaction)
+            databaseStorage.write { (transaction) in
+                self.completePerMessageExpiration(forMessage: message,
+                                                  transaction: transaction)
+            }
         }
     }
 
     private class func completePerMessageExpiration(forMessage message: TSMessage,
                                                     transaction: SDSAnyWriteTransaction) {
-        guard message.perMessageExpirationHasExpired else {
+        Logger.verbose("Expiring message per schedule.")
+        guard let uniqueId = message.uniqueId else {
+            owsFailDebug("Missing uniqueId.")
+            return
+        }
+        guard let latestCopy = TSMessage.anyFetch(uniqueId: uniqueId, transaction: transaction) as? TSMessage else {
+            // Message has been deleted.
+            Logger.warn("Couldn't expire message; not in database.")
+            return
+        }
+        guard latestCopy.perMessageExpirationHasExpired else {
             // Already expired, no need to expire again.
             return
         }
-        message.updateWithHasPerMessageExpiredAndRemoveRenderableContent(with: transaction)
+        latestCopy.updateWithHasPerMessageExpiredAndRemoveRenderableContent(with: transaction)
     }
 
     // MARK: - Events

SignalServiceKit/tests/Util/PerMessageExpirationTest.swift

@@ -0,0 +1,101 @@
+//
+//  Copyright (c) 2019 Open Whisper Systems. All rights reserved.
+//
+
+import Foundation
+import XCTest
+
+@testable import SignalServiceKit
+
+class PerMessageExpirationTest: SSKBaseTestSwift {
+
+    override func setUp() {
+        super.setUp()
+    }
+
+    override func tearDown() {
+        super.tearDown()
+    }
+
+    // MARK: 
+
+    func test_expiration() {
+
+        let messageCount = { () -> Int in
+            return self.databaseStorage.readReturningResult { transaction in
+                return TSInteraction.anyFetchAll(transaction: transaction).count
+            }
+        }
+        let latestCopy = { (message: TSMessage) -> TSMessage in
+            let uniqueId = message.uniqueId!
+            return self.databaseStorage.readReturningResult { transaction in
+                return TSMessage.anyFetch(uniqueId: uniqueId, transaction: transaction) as! TSMessage
+            }
+        }
+
+        // Factory 1 builds messages without per-message expiration.
+        let factory1 = IncomingMessageFactory()
+        factory1.perMessageExpirationDurationSecondsBuilder = {
+            return 0
+        }
+
+        self.write { transaction in
+            _ = factory1.create(transaction: transaction)
+        }
+
+        XCTAssertEqual(1, messageCount())
+
+        // Factory 2 builds messages with per-message expiration.
+        let expirationSeconds: UInt32 = 1
+        let factory2 = IncomingMessageFactory()
+        factory2.perMessageExpirationDurationSecondsBuilder = {
+            return expirationSeconds
+        }
+
+        var disappearingMessage: TSMessage?
+        self.write { transaction in
+            disappearingMessage = factory2.create(transaction: transaction)
+        }
+
+        XCTAssertEqual(2, messageCount())
+
+        guard let message = disappearingMessage else {
+            XCTFail("Missing message.")
+            return
+        }
+
+        XCTAssertTrue(message.hasPerMessageExpiration)
+        XCTAssertFalse(message.hasPerMessageExpirationStarted)
+        XCTAssertFalse(message.perMessageExpirationHasExpired)
+
+        XCTAssertTrue(latestCopy(message).hasPerMessageExpiration)
+        XCTAssertFalse(latestCopy(message).hasPerMessageExpirationStarted)
+        XCTAssertFalse(latestCopy(message).perMessageExpirationHasExpired)
+
+        self.write { transaction in
+            PerMessageExpiration.startPerMessageExpiration(forMessage: message,
+                                                           transaction: transaction)
+        }
+
+        XCTAssertTrue(message.hasPerMessageExpiration)
+        XCTAssertTrue(message.hasPerMessageExpirationStarted)
+        XCTAssertFalse(message.perMessageExpirationHasExpired)
+
+        XCTAssertTrue(latestCopy(message).hasPerMessageExpiration)
+        XCTAssertTrue(latestCopy(message).hasPerMessageExpirationStarted)
+        XCTAssertFalse(latestCopy(message).perMessageExpirationHasExpired)
+
+        XCTAssertEqual(2, messageCount())
+
+        Logger.verbose("Sleeping for \(expirationSeconds + 1)")
+        // Sleep for a little extra time to avoid races.
+        sleep(expirationSeconds + 1)
+        Logger.verbose("Slept for \(expirationSeconds + 1)")
+
+        XCTAssertTrue(latestCopy(message).hasPerMessageExpiration)
+        XCTAssertTrue(latestCopy(message).hasPerMessageExpirationStarted)
+        XCTAssertTrue(latestCopy(message).perMessageExpirationHasExpired)
+
+        XCTAssertEqual(2, messageCount())
+    }
+}