Improved teardowns

mullvad · May 15, 2024 · 3fd09c3 · 3fd09c3
1 parent cb52f80
commit 3fd09c3
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 6 deletions.
diff --git a/ios/MullvadVPNUITests/RelayTests.swift b/ios/MullvadVPNUITests/RelayTests.swift
@@ -9,7 +9,7 @@
 import Foundation
 import XCTest
 
-fileprivate struct RelayInfo {
+private struct RelayInfo {
     var name: String
     var ipAddress: String
 }
@@ -93,6 +93,17 @@ class RelayTests: LoggedInWithTimeUITestCase {
         FirewallAPIClient().removeRules()
         removeFirewallRulesInTearDown = true
 
+        addTeardownBlock {
+            TunnelControlPage(self.app)
+                .tapSelectLocationButton()
+
+            SelectLocationPage(self.app)
+                .tapLocationCell(withName: "Sweden")
+
+            TunnelControlPage(self.app)
+                .tapCancelButton()
+        }
+
         // First get relay IP address
         let relayInfo = getGothenburgRelayInfo()
 
@@ -160,6 +171,17 @@ class RelayTests: LoggedInWithTimeUITestCase {
         FirewallAPIClient().removeRules()
         removeFirewallRulesInTearDown = true
 
+        addTeardownBlock {
+            TunnelControlPage(self.app)
+                .tapSelectLocationButton()
+
+            SelectLocationPage(self.app)
+                .tapLocationCell(withName: "Sweden")
+
+            TunnelControlPage(self.app)
+                .tapCancelButton()
+        }
+
         // First get relay IP address
         let relayInfo = getGothenburgRelayInfo()
 

diff --git a/ios/MullvadVPNUITests/Test base classes/LoggedInWithTimeUITestCase.swift b/ios/MullvadVPNUITests/Test base classes/LoggedInWithTimeUITestCase.swift
@@ -7,26 +7,41 @@
 //
 
 import Foundation
+import XCTest
 
 /// Base class for tests that should start from a state of being logged on to an account with time left
 class LoggedInWithTimeUITestCase: BaseUITestCase {
+    var hasTimeAccountNumber: String?
+
     override func setUp() {
         super.setUp()
 
-        let hasTimeAccountNumber = getAccountWithTime()
-
-        addTeardownBlock {
-            self.returnAccountWithTime(accountNumber: hasTimeAccountNumber)
-        }
+        hasTimeAccountNumber = getAccountWithTime()
 
         agreeToTermsOfServiceIfShown()
         dismissChangeLogIfShown()
         logoutIfLoggedIn()
 
+        guard let hasTimeAccountNumber = self.hasTimeAccountNumber else {
+            XCTFail("hasTimeAccountNumber unexpectedly not set")
+            return
+        }
+
         login(accountNumber: hasTimeAccountNumber)
 
         // Relaunch app so that tests start from a deterministic state
         app.terminate()
         app.launch()
     }
+
+    override func tearDown() {
+        guard let hasTimeAccountNumber = self.hasTimeAccountNumber else {
+            XCTFail("hasTimeAccountNumber unexpectedly not set")
+            return
+        }
+
+        app.terminate() // Terminate app to make sure we have Internet connectivity
+        self.returnAccountWithTime(accountNumber: hasTimeAccountNumber)
+        app.launch()
+    }
 }