Add connection retry logic test

8 files changed, 129 insertions, 28 deletions

diff --git a/ios/MullvadVPN/Classes/AccessbilityIdentifier.swift b/ios/MullvadVPN/Classes/AccessbilityIdentifier.swift
index 25d7d2d270..371412c975 100644
--- a/ios/MullvadVPN/Classes/AccessbilityIdentifier.swift
+++ b/ios/MullvadVPN/Classes/AccessbilityIdentifier.swift
@@ -17,6 +17,7 @@ public enum AccessibilityIdentifier: String {
     case cancelButton
     case connectionPanelButton
     case collapseButton
+    case expandButton
     case createAccountButton
     case deleteButton
     case disconnectButton
@@ -57,7 +58,8 @@ public enum AccessibilityIdentifier: String {
 
     // Labels
     case headerDeviceNameLabel
-    case connectionStatusLabel
+    case connectionStatusConnectedLabel
+    case connectionStatusNotConnectedLabel
     case welcomeAccountNumberLabel
     case connectionPanelDetailLabel
 
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift
index 24a1ce6fac..d68417f4c3 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift
@@ -49,7 +49,6 @@ class LocationCell: UITableViewCell {
 
     private let collapseButton: UIButton = {
         let button = UIButton(type: .custom)
-        button.accessibilityIdentifier = .collapseButton
         button.isAccessibilityElement = false
         button.tintColor = .white
         return button
@@ -267,6 +266,7 @@ class LocationCell: UITableViewCell {
     private func updateCollapseImage() {
         let image = isExpanded ? chevronUp : chevronDown
 
+        collapseButton.accessibilityIdentifier = isExpanded ? .collapseButton : .expandButton
         collapseButton.setImage(image, for: .normal)
     }
 
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationViewController.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationViewController.swift
index 6b27418aa5..11b5c0bbb7 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationViewController.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationViewController.swift
@@ -55,6 +55,7 @@ final class LocationViewController: UIViewController {
     override func viewDidLoad() {
         super.viewDidLoad()
 
+        view.accessibilityIdentifier = .selectLocationView
         view.backgroundColor = .secondaryColor
 
         navigationItem.title = NSLocalizedString(
diff --git a/ios/MullvadVPN/View controllers/Tunnel/TunnelControlView.swift b/ios/MullvadVPN/View controllers/Tunnel/TunnelControlView.swift
index 1472515927..8635bc1b9d 100644
--- a/ios/MullvadVPN/View controllers/Tunnel/TunnelControlView.swift
+++ b/ios/MullvadVPN/View controllers/Tunnel/TunnelControlView.swift
@@ -121,8 +121,6 @@ final class TunnelControlView: UIView {
         accessibilityContainerType = .semanticGroup
         accessibilityIdentifier = .tunnelControlView
 
-        secureLabel.accessibilityIdentifier = .connectionStatusLabel
-
         addSubviews()
         addButtonHandlers()
     }
@@ -187,6 +185,13 @@ final class TunnelControlView: UIView {
     private func updateSecureLabel(tunnelState: TunnelState) {
         secureLabel.text = tunnelState.localizedTitleForSecureLabel.uppercased()
         secureLabel.textColor = tunnelState.textColorForSecureLabel
+
+        switch tunnelState {
+        case .connected:
+            secureLabel.accessibilityIdentifier = .connectionStatusConnectedLabel
+        default:
+            secureLabel.accessibilityIdentifier = .connectionStatusNotConnectedLabel
+        }
     }
 
     private func updateButtonTitles(tunnelState: TunnelState) {
diff --git a/ios/MullvadVPNUITests/Networking/FirewallRule.swift b/ios/MullvadVPNUITests/Networking/FirewallRule.swift
index 1b18baa91c..84f9a56820 100644
--- a/ios/MullvadVPNUITests/Networking/FirewallRule.swift
+++ b/ios/MullvadVPNUITests/Networking/FirewallRule.swift
@@ -45,6 +45,16 @@ struct FirewallRule {
         )
     }
 
+    public static func makeBlockAllTrafficRule(toIPAddress: String) throws -> FirewallRule {
+        let deviceIPAddress = try Networking.getIPAddress()
+
+        return FirewallRule(
+            fromIPAddress: deviceIPAddress,
+            toIPAddress: toIPAddress,
+            protocols: [.ICMP, .TCP, .UDP]
+        )
+    }
+
     public static func makeBlockUDPTrafficRule(toIPAddress: String) throws -> FirewallRule {
         let deviceIPAddress = try Networking.getIPAddress()
 
diff --git a/ios/MullvadVPNUITests/Pages/SelectLocationPage.swift b/ios/MullvadVPNUITests/Pages/SelectLocationPage.swift
index ba15a7323f..ad0ee06c4b 100644
--- a/ios/MullvadVPNUITests/Pages/SelectLocationPage.swift
+++ b/ios/MullvadVPNUITests/Pages/SelectLocationPage.swift
@@ -22,14 +22,19 @@ class SelectLocationPage: Page {
         return self
     }
 
-    @discardableResult func tapLocationCellExpandCollapseButton(withName name: String) -> Self {
+    @discardableResult func tapLocationCellExpandButton(withName name: String) -> Self {
         let table = app.tables[AccessibilityIdentifier.selectLocationTableView]
         let matchingCells = table.cells.containing(.any, identifier: name)
         let buttons = matchingCells.buttons
-        let expandButton = buttons[AccessibilityIdentifier.collapseButton]
+        let expandButton = buttons[AccessibilityIdentifier.expandButton]
 
         expandButton.tap()
 
         return self
     }
+
+    func locationCellIsExpanded(_ name: String) -> Bool {
+        let matchingCells = app.cells.containing(.any, identifier: name)
+        return matchingCells.buttons[AccessibilityIdentifier.expandButton].exists ? false : true
+    }
 }
diff --git a/ios/MullvadVPNUITests/Pages/TunnelControlPage.swift b/ios/MullvadVPNUITests/Pages/TunnelControlPage.swift
index daef1a30e7..4a2eeee4c7 100644
--- a/ios/MullvadVPNUITests/Pages/TunnelControlPage.swift
+++ b/ios/MullvadVPNUITests/Pages/TunnelControlPage.swift
@@ -19,9 +19,10 @@ class TunnelControlPage: Page {
     /// Poll the "in address row" label for its updated values and output an array of ConnectionAttempt objects representing the connection attempts that have been communicated through the UI.
     /// - Parameters:
     ///   - attemptsCount: number of connection attempts to look for
-    ///   - timeout: timeout after this many seconds if attemptsCount haven't been reached yet
+    ///   - timeout: return the attemps found so far after this many seconds if `attemptsCount` haven't been reached yet
     private func waitForConnectionAttempts(_ attemptsCount: Int, timeout: TimeInterval) -> [ConnectionAttempt] {
         var connectionAttempts: [ConnectionAttempt] = []
+        var lastConnectionAttempt: ConnectionAttempt?
         let startTime = Date()
         let pollingInterval = TimeInterval(0.5) // How often to check for changes
 
@@ -52,8 +53,9 @@ class TunnelControlPage: Page {
                     protocolName: protocolName
                 )
 
-                if connectionAttempts.contains(connectionAttempt) == false {
+                if connectionAttempt != lastConnectionAttempt {
                     connectionAttempts.append(connectionAttempt)
+                    lastConnectionAttempt = connectionAttempt
 
                     if connectionAttempts.count == attemptsCount {
                         break
@@ -87,8 +89,13 @@ class TunnelControlPage: Page {
         return self
     }
 
+    @discardableResult func tapCancelButton() -> Self {
+        app.buttons[AccessibilityIdentifier.cancelButton].tap()
+        return self
+    }
+
     @discardableResult func waitForSecureConnectionLabel() -> Self {
-        _ = app.staticTexts[AccessibilityIdentifier.connectionStatusLabel]
+        _ = app.staticTexts[AccessibilityIdentifier.connectionStatusConnectedLabel]
             .waitForExistence(timeout: BaseUITestCase.defaultTimeout)
         return self
     }
@@ -130,6 +137,43 @@ class TunnelControlPage: Page {
         return self
     }
 
+    /// Verify that connection attempts are made in the correct order
+    @discardableResult func verifyConnectionAttemptsOrder() -> Self {
+        let connectionAttempts = waitForConnectionAttempts(4, timeout: 50)
+        XCTAssertEqual(connectionAttempts.count, 4)
+
+        if connectionAttempts.last?.protocolName == "UDP" {
+            // If last attempt is over UDP it means we have encountered the UI bug where only one UDP attempt is shown and then the two TCP attempts
+            for (attemptIndex, attempt) in connectionAttempts.enumerated() {
+                if attemptIndex == 0 {
+                    XCTAssertEqual(attempt.protocolName, "UDP")
+                } else if attemptIndex == 1 {
+                    XCTAssertEqual(attempt.protocolName, "TCP")
+                    XCTAssertEqual(attempt.port, "80")
+                } else if attemptIndex == 2 {
+                    XCTAssertEqual(attempt.protocolName, "TCP")
+                    XCTAssertEqual(attempt.port, "5001")
+                } // Ignore the 4th attempt which is the first attempt of new attempt cycle
+            }
+        } else {
+            for (attemptIndex, attempt) in connectionAttempts.enumerated() {
+                if attemptIndex == 0 {
+                    XCTAssertEqual(attempt.protocolName, "UDP")
+                } else if attemptIndex == 1 {
+                    XCTAssertEqual(attempt.protocolName, "UDP")
+                } else if attemptIndex == 2 {
+                    XCTAssertEqual(attempt.protocolName, "TCP")
+                    XCTAssertEqual(attempt.port, "80")
+                } else if attemptIndex == 3 {
+                    XCTAssertEqual(attempt.protocolName, "TCP")
+                    XCTAssertEqual(attempt.port, "5001")
+                }
+            }
+        }
+
+        return self
+    }
+
     @discardableResult func verifyConnectingToPort(_ port: String) -> Self {
         let connectionAttempts = waitForConnectionAttempts(1, timeout: 10)
         XCTAssertEqual(connectionAttempts.count, 1)
diff --git a/ios/MullvadVPNUITests/RelayTests.swift b/ios/MullvadVPNUITests/RelayTests.swift
index 04df1e3454..5df5b05a06 100644
--- a/ios/MullvadVPNUITests/RelayTests.swift
+++ b/ios/MullvadVPNUITests/RelayTests.swift
@@ -55,6 +55,27 @@ class RelayTests: LoggedInWithTimeUITestCase {
             .tapDisconnectButton()
     }
 
+    func testConnectionRetryLogic() throws {
+        FirewallAPIClient().removeRules()
+        removeFirewallRulesInTearDown = true
+
+        // First get relay IP address
+        let relayIPAddress = getGot001WireGuardRelayIPAddress()
+
+        // Run actual test
+        try FirewallAPIClient().createRule(
+            FirewallRule.makeBlockAllTrafficRule(toIPAddress: relayIPAddress)
+        )
+
+        TunnelControlPage(app)
+            .tapSecureConnectionButton()
+
+        // Should be two UDP connection attempts but sometimes only one is shown in the UI
+        TunnelControlPage(app)
+            .verifyConnectionAttemptsOrder()
+            .tapCancelButton()
+    }
+
     func testWireGuardOverTCPManually() throws {
         HeaderBar(app)
             .tapSettingsButton()
@@ -86,29 +107,11 @@ class RelayTests: LoggedInWithTimeUITestCase {
 
     /// Test automatic switching to TCP is functioning when UDP traffic to relay is blocked. This test first connects to a realy to get the IP address of it, in order to block UDP traffic to this relay.
     func testWireGuardOverTCPAutomatically() throws {
-        let wireGuardGot001RelayName = "se-got-wg-001"
-
         FirewallAPIClient().removeRules()
         removeFirewallRulesInTearDown = true
 
         // First get relay IP address
-        TunnelControlPage(app)
-            .tapSelectLocationButton()
-
-        SelectLocationPage(app)
-            .tapLocationCellExpandCollapseButton(withName: "Sweden")
-            .tapLocationCellExpandCollapseButton(withName: "Gothenburg")
-            .tapLocationCell(withName: wireGuardGot001RelayName)
-
-        allowAddVPNConfigurationsIfAsked()
-
-        let relayIPAddress = TunnelControlPage(app)
-            .waitForSecureConnectionLabel()
-            .tapRelayStatusExpandCollapseButton()
-            .getInIPAddressFromConnectionStatus()
-
-        TunnelControlPage(app)
-            .tapDisconnectButton()
+        let relayIPAddress = getGot001WireGuardRelayIPAddress()
 
         // Run actual test
         try FirewallAPIClient().createRule(
@@ -163,4 +166,35 @@ class RelayTests: LoggedInWithTimeUITestCase {
             .verifyConnectingToPort("4001")
             .tapDisconnectButton()
     }
+
+    /// Get got001 WireGuard relay IP address by connecting to it and checking which IP address the app connects to. Assumes user is logged on and at tunnel control page.
+    private func getGot001WireGuardRelayIPAddress() -> String {
+        let wireGuardGot001RelayName = "se-got-wg-001"
+
+        TunnelControlPage(app)
+            .tapSelectLocationButton()
+
+        if SelectLocationPage(app).locationCellIsExpanded("Sweden") {
+            // Already expanded - just make sure correct relay is selected
+            SelectLocationPage(app)
+                .tapLocationCell(withName: wireGuardGot001RelayName)
+        } else {
+            SelectLocationPage(app)
+                .tapLocationCellExpandButton(withName: "Sweden")
+                .tapLocationCellExpandButton(withName: "Gothenburg")
+                .tapLocationCell(withName: wireGuardGot001RelayName)
+        }
+
+        allowAddVPNConfigurationsIfAsked()
+
+        let relayIPAddress = TunnelControlPage(app)
+            .waitForSecureConnectionLabel()
+            .tapRelayStatusExpandCollapseButton()
+            .getInIPAddressFromConnectionStatus()
+
+        TunnelControlPage(app)
+            .tapDisconnectButton()
+
+        return relayIPAddress
+    }
 }