Implement initial tests: start rotation on account set, stop on account unset

9 files changed, 148 insertions, 32 deletions

diff --git a/ios/MullvadREST/ApiHandlers/RESTAccountsProxy.swift b/ios/MullvadREST/ApiHandlers/RESTAccountsProxy.swift
index 4088a2abb4..568e24face 100644
--- a/ios/MullvadREST/ApiHandlers/RESTAccountsProxy.swift
+++ b/ios/MullvadREST/ApiHandlers/RESTAccountsProxy.swift
@@ -145,3 +145,17 @@ extension REST {
         public let number: String
     }
 }
+
+extension REST.NewAccountData {
+    public static func mockValue() -> REST.NewAccountData {
+        return REST.NewAccountData(
+            id: UUID().uuidString,
+            expiry: Date().addingTimeInterval(3600),
+            maxPorts: 2,
+            canAddPorts: false,
+            maxDevices: 5,
+            canAddDevices: false,
+            number: "1234567890123456"
+        )
+    }
+}
diff --git a/ios/MullvadVPN.xcodeproj/project.pbxproj b/ios/MullvadVPN.xcodeproj/project.pbxproj
index 6e6b63475b..9536142933 100644
--- a/ios/MullvadVPN.xcodeproj/project.pbxproj
+++ b/ios/MullvadVPN.xcodeproj/project.pbxproj
@@ -42,6 +42,8 @@
 		06AC116228F94C450037AF9A /* ApplicationConfiguration.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58BFA5CB22A7CE1F00A6173D /* ApplicationConfiguration.swift */; };
 		449872E12B7BBC5400094DDC /* TunnelSettingsUpdate.swift in Sources */ = {isa = PBXBuildFile; fileRef = 449872E02B7BBC5400094DDC /* TunnelSettingsUpdate.swift */; };
 		449872E42B7CB96300094DDC /* TunnelSettingsUpdateTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 449872E32B7CB96300094DDC /* TunnelSettingsUpdateTests.swift */; };
+		449EB9FD2B95F8AD00DFA4EB /* DeviceMock.swift in Sources */ = {isa = PBXBuildFile; fileRef = 449EB9FC2B95F8AD00DFA4EB /* DeviceMock.swift */; };
+		449EB9FF2B95FF2500DFA4EB /* AccountMock.swift in Sources */ = {isa = PBXBuildFile; fileRef = 449EB9FE2B95FF2500DFA4EB /* AccountMock.swift */; };
 		44DD7D242B6CFFD70005F67F /* StartTunnelOperationTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 44DD7D232B6CFFD70005F67F /* StartTunnelOperationTests.swift */; };
 		44DD7D272B6D18FB0005F67F /* MockTunnelInteractor.swift in Sources */ = {isa = PBXBuildFile; fileRef = 44DD7D262B6D18FB0005F67F /* MockTunnelInteractor.swift */; };
 		44DD7D292B7113CA0005F67F /* MockTunnel.swift in Sources */ = {isa = PBXBuildFile; fileRef = 44DD7D282B7113CA0005F67F /* MockTunnel.swift */; };
