Add test for login to account with too many devices

8 files changed, 116 insertions, 9 deletions

diff --git a/ios/MullvadVPN.xcodeproj/project.pbxproj b/ios/MullvadVPN.xcodeproj/project.pbxproj
index ea4432d9ef..f178e1cdcd 100644
--- a/ios/MullvadVPN.xcodeproj/project.pbxproj
+++ b/ios/MullvadVPN.xcodeproj/project.pbxproj
@@ -643,6 +643,7 @@
 		8587A05D2B84D43100152938 /* ChangeLogAlert.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8587A05C2B84D43100152938 /* ChangeLogAlert.swift */; };
 		8590896F2B61763B003AF5F5 /* LoggedOutUITestCase.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8590896B2B61763B003AF5F5 /* LoggedOutUITestCase.swift */; };
 		85A42B862BB1D627007BABF7 /* XCUIElement+Extensions.swift in Sources */ = {isa = PBXBuildFile; fileRef = 85A42B852BB1D627007BABF7 /* XCUIElement+Extensions.swift */; };
+		85A42B882BB44D31007BABF7 /* DeviceManagementPage.swift in Sources */ = {isa = PBXBuildFile; fileRef = 85A42B872BB44D31007BABF7 /* DeviceManagementPage.swift */; };
 		85B267612B849ADB0098E3CD /* mullvad-api.h in Headers */ = {isa = PBXBuildFile; fileRef = 85B267602B849ADB0098E3CD /* mullvad-api.h */; };
 		85C7A2E92B89024B00035D5A /* SettingsTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 85C7A2E82B89024B00035D5A /* SettingsTests.swift */; };
 		85D039982BA4711800940E7F /* SettingsMigrationTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 85D039972BA4711800940E7F /* SettingsMigrationTests.swift */; };
@@ -1906,6 +1907,7 @@
 		8590896A2B61763B003AF5F5 /* BaseUITestCase.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = BaseUITestCase.swift; sourceTree = "<group>"; };
 		8590896B2B61763B003AF5F5 /* LoggedOutUITestCase.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = LoggedOutUITestCase.swift; sourceTree = "<group>"; };
 		85A42B852BB1D627007BABF7 /* XCUIElement+Extensions.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "XCUIElement+Extensions.swift"; sourceTree = "<group>"; };
+		85A42B872BB44D31007BABF7 /* DeviceManagementPage.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = DeviceManagementPage.swift; sourceTree = "<group>"; };
 		85B267602B849ADB0098E3CD /* mullvad-api.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; name = "mullvad-api.h"; path = "../../mullvad-api/include/mullvad-api.h"; sourceTree = "<group>"; };
 		85C7A2E82B89024B00035D5A /* SettingsTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SettingsTests.swift; sourceTree = "<group>"; };
 		85D039972BA4711800940E7F /* SettingsMigrationTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SettingsMigrationTests.swift; sourceTree = "<group>"; };
@@ -3646,6 +3648,7 @@
 				8529693B2B4F0257007EAD4C /* Alert.swift */,
 				8587A05C2B84D43100152938 /* ChangeLogAlert.swift */,
 				852A26452BA9C9CB006EB9C8 /* DNSSettingsPage.swift */,
+				85A42B872BB44D31007BABF7 /* DeviceManagementPage.swift */,
 				85557B1D2B5FB8C700795FE1 /* HeaderBar.swift */,
 				852969342B4E9270007EAD4C /* LoginPage.swift */,
 				85139B2C2B84B4A700734217 /* OutOfTimePage.swift */,
@@ -5661,6 +5664,7 @@
 				85557B162B5ABBBE00795FE1 /* XCUIElementQuery+Extensions.swift in Sources */,
 				855D9F5B2B63E56B00D7C64D /* ProblemReportPage.swift in Sources */,
 				8529693A2B4F0238007EAD4C /* TermsOfServicePage.swift in Sources */,
+				85A42B882BB44D31007BABF7 /* DeviceManagementPage.swift in Sources */,
 				8532E6872B8CCED600ACECD1 /* ProblemReportSubmittedPage.swift in Sources */,
 				85FB5A0C2B6903990015DCED /* WelcomePage.swift in Sources */,
 				852A26462BA9C9CB006EB9C8 /* DNSSettingsPage.swift in Sources */,
