Issue #99: Test rate limit for max daily upload (#101)

* - add mockable timestamp function
- add error type for max daily upload count
* Add upload limit test. I did some file reorg that has made the commit look bigger.

Telemetry/TelemetryError.swift

@@ -5,11 +5,12 @@
 import Foundation
 
 public class TelemetryError {
-    public static let ErrorDomain: String = "TelemetryErrorDomain"
+    public static let ErrorDomain = "TelemetryErrorDomain"
 
-    public static let SessionAlreadyStarted: Int = 101
-    public static let SessionNotStarted: Int = 102
-    public static let InvalidUploadURL: Int = 103
-    public static let CannotGenerateJSON: Int = 104
-    public static let UnknownUploadError: Int = 105
+    public static let SessionAlreadyStarted = 101
+    public static let SessionNotStarted = 102
+    public static let InvalidUploadURL = 103
+    public static let CannotGenerateJSON = 104
+    public static let UnknownUploadError = 105
+    public static let MaxDailyUploadReached = 106
 }

Telemetry/TelemetryMeasurements.swift

@@ -96,7 +96,7 @@ public class CreatedTimestampMeasurement: TelemetryMeasurement {
     }
 
     override func flush() -> Any? {
-        return UInt64.safeConvert(Date().timeIntervalSince1970 * 1000)
+        return UInt64.safeConvert(TelemetryUtils.timestamp() * 1000)
     }
 }
 
@@ -244,7 +244,7 @@ public class ProfileDateMeasurement: TelemetryMeasurement {
         }
 
         // Fallback to current date if profile cannot be found
-        let seconds = UInt64.safeConvert(Date().timeIntervalSince1970)
+        let seconds = UInt64.safeConvert(TelemetryUtils.timestamp())
         let days = seconds * oneSecondInMilliseconds / oneDayInMilliseconds
 
         return days

Telemetry/TelemetryPingBuilder.swift

@@ -43,7 +43,7 @@ public class TelemetryPingBuilder {
                 data = handler(data)
             }
         }