@@ -1294,6 +1296,8 @@
 		06FAE67D28F83CA50033DD93 /* RESTTransport.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = RESTTransport.swift; sourceTree = "<group>"; };
 		449872E02B7BBC5400094DDC /* TunnelSettingsUpdate.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TunnelSettingsUpdate.swift; sourceTree = "<group>"; };
 		449872E32B7CB96300094DDC /* TunnelSettingsUpdateTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TunnelSettingsUpdateTests.swift; sourceTree = "<group>"; };
+		449EB9FC2B95F8AD00DFA4EB /* DeviceMock.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = DeviceMock.swift; sourceTree = "<group>"; };
+		449EB9FE2B95FF2500DFA4EB /* AccountMock.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AccountMock.swift; sourceTree = "<group>"; };
 		44DD7D232B6CFFD70005F67F /* StartTunnelOperationTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = StartTunnelOperationTests.swift; sourceTree = "<group>"; };
 		44DD7D262B6D18FB0005F67F /* MockTunnelInteractor.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = MockTunnelInteractor.swift; sourceTree = "<group>"; };
 		44DD7D282B7113CA0005F67F /* MockTunnel.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = MockTunnel.swift; sourceTree = "<group>"; };
@@ -2202,7 +2206,9 @@
 		44DD7D252B6D18E90005F67F /* Mocks */ = {
 			isa = PBXGroup;
 			children = (
+				449EB9FE2B95FF2500DFA4EB /* AccountMock.swift */,
 				7A9BE5A82B90806800E2A7D0 /* CustomListsRepositoryStub.swift */,
+				449EB9FC2B95F8AD00DFA4EB /* DeviceMock.swift */,
 				44DD7D282B7113CA0005F67F /* MockTunnel.swift */,
 				44DD7D262B6D18FB0005F67F /* MockTunnelInteractor.swift */,
 			);
@@ -4854,6 +4860,7 @@
 				A9A5FA272ACB05160083449F /* VPNConnectionProtocol.swift in Sources */,
 				A9A5FA282ACB05160083449F /* WgKeyRotation.swift in Sources */,
 				449872E42B7CB96300094DDC /* TunnelSettingsUpdateTests.swift in Sources */,
+				449EB9FD2B95F8AD00DFA4EB /* DeviceMock.swift in Sources */,
 				A9A5FA292ACB05160083449F /* AddressCacheTests.swift in Sources */,
 				A9B6AC182ADE8F4300F7802A /* MigrationManagerTests.swift in Sources */,
 				7A9BE5AB2B909A1700E2A7D0 /* LocationDataSourceProtocol.swift in Sources */,
@@ -4866,6 +4873,7 @@
 				A9A5FA2F2ACB05160083449F /* FixedWidthIntegerArithmeticsTests.swift in Sources */,
 				A9A5FA302ACB05160083449F /* InputTextFormatterTests.swift in Sources */,
 				F0B0E6972AFE6E7E001DC66B /* XCTest+Async.swift in Sources */,
+				449EB9FF2B95FF2500DFA4EB /* AccountMock.swift in Sources */,
 				7ADCB2DA2B6A730400C88F89 /* IPOverrideRepositoryStub.swift in Sources */,
 				A9A5FA312ACB05160083449F /* MockFileCache.swift in Sources */,
 				A9A5FA322ACB05160083449F /* RelayCacheTests.swift in Sources */,
diff --git a/ios/MullvadVPN/TunnelManager/TunnelManager.swift b/ios/MullvadVPN/TunnelManager/TunnelManager.swift
index 3f0e69459f..b3d53b5613 100644
--- a/ios/MullvadVPN/TunnelManager/TunnelManager.swift
+++ b/ios/MullvadVPN/TunnelManager/TunnelManager.swift
@@ -61,7 +61,8 @@ final class TunnelManager: StorePaymentObserver {
     private var networkMonitor: NWPathMonitor?
 
     private var privateKeyRotationTimer: DispatchSourceTimer?
-    private var isRunningPeriodicPrivateKeyRotation = false
+    public private(set) var isRunningPeriodicPrivateKeyRotation = false
+    public private(set) var nextKeyRotationDate: Date?
 
     private var tunnelStatusPollTimer: DispatchSourceTimer?
     private var isPolling = false
@@ -111,7 +112,7 @@ final class TunnelManager: StorePaymentObserver {
         nslock.lock()
         defer { nslock.unlock() }
 
-        guard !isRunningPeriodicPrivateKeyRotation else { return }
+        guard !isRunningPeriodicPrivateKeyRotation, deviceState.isLoggedIn else { return }
 
         logger.debug("Start periodic private key rotation.")
 
@@ -131,6 +132,14 @@ final class TunnelManager: StorePaymentObserver {
         updatePrivateKeyRotationTimer()
     }
 
+    func startOrStopPeriodicPrivateKeyRotation() {
+        if deviceState.isLoggedIn {
+            startPeriodicPrivateKeyRotation()
+        } else {
+            stopPeriodicPrivateKeyRotation()
+        }
+    }
+
     func getNextKeyRotationDate() -> Date? {
         nslock.lock()
         defer { nslock.unlock() }
@@ -144,9 +153,11 @@ final class TunnelManager: StorePaymentObserver {
 
         privateKeyRotationTimer?.cancel()
         privateKeyRotationTimer = nil
+        nextKeyRotationDate = nil
 
         guard isRunningPeriodicPrivateKeyRotation,
               let scheduleDate = getNextKeyRotationDate() else { return }
+        nextKeyRotationDate = scheduleDate
 
         let timer = DispatchSource.makeTimerSource(queue: .main)
 
@@ -334,7 +345,8 @@ final class TunnelManager: StorePaymentObserver {
 
         operation.completionQueue = .main
         operation.completionHandler = { [weak self] result in
-            self?.updatePrivateKeyRotationTimer()
+            guard let self else { return }
+            startOrStopPeriodicPrivateKeyRotation()
 
             completionHandler(result)
         }
diff --git a/ios/MullvadVPNTests/AccountsProxy+Stubs.swift b/ios/MullvadVPNTests/AccountsProxy+Stubs.swift
index 18e940e8ae..76e5d17ccb 100644
--- a/ios/MullvadVPNTests/AccountsProxy+Stubs.swift
+++ b/ios/MullvadVPNTests/AccountsProxy+Stubs.swift
@@ -11,11 +11,15 @@ import Foundation
 @testable import MullvadTypes
 
 struct AccountsProxyStub: RESTAccountHandling {
+    var createAccountResult: Result<REST.NewAccountData, Error>?
     func createAccount(
         retryStrategy: REST.RetryStrategy,
         completion: @escaping MullvadREST.ProxyCompletionHandler<REST.NewAccountData>
     ) -> Cancellable {
-        AnyCancellable()
+        if let createAccountResult = createAccountResult {
+            completion(createAccountResult)
+        }
+        return AnyCancellable()
     }
 
     func getAccountData(accountNumber: String) -> any RESTRequestExecutor<Account> {
diff --git a/ios/MullvadVPNTests/DeviceCheckOperationTests.swift b/ios/MullvadVPNTests/DeviceCheckOperationTests.swift
index aba1eeb71f..2dc3914db9 100644
--- a/ios/MullvadVPNTests/DeviceCheckOperationTests.swift
+++ b/ios/MullvadVPNTests/DeviceCheckOperationTests.swift
@@ -526,31 +526,6 @@ private extension StoredDeviceData {
     }
 }
 
-private extension Device {
-    static func mock(publicKey: PublicKey) -> Device {
-        Device(
-            id: "device-id",
-            name: "device-name",
-            pubkey: publicKey,
-            hijackDNS: false,
-            created: Date(),
-            ipv4Address: IPAddressRange(from: "127.0.0.1/32")!,
-            ipv6Address: IPAddressRange(from: "::ff/64")!
-        )
-    }
-}
-
-private extension Account {
-    static func mock(expiry: Date = .distantFuture) -> Account {
-        Account(
-            id: "account-id",
-            expiry: expiry,
-            maxDevices: 5,
-            canAddDevices: true
-        )
-    }
-}
-
 private extension KeyRotationStatus {
     /// Returns `true` if key rotation status is `.attempted`.
     var isAttempted: Bool {
diff --git a/ios/MullvadVPNTests/DevicesProxy+Stubs.swift b/ios/MullvadVPNTests/DevicesProxy+Stubs.swift
index 3ec0f3faff..69d0251f81 100644
--- a/ios/MullvadVPNTests/DevicesProxy+Stubs.swift
+++ b/ios/MullvadVPNTests/DevicesProxy+Stubs.swift
@@ -12,6 +12,7 @@ import Foundation
 @testable import WireGuardKitTypes
 
 struct DevicesProxyStub: DeviceHandling {
+    let mockDevice = Device.mock(publicKey: PrivateKey().publicKey)
     func getDevice(
         accountNumber: String,
         identifier: String,
@@ -35,7 +36,8 @@ struct DevicesProxyStub: DeviceHandling {
         retryStrategy: REST.RetryStrategy,
         completion: @escaping ProxyCompletionHandler<Device>
     ) -> Cancellable {
-        AnyCancellable()
+        completion(.success(mockDevice))
+        return AnyCancellable()
     }
 
     func deleteDevice(
@@ -44,7 +46,8 @@ struct DevicesProxyStub: DeviceHandling {
         retryStrategy: REST.RetryStrategy,
         completion: @escaping ProxyCompletionHandler<Bool>
     ) -> Cancellable {
-        AnyCancellable()
+        completion(.success(true))
+        return AnyCancellable()
     }
 
     func rotateDeviceKey(
diff --git a/ios/MullvadVPNTests/Mocks/AccountMock.swift b/ios/MullvadVPNTests/Mocks/AccountMock.swift
new file mode 100644
index 0000000000..70ab1d10b3
--- /dev/null
+++ b/ios/MullvadVPNTests/Mocks/AccountMock.swift
@@ -0,0 +1,20 @@
+//
+//  AccountMock.swift
+//  MullvadVPNTests
+//
+//  Created by Andrew Bulhak on 2024-03-04.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import MullvadTypes
+
+extension Account {
+    static func mock(expiry: Date = .distantFuture) -> Account {
+        Account(
+            id: "account-id",
+            expiry: expiry,
+            maxDevices: 5,
+            canAddDevices: true
+        )
+    }
+}
diff --git a/ios/MullvadVPNTests/Mocks/DeviceMock.swift b/ios/MullvadVPNTests/Mocks/DeviceMock.swift
new file mode 100644
index 0000000000..9f17410dad
--- /dev/null
+++ b/ios/MullvadVPNTests/Mocks/DeviceMock.swift
@@ -0,0 +1,25 @@
+//
+//  DeviceMock.swift
+//  MullvadVPNTests
+//
+//  Created by Andrew Bulhak on 2024-03-04.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import Foundation
+import MullvadTypes
+import WireGuardKitTypes
+
+extension Device {
+    static func mock(publicKey: PublicKey) -> Device {
+        Device(
+            id: "device-id",
+            name: "Devicey McDeviceface",
+            pubkey: publicKey,
+            hijackDNS: false,
+            created: Date(),
+            ipv4Address: IPAddressRange(from: "127.0.0.1/32")!,
+            ipv6Address: IPAddressRange(from: "::ff/64")!
+        )
+    }
+}
diff --git a/ios/MullvadVPNTests/TunnelManagerTests.swift b/ios/MullvadVPNTests/TunnelManagerTests.swift
index 958c50e252..af51b48468 100644
--- a/ios/MullvadVPNTests/TunnelManagerTests.swift
+++ b/ios/MullvadVPNTests/TunnelManagerTests.swift
@@ -6,9 +6,21 @@
 //  Copyright © 2023 Mullvad VPN AB. All rights reserved.
 //
 
+import MullvadREST
+@testable import MullvadSettings
 import XCTest
 
 final class TunnelManagerTests: XCTestCase {
+    static let store = InMemorySettingsStore<SettingNotFound>()
+
+    override class func setUp() {
+        SettingsManager.unitTestStore = store
+    }
+
+    override class func tearDown() {
+        SettingsManager.unitTestStore = nil
+    }
+
     func testTunnelManager() {
         let application = UIApplicationStub()
         let tunnelStore = TunnelStoreStub()
@@ -17,7 +29,27 @@ final class TunnelManagerTests: XCTestCase {
         let devicesProxy = DevicesProxyStub()
         let apiProxy = APIProxyStub()
         let accessTokenManager = AccessTokenManagerStub()
+        let tunnelManager = TunnelManager(
+            application: application,
+            tunnelStore: tunnelStore,
+            relayCacheTracker: relayCacheTracker,
+            accountsProxy: accountProxy,
+            devicesProxy: devicesProxy,
+            apiProxy: apiProxy,
+            accessTokenManager: accessTokenManager
+        )
+        XCTAssertNotNil(tunnelManager)
+    }
 
+    func testLogInStartsKeyRotations() async throws {
+        let application = UIApplicationStub()
+        let tunnelStore = TunnelStoreStub()
+        let relayCacheTracker = RelayCacheTrackerStub()
+        var accountProxy = AccountsProxyStub()
+        let devicesProxy = DevicesProxyStub()
+        let apiProxy = APIProxyStub()
+        let accessTokenManager = AccessTokenManagerStub()
+        accountProxy.createAccountResult = .success(REST.NewAccountData.mockValue())
         let tunnelManager = TunnelManager(
             application: application,
             tunnelStore: tunnelStore,
@@ -27,7 +59,30 @@ final class TunnelManagerTests: XCTestCase {
             apiProxy: apiProxy,
             accessTokenManager: accessTokenManager
         )
+        _ = try await tunnelManager.setNewAccount()
+        XCTAssertEqual(tunnelManager.isRunningPeriodicPrivateKeyRotation, true)
+    }
 
-        XCTAssertNotNil(tunnelManager)
+    func testLogOutStopsKeyRotations() async throws {
+        let application = UIApplicationStub()
+        let tunnelStore = TunnelStoreStub()
+        let relayCacheTracker = RelayCacheTrackerStub()
+        var accountProxy = AccountsProxyStub()
+        let devicesProxy = DevicesProxyStub()
+        let apiProxy = APIProxyStub()
+        let accessTokenManager = AccessTokenManagerStub()
+        accountProxy.createAccountResult = .success(REST.NewAccountData.mockValue())
+        let tunnelManager = TunnelManager(
+            application: application,
+            tunnelStore: tunnelStore,
+            relayCacheTracker: relayCacheTracker,
+            accountsProxy: accountProxy,
+            devicesProxy: devicesProxy,
+            apiProxy: apiProxy,
+            accessTokenManager: accessTokenManager
+        )
+        _ = try await tunnelManager.setNewAccount()
+        await tunnelManager.unsetAccount()
+        XCTAssertEqual(tunnelManager.isRunningPeriodicPrivateKeyRotation, false)
     }
 }