diff --git a/ios/MullvadVPN/Classes/AccessbilityIdentifier.swift b/ios/MullvadVPN/Classes/AccessbilityIdentifier.swift
index cda791d042..174408f0af 100644
--- a/ios/MullvadVPN/Classes/AccessbilityIdentifier.swift
+++ b/ios/MullvadVPN/Classes/AccessbilityIdentifier.swift
@@ -16,15 +16,19 @@ public enum AccessibilityIdentifier: String {
     case applyButton
     case cancelButton
     case connectionPanelButton
+    case continueWithLoginButton
     case collapseButton
     case expandButton
     case createAccountButton
     case deleteButton
+    case deviceCellRemoveButton
     case disconnectButton
     case revokedDeviceLoginButton
     case dnsSettingsEditButton
     case infoButton
     case learnAboutPrivacyButton
+    case logOutDeviceConfirmButton
+    case logOutDeviceCancelButton
     case loginBarButton
     case loginTextFieldButton
     case logoutButton
@@ -41,6 +45,7 @@ public enum AccessibilityIdentifier: String {
     case settingsDoneButton
 
     // Cells
+    case deviceCell
     case vpnSettingsCell
     case dnsSettingsAddServerCell
     case dnsSettingsUseCustomDNSCell
@@ -70,6 +75,7 @@ public enum AccessibilityIdentifier: String {
     case alertContainerView
     case alertTitle
     case changeLogAlert
+    case deviceManagementView
     case headerBarView
     case loginView
     case outOfTimeView
diff --git a/ios/MullvadVPN/View controllers/DeviceList/DeviceManagementContentView.swift b/ios/MullvadVPN/View controllers/DeviceList/DeviceManagementContentView.swift
index 9aab160be5..5effb05beb 100644
--- a/ios/MullvadVPN/View controllers/DeviceList/DeviceManagementContentView.swift
+++ b/ios/MullvadVPN/View controllers/DeviceList/DeviceManagementContentView.swift
@@ -69,6 +69,7 @@ class DeviceManagementContentView: UIView {
             for: .normal
         )
         button.isEnabled = false
+        button.accessibilityIdentifier = .continueWithLoginButton
         return button
     }()
 
@@ -112,6 +113,8 @@ class DeviceManagementContentView: UIView {
         addViews()
         constraintViews()
         updateView()
+
+        accessibilityIdentifier = .deviceManagementView
     }
 
     private func addViews() {
diff --git a/ios/MullvadVPN/View controllers/DeviceList/DeviceManagementViewController.swift b/ios/MullvadVPN/View controllers/DeviceList/DeviceManagementViewController.swift
index c54c55836b..b36c48b2e7 100644
--- a/ios/MullvadVPN/View controllers/DeviceList/DeviceManagementViewController.swift
+++ b/ios/MullvadVPN/View controllers/DeviceList/DeviceManagementViewController.swift
@@ -212,6 +212,7 @@ class DeviceManagementViewController: UIViewController, RootContainment {
                         comment: ""
                     ),
                     style: .destructive,
+                    accessibilityId: .logOutDeviceConfirmButton,
                     handler: {
                         completion(true)
                     }
@@ -224,6 +225,7 @@ class DeviceManagementViewController: UIViewController, RootContainment {
                         comment: ""
                     ),
                     style: .default,
+                    accessibilityId: .logOutDeviceCancelButton,
                     handler: {
                         completion(false)
                     }
diff --git a/ios/MullvadVPN/View controllers/DeviceList/DeviceRowView.swift b/ios/MullvadVPN/View controllers/DeviceList/DeviceRowView.swift
index 2cb3dcb731..1fec0aeffe 100644
--- a/ios/MullvadVPN/View controllers/DeviceList/DeviceRowView.swift
+++ b/ios/MullvadVPN/View controllers/DeviceList/DeviceRowView.swift
@@ -70,6 +70,7 @@ class DeviceRowView: UIView {
 
         super.init(frame: .zero)
 
+        accessibilityIdentifier = .deviceCell
         backgroundColor = .primaryColor
         directionalLayoutMargins = UIMetrics.TableView.rowViewLayoutMargins
 
@@ -90,6 +91,7 @@ class DeviceRowView: UIView {
         )
 
         removeButton.addTarget(self, action: #selector(handleButtonTap(_:)), for: .touchUpInside)
+        removeButton.accessibilityIdentifier = .deviceCellRemoveButton
 
         NSLayoutConstraint.activate([
             textLabel.topAnchor.constraint(equalTo: layoutMarginsGuide.topAnchor),
diff --git a/ios/MullvadVPNUITests/AccountTests.swift b/ios/MullvadVPNUITests/AccountTests.swift
index 6c8d9094b9..6ca45169f3 100644
--- a/ios/MullvadVPNUITests/AccountTests.swift
+++ b/ios/MullvadVPNUITests/AccountTests.swift
@@ -15,8 +15,6 @@ class AccountTests: LoggedOutUITestCase {
         try super.setUpWithError()
     }
 
-    override func tearDownWithError() throws {}
-
     func testCreateAccount() throws {
         LoginPage(app)
             .tapCreateAccountButton()
@@ -73,6 +71,43 @@ class AccountTests: LoggedOutUITestCase {
             .waitForPageToBeShown() // Verify still on login page
     }
 
+    func testLoginToAccountWithTooManyDevices() throws {
+        // Setup
+        let temporaryAccountNumber = try MullvadAPIWrapper().createAccount()
+        try MullvadAPIWrapper().addDevices(5, account: temporaryAccountNumber)
+
+        // Teardown
+        addTeardownBlock {
+            do {
+                try MullvadAPIWrapper().deleteAccount(temporaryAccountNumber)
+            } catch {
+                XCTFail("Failed to delete account using app API")
+            }
+        }
+
+        LoginPage(app)
+            .tapAccountNumberTextField()
+            .enterText(temporaryAccountNumber)
+            .tapAccountNumberSubmitButton()
+
+        DeviceManagementPage(app)
+            .tapRemoveDeviceButton(cellIndex: 0)
+
+        DeviceManagementLogOutDeviceConfirmationAlert(app)
+            .tapYesLogOutDeviceButton()
+
+        DeviceManagementPage(app)
+            .tapContinueWithLoginButton()
+
+        // First taken back to login page and automatically being logged in
+        LoginPage(app)
+            .verifySuccessIconShown()
+            .verifyDeviceLabelShown()
+
+        // And then taken to out of time page because this account don't have any time added to it
+        OutOfTimePage(app)
+    }
+
     func testLogOut() throws {
         let newAccountNumber = try MullvadAPIWrapper().createAccount()
         login(accountNumber: newAccountNumber)
diff --git a/ios/MullvadVPNUITests/Networking/MullvadAPIWrapper.swift b/ios/MullvadVPNUITests/Networking/MullvadAPIWrapper.swift
index 25439c987a..588b692a7a 100644
--- a/ios/MullvadVPNUITests/Networking/MullvadAPIWrapper.swift
+++ b/ios/MullvadVPNUITests/Networking/MullvadAPIWrapper.swift
@@ -6,6 +6,7 @@
 //  Copyright © 2024 Mullvad VPN AB. All rights reserved.
 //
 
+import CryptoKit
 import Foundation
 import XCTest
 
@@ -48,15 +49,13 @@ class MullvadAPIWrapper {
         return port
     }
 
-    /// Generate a mock WireGuard key
+    /// Generate a mock public WireGuard key
     private func generateMockWireGuardKey() -> Data {
-        var bytes = [UInt8]()
+        let privateKey = Curve25519.KeyAgreement.PrivateKey()
+        let publicKey = privateKey.publicKey
+        let publicKeyData = publicKey.rawRepresentation
 
-        for _ in 0 ..< 44 {
-            bytes.append(UInt8.random(in: 0 ..< 255))
-        }
-
-        return Data(bytes)
+        return publicKeyData
     }
 
     func createAccount() -> String {
@@ -88,6 +87,13 @@ class MullvadAPIWrapper {
         }
     }
 
+    /// Add multiple devices to specified account. Dummy WireGuard keys will be generated.
+    func addDevices(_ numberOfDevices: Int, account: String) throws {
+        for _ in 0 ..< numberOfDevices {
+            try self.addDevice(account)
+        }
+    }
+
     func getAccountExpiry(_ account: String) throws -> Date {
         do {
             let accountExpiryTimestamp = Double(try mullvadAPI.getExpiry(forAccount: account))
diff --git a/ios/MullvadVPNUITests/Pages/DeviceManagementPage.swift b/ios/MullvadVPNUITests/Pages/DeviceManagementPage.swift
new file mode 100644
index 0000000000..571fb1cfd6
--- /dev/null
+++ b/ios/MullvadVPNUITests/Pages/DeviceManagementPage.swift
@@ -0,0 +1,49 @@
+//
+//  DeviceManagementPage.swift
+//  MullvadVPNUITests
+//
+//  Created by Niklas Berglund on 2024-03-27.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import XCTest
+
+/// Page class for the "too many devices" page shown when logging on to an account with too many devices
+class DeviceManagementPage: Page {
+    override init(_ app: XCUIApplication) {
+        super.init(app)
+
+        self.pageAccessibilityIdentifier = .deviceManagementView
+        waitForPageToBeShown()
+    }
+
+    @discardableResult func tapRemoveDeviceButton(cellIndex: Int) -> Self {
+        app
+            .otherElements.matching(identifier: AccessibilityIdentifier.deviceCell.rawValue).element(boundBy: cellIndex)
+            .buttons[AccessibilityIdentifier.deviceCellRemoveButton]
+            .tap()
+
+        return self
+    }
+
+    @discardableResult func tapContinueWithLoginButton() -> Self {
+        app.buttons[AccessibilityIdentifier.continueWithLoginButton].tap()
+        return self
+    }
+}
+
+/// Confirmation alert displayed when removing a device
+class DeviceManagementLogOutDeviceConfirmationAlert: Page {
+    override init(_ app: XCUIApplication) {
+        super.init(app)
+
+        self.pageAccessibilityIdentifier = .alertContainerView
+        waitForPageToBeShown()
+    }
+
+    @discardableResult func tapYesLogOutDeviceButton() -> Self {
+        app.buttons[AccessibilityIdentifier.logOutDeviceConfirmButton]
+            .tap()
+        return self
+    }
+}