-        return TelemetryPing(pingType: pingType, documentId: documentId, uploadPath: uploadPath, measurements: data, timestamp: Date().timeIntervalSince1970)
+        return TelemetryPing(pingType: pingType, documentId: documentId, uploadPath: uploadPath, measurements: data, timestamp: TelemetryUtils.timestamp())
     }
 
     func getUploadPath(withDocumentId documentId: String) -> String {

Telemetry/TelemetryScheduler.swift

@@ -20,8 +20,14 @@ class TelemetryScheduler {
         var pingSequence = storage.sequence(forPingType: pingType)
 
         func uploadNextPing() {
-            guard !hasReachedDailyUploadLimit(forPingType: pingType),
-                let ping = pingSequence.next() else {
+            guard let ping = pingSequence.next() else {
+                completionHandler()
+                return
+            }
+
+            guard !hasReachedDailyUploadLimit(forPingType: pingType) else {
+                let error = NSError(domain: TelemetryError.ErrorDomain, code: TelemetryError.MaxDailyUploadReached, userInfo: [NSLocalizedDescriptionKey: "Max daily upload reached."])
+                NotificationCenter.default.post(name: Telemetry.notificationReportError, object: nil, userInfo: ["error": error])
                 completionHandler()
                 return
             }
@@ -50,14 +56,14 @@ class TelemetryScheduler {
     }
 
     private func lastUploadTimestamp(forPingType pingType: String) -> TimeInterval {
-        return storage.get(valueFor: "\(pingType)-lastUploadTimestamp") as? TimeInterval ?? Date().timeIntervalSince1970
+        return storage.get(valueFor: "\(pingType)-lastUploadTimestamp") as? TimeInterval ?? TelemetryUtils.timestamp()
     }
 
     private func incrementDailyUploadCount(forPingType pingType: String) {
         let uploadCount = dailyUploadCount(forPingType: pingType) + 1
         storage.set(key: "\(pingType)-dailyUploadCount", value: uploadCount)
 
-        let lastUploadTimestamp = Date().timeIntervalSince1970
+        let lastUploadTimestamp = TelemetryUtils.timestamp()
         storage.set(key: "\(pingType)-lastUploadTimestamp", value: lastUploadTimestamp)
     }
 
@@ -66,7 +72,7 @@ class TelemetryScheduler {
             storage.set(key: "\(pingType)-dailyUploadCount", value: 0)
             return false
         }
-        
+
         return dailyUploadCount(forPingType: pingType) >= configuration.maximumNumberOfPingUploadsPerDay
     }
 
@@ -76,7 +82,7 @@ class TelemetryScheduler {
         let monthA = Calendar.current.component(.month, from: dateA)
         let yearA = Calendar.current.component(.year, from: dateA)
 
-        let dateB = Date()
+        let dateB = Date(timeIntervalSince1970: TelemetryUtils.timestamp())
         let dayB = Calendar.current.component(.day, from: dateB)
         let monthB = Calendar.current.component(.month, from: dateB)
         let yearB = Calendar.current.component(.year, from: dateB)

Telemetry/TelemetryStorage.swift

@@ -23,7 +23,7 @@ class TelemetryStorageSequence : Sequence, IteratorProtocol {
             return false
         }
 
-        let days = TelemetryUtils.daysBetween(start: time, end: Date())
+        let days = TelemetryUtils.daysOld(date: time)
         return days > configuration.maximumAgeOfPingInDays
     }
 
@@ -115,19 +115,19 @@ public class TelemetryStorage {
     fileprivate let configuration: TelemetryConfiguration
 
     // Prepend to all key usage to avoid UserDefaults name collisions
-    private let keyPrefix = "telemetry-key-prefix-"
+    static let keyPrefix = "telemetry-key-prefix-"
 
     init(name: String, configuration: TelemetryConfiguration) {
         self.name = name
         self.configuration = configuration
     }
 
     func get(valueFor key: String) -> Any? {
-        return UserDefaults.standard.object(forKey: keyPrefix + key)
+        return UserDefaults.standard.object(forKey: TelemetryStorage.keyPrefix + key)
     }
 
     func set(key: String, value: Any) {
-        UserDefaults.standard.set(value, forKey: keyPrefix + key)
+        UserDefaults.standard.set(value, forKey: TelemetryStorage.keyPrefix + key)
     }
 
     func enqueue(ping: TelemetryPing) {
@@ -136,7 +136,7 @@ public class TelemetryStorage {
             return
         }
 
-        var url = directory.appendingPathComponent("-t-\(Date().timeIntervalSince1970).json")
+        var url = directory.appendingPathComponent("-t-\(TelemetryUtils.timestamp()).json")
 
         do {
             // TODO: Check `configuration.maximumNumberOfPingsPerType` and remove oldest ping if necessary.

Telemetry/TelemetryUtils.swift

@@ -37,7 +37,10 @@ extension UInt64 {
     }
 }
 
+
 class TelemetryUtils {
+    static let isUnitTesting = ProcessInfo.processInfo.environment["XCTestConfigurationFilePath"] != nil || (ProcessInfo.processInfo.environment["DYLD_INSERT_LIBRARIES"] ?? "").contains("libXCTTargetBootstrapInject.dylib")
+
     static func asString(_ object: Any?) -> String {
         if let string = object as? String {
             return string
@@ -57,7 +60,29 @@ class TelemetryUtils {
         return string
     }
 
-    static func daysBetween(start: Date, end: Date) -> Int {
-        return Calendar.current.dateComponents([.day], from: start, to: end).day ?? 0
+    static func daysOld(date: Date) -> Int {
+        let end = TelemetryUtils.dateFromTimestamp(TelemetryUtils.timestamp())
+        return Calendar.current.dateComponents([.day], from: date, to: end).day ?? 0
     }
 }
+
+// Allows for adjusting the time when testing
+extension TelemetryUtils {
+    static var mockableOffset: TimeInterval? {
+        didSet {
+            if !isUnitTesting {
+                // Testing only!!
+                mockableOffset = nil
+            }
+        }
+    }
+
+    static func timestamp() -> TimeInterval {
+        return Date().timeIntervalSince1970 + (mockableOffset ?? 0)
+    }
+
+    static func dateFromTimestamp(_ timestampSince1970: TimeInterval) -> Date {
+        return Date(timeIntervalSince1970: timestampSince1970 + (mockableOffset ?? 0))
+    }
+}
+

TelemetryTests/TelemetryTests.swift

@@ -9,10 +9,16 @@ import OHHTTPStubs
 
 class TelemetryTests: XCTestCase {
     var expectation: XCTestExpectation? = nil
+    var errorNotificationExpectation: XCTestExpectation? = nil
 
     override func setUp() {
         super.setUp()
 
+        let prefKeys = UserDefaults.standard.dictionaryRepresentation().keys
+        prefKeys.filter { $0.hasPrefix(TelemetryStorage.keyPrefix) }.forEach {
+            UserDefaults.standard.removeObject(forKey: $0)
+        }
+
         let telemetryConfig = Telemetry.default.configuration
         telemetryConfig.appName = "AppInfo.displayName"
         telemetryConfig.userDefaultsSuiteName = "AppInfo.sharedContainerIdentifier"
@@ -28,6 +34,17 @@ class TelemetryTests: XCTestCase {
 
             Telemetry.default.storage.deleteEventArrayFile(forPingType: t)
         }
+
+        NotificationCenter.default.removeObserver(self)
+        NotificationCenter.default.addObserver(self, selector: #selector(TelemetryTests.uploadError(notification:)), name: Telemetry.notificationReportError, object: nil)
+
+    }
+
+    @objc func uploadError(notification: NSNotification) {
+        guard let error = notification.userInfo?["error"] as? NSError else { return }
+        print("Upload error notification: \(error.localizedDescription)")
+        errorNotificationExpectation?.fulfill()
+        errorNotificationExpectation = nil
     }
 
     override func tearDown() {
@@ -76,8 +93,13 @@ class TelemetryTests: XCTestCase {
         }
     }
 
-    private func upload(pingType: String) {
-        expectation = expectation(description: "Completed upload")
+    private func upload(pingType: String, expectUploadError: Bool = false) {
+        if expectUploadError {
+            errorNotificationExpectation = expectation(description: "Expect upload error notification")
+        } else {
+            expectation = expectation(description: "Completed upload")
+        }
+
         Telemetry.default.scheduleUpload(pingType: pingType)
         waitForExpectations(timeout: 10.0) { error in
             if error != nil {
@@ -86,7 +108,7 @@ class TelemetryTests: XCTestCase {
         }
     }
 
-    private func storeOnDiskAndUpload(corePingFilesToWrite: Int) {
+    private func storeOnDiskAndUpload(corePingFilesToWrite: Int, expectUploadError: Bool = false) {
         let startSeq = Telemetry.default.storage.get(valueFor: "\(CorePingBuilder.PingType)-seq") as? Int ?? 0
 
         for _ in 0..<corePingFilesToWrite {
@@ -96,15 +118,32 @@ class TelemetryTests: XCTestCase {
         }
 
         waitForFilesOnDisk(count: corePingFilesToWrite)
-        upload(pingType: CorePingBuilder.PingType)
+        upload(pingType: CorePingBuilder.PingType, expectUploadError: expectUploadError)
 
         let endSeq = Telemetry.default.storage.get(valueFor: "\(CorePingBuilder.PingType)-seq") as! Int
         XCTAssert(endSeq - startSeq == corePingFilesToWrite)
     }
 
     private func waitForFilesOnDisk(count: Int, pingType: String = CorePingBuilder.PingType) {
         wait()
-        XCTAssert(countFilesOnDisk(forPingType: pingType) == count, "waitForFilesOnDisk matching")
+        XCTAssert(countFilesOnDisk(forPingType: pingType) >= count, "waitForFilesOnDisk \(countFilesOnDisk(forPingType: pingType)) >= \(count)")
+    }
+
+    private func countFilesOnDisk(forPingType pingType: String = CorePingBuilder.PingType) -> Int {
+        var result = 0
+        let seq = Telemetry.default.storage.sequence(forPingType: pingType)
+        while seq.next() != nil {
+            result += 1
+        }
+        return result
+    }
+
+    private func serverError(code: URLError.Code) {
+        setupHttpErrorStub(expectedFilesUploaded: 1, statusCode: code)
+        let filesOnDisk = 3
+        // Only one attempted upload, but all 3 files should remain on disk.
+        storeOnDiskAndUpload(corePingFilesToWrite: filesOnDisk)
+        waitForFilesOnDisk(count: filesOnDisk)
     }
 
     private func wait() {
@@ -141,23 +180,6 @@ class TelemetryTests: XCTestCase {
         serverError(code: URLError.Code.badServerResponse)
     }
 
-    private func countFilesOnDisk(forPingType pingType: String = CorePingBuilder.PingType) -> Int {
-        var result = 0
-        let seq = Telemetry.default.storage.sequence(forPingType: pingType)
-        while seq.next() != nil {
-            result += 1
-        }
-        return result
-    }
-
-    private func serverError(code: URLError.Code) {
-        setupHttpErrorStub(expectedFilesUploaded: 1, statusCode: code)
-        let filesOnDisk = 3
-        // Only one attempted upload, but all 3 files should remain on disk.
-        storeOnDiskAndUpload(corePingFilesToWrite: filesOnDisk)
-        waitForFilesOnDisk(count: filesOnDisk)
-    }
-
 
     func test4xxCode() {
         setupHttpResponseStub(expectedFilesUploaded: 3, statusCode: 400)
@@ -176,7 +198,31 @@ class TelemetryTests: XCTestCase {
         let fileDate = TelemetryStorage.extractTimestampFromName(pingFile: URL(string: "a/b/c/foo-t-\(lastWeek.timeIntervalSince1970).json")!)!
         XCTAssert(fileDate.timeIntervalSince1970 > 0)
         XCTAssert(fabs(fileDate.timeIntervalSince1970 - lastWeek.timeIntervalSince1970) < 0.1 /* epsilon */)
-        XCTAssert(TelemetryUtils.daysBetween(start: fileDate, end: Date()) == 7)
+        XCTAssert(TelemetryUtils.daysOld(date: fileDate) == 7)
+    }
+
+
+    func testDailyUploadLimit() {
+        let max = 5
+        Telemetry.default.configuration.maximumNumberOfPingUploadsPerDay = max
+        for _ in 0..<max {
+            setupHttpResponseStub(expectedFilesUploaded: 1, statusCode: 200)
+            storeOnDiskAndUpload(corePingFilesToWrite: 1)
+            waitForFilesOnDisk(count: 0)
+        }
+
+        setupHttpResponseStub(expectedFilesUploaded: 1, statusCode: 200)
+        storeOnDiskAndUpload(corePingFilesToWrite: 1, expectUploadError: true)
+        waitForFilesOnDisk(count: 1)
+
+        TelemetryUtils.mockableOffset = 86400 // 1 day of seconds
+
+        setupHttpResponseStub(expectedFilesUploaded: 1, statusCode: 200)
+        storeOnDiskAndUpload(corePingFilesToWrite: 1)
+        waitForFilesOnDisk(count: 0)
+
+        TelemetryUtils.mockableOffset = nil
+        Telemetry.default.configuration.maximumNumberOfPingUploadsPerDay = 100
     }
 }