Migrate ConnectionTest e2e tests to use POP

13 files changed, 371 insertions, 101 deletions

diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/component/TopBar.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/component/TopBar.kt
index 5bdb8e71ca..91cb49ae53 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/component/TopBar.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/component/TopBar.kt
@@ -49,6 +49,7 @@ import androidx.compose.ui.unit.dp
 import net.mullvad.mullvadvpn.R
 import net.mullvad.mullvadvpn.compose.test.TOP_BAR_ACCOUNT_BUTTON
 import net.mullvad.mullvadvpn.compose.test.TOP_BAR_SETTINGS_BUTTON
+import net.mullvad.mullvadvpn.compose.test.TOP_BAR_TEST_TAG
 import net.mullvad.mullvadvpn.lib.theme.AppTheme
 import net.mullvad.mullvadvpn.lib.theme.Dimens
 
@@ -117,7 +118,7 @@ fun MullvadTopBar(
     isIconAndLogoVisible: Boolean = true,
 ) {
     TopAppBar(
-        modifier = modifier,
+        modifier = modifier.testTag(TOP_BAR_TEST_TAG),
         title = {
             if (isIconAndLogoVisible) {
                 Row(verticalAlignment = Alignment.CenterVertically) {
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/SettingsScreen.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/SettingsScreen.kt
index 75ba5abdd8..15754990f2 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/SettingsScreen.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/SettingsScreen.kt
@@ -43,6 +43,7 @@ import net.mullvad.mullvadvpn.compose.extensions.itemWithDivider
 import net.mullvad.mullvadvpn.compose.preview.SettingsUiStatePreviewParameterProvider
 import net.mullvad.mullvadvpn.compose.state.SettingsUiState
 import net.mullvad.mullvadvpn.compose.test.LAZY_LIST_TEST_TAG
+import net.mullvad.mullvadvpn.compose.test.VPN_SETTINGS_CELL_TEST_TAG
 import net.mullvad.mullvadvpn.compose.transitions.TopLevelTransition
 import net.mullvad.mullvadvpn.lib.theme.AppTheme
 import net.mullvad.mullvadvpn.lib.theme.Dimens
@@ -115,6 +116,7 @@ fun SettingsScreen(
                     NavigationComposeCell(
                         title = stringResource(id = R.string.settings_vpn),
                         onClick = onVpnSettingCellClick,
+                        testTag = VPN_SETTINGS_CELL_TEST_TAG,
                     )
                 }
                 item { Spacer(modifier = Modifier.height(Dimens.cellVerticalSpacing)) }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/location/SelectLocationScreen.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/location/SelectLocationScreen.kt
index d6d4721f20..36a3ed2eee 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/location/SelectLocationScreen.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/location/SelectLocationScreen.kt
@@ -32,6 +32,7 @@ import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.platform.LocalContext
+import androidx.compose.ui.platform.testTag
 import androidx.compose.ui.res.stringResource
 import androidx.compose.ui.tooling.preview.Preview
 import androidx.compose.ui.tooling.preview.PreviewParameter
@@ -63,6 +64,7 @@ import net.mullvad.mullvadvpn.compose.extensions.dropUnlessResumed
 import net.mullvad.mullvadvpn.compose.preview.SelectLocationsUiStatePreviewParameterProvider
 import net.mullvad.mullvadvpn.compose.state.RelayListType
 import net.mullvad.mullvadvpn.compose.state.SelectLocationUiState
+import net.mullvad.mullvadvpn.compose.test.SELECT_LOCATION_SCREEN_TEST_TAG
 import net.mullvad.mullvadvpn.compose.transitions.TopLevelTransition
 import net.mullvad.mullvadvpn.compose.util.CollectSideEffectWithLifecycle
 import net.mullvad.mullvadvpn.compose.util.showSnackbarImmediately
@@ -237,6 +239,7 @@ fun SelectLocationScreen(
                 )
             }
         },
+        modifier = Modifier.testTag(SELECT_LOCATION_SCREEN_TEST_TAG),
         snackbarHostState = snackbarHostState,
         actions = {
             IconButton(
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/test/ComposeTestTagConstants.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/test/ComposeTestTagConstants.kt
index b124ffcc61..eef89c5ea2 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/test/ComposeTestTagConstants.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/test/ComposeTestTagConstants.kt
@@ -3,6 +3,10 @@ package net.mullvad.mullvadvpn.compose.test
 // Top Bar
 const val TOP_BAR_ACCOUNT_BUTTON = "top_bar_account_button"
 const val TOP_BAR_SETTINGS_BUTTON = "top_bar_settings_button"
+const val TOP_BAR_TEST_TAG = "top_bar_test_tag"
+
+// Settings screen
+const val VPN_SETTINGS_CELL_TEST_TAG = "vpn_settings_cell_test_tag"
 
 // VpnSettingsScreen
 const val LAZY_LIST_VPN_SETTINGS_TEST_TAG = "lazy_list_vpn_settings_test_tag"
@@ -24,6 +28,7 @@ const val WIREGUARD_OBFUSCATION_UDP_OVER_TCP_CELL =
     "wireguard_obfuscation_udp_over_tcp_cell_test_tag"
 
 // SelectLocationScreen, ConnectScreen, CustomListLocationsScreen
+const val SELECT_LOCATION_SCREEN_TEST_TAG = "select_location_screen_test_tag"
 const val CIRCULAR_PROGRESS_INDICATOR = "circular_progress_indicator"
 const val EXPAND_BUTTON_TEST_TAG = "expand_button_test_tag"
 const val LOCATION_CELL_TEST_TAG = "location_cell_test_tag"
diff --git a/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/ConnectPage.kt b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/ConnectPage.kt
index fe1fafcc7f..320c01d7aa 100644
--- a/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/ConnectPage.kt
+++ b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/ConnectPage.kt
@@ -1,10 +1,79 @@
 package net.mullvad.mullvadvpn.test.common.page
 
 import androidx.test.uiautomator.By
+import net.mullvad.mullvadvpn.test.common.constant.VERY_LONG_TIMEOUT
 import net.mullvad.mullvadvpn.test.common.extension.findObjectWithTimeout
 
 class ConnectPage internal constructor() : Page() {
+    private val disconnectSelector = By.text("Disconnect")
+    private val cancelSelector = By.text("Cancel")
+    private val connectedSelector = By.text("CONNECTED")
+    private val connectingSelector = By.text("CONNECTING...")
+
     override fun assertIsDisplayed() {
-        uiDevice.findObjectWithTimeout(By.res("connect_card_header_test_tag"))
+        uiDevice.findObjectWithTimeout(By.res(CONNECT_CARD_HEADER_TEST_TAG))
+    }
+
+    fun clickSelectLocation() {
+        uiDevice.findObjectWithTimeout(By.res(SELECT_LOCATION_BUTTON_TEST_TAG)).click()
+    }
+
+    fun clickConnect() {
+        uiDevice.findObjectWithTimeout(By.res(CONNECT_BUTTON_TEST_TAG)).click()
+    }
+
+    fun clickDisconnect() {
+        uiDevice.findObjectWithTimeout(disconnectSelector).click()
+    }
+
+    fun clickCancel() {
+        uiDevice.findObjectWithTimeout(cancelSelector).click()
+    }
+
+    fun waitForConnectedLabel(timeout: Long = VERY_LONG_TIMEOUT) {
+        uiDevice.findObjectWithTimeout(connectedSelector, timeout)
+    }
+
+    fun waitForConnectingLabel() {
+        uiDevice.findObjectWithTimeout(connectingSelector)
+    }
+
+    /**
+     * Extracts the in IPv4 address from the connection card. It is a prerequisite that the
+     * connection card is in collapsed state.
+     */
+    fun extractInIpv4Address(): String {
+        uiDevice.findObjectWithTimeout(By.res("connect_card_header_test_tag")).click()
+        val inString =
+            uiDevice
+                .findObjectWithTimeout(
+                    By.res("location_info_connection_in_test_tag"),
+                    VERY_LONG_TIMEOUT,
+                )
+                .text
+
+        val extractedIpAddress = inString.split(" ")[0].split(":")[0]
+        return extractedIpAddress
+    }
+
+    /**
+     * Extracts the out IPv4 address from the connection card. It is a prerequisite that the
+     * connection card is in collapsed state.
+     */
+    fun extractOutIpv4Address(): String {
+        uiDevice.findObjectWithTimeout(By.res("connect_card_header_test_tag")).click()
+        return uiDevice
+            .findObjectWithTimeout(
+                // Text exist and contains IP address
+                By.res("location_info_connection_out_test_tag").textContains("."),
+                VERY_LONG_TIMEOUT,
+            )
+            .text
+    }
+
+    companion object {
+        const val CONNECT_CARD_HEADER_TEST_TAG = "connect_card_header_test_tag"
+        const val SELECT_LOCATION_BUTTON_TEST_TAG = "select_location_button_test_tag"
+        const val CONNECT_BUTTON_TEST_TAG = "connect_button_test_tag"
     }
 }
diff --git a/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/LoginPage.kt b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/LoginPage.kt
index 9fdffe1eae..1098e4d1cc 100644
--- a/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/LoginPage.kt
+++ b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/LoginPage.kt
@@ -8,6 +8,9 @@ import net.mullvad.mullvadvpn.test.common.constant.EXTREMELY_LONG_TIMEOUT
 import net.mullvad.mullvadvpn.test.common.extension.findObjectWithTimeout
 
 class LoginPage internal constructor() : Page() {
+    private val invalidAccountNumberSelector = By.text("Invalid account number")
+    private val loginSelector = By.text("Login")
+
     fun enterAccountNumber(accountNumber: String) {
         uiDevice.findObjectWithTimeout(By.clazz("android.widget.EditText")).text = accountNumber
     }
@@ -20,10 +23,10 @@ class LoginPage internal constructor() : Page() {
     }
 
     fun verifyShowingInvalidAccount() {
-        uiDevice.findObjectWithTimeout(By.text("Invalid account number"), EXTREMELY_LONG_TIMEOUT)
+        uiDevice.findObjectWithTimeout(invalidAccountNumberSelector, EXTREMELY_LONG_TIMEOUT)
     }
 
     override fun assertIsDisplayed() {
-        uiDevice.findObjectWithTimeout(By.text("Login"))
+        uiDevice.findObjectWithTimeout(loginSelector)
     }
 }
diff --git a/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/PrivacyPage.kt b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/PrivacyPage.kt
index 43ec183c50..f5ca662113 100644
--- a/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/PrivacyPage.kt
+++ b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/PrivacyPage.kt
@@ -7,12 +7,16 @@ import net.mullvad.mullvadvpn.test.common.constant.DEFAULT_TIMEOUT
 import net.mullvad.mullvadvpn.test.common.extension.findObjectWithTimeout
 
 class PrivacyPage internal constructor() : Page() {
+    private val privacySelector = By.text("Privacy")
+    private val agreeSelector = By.text("Agree and continue")
+    private val allowSelector = By.text("Allow")
+
     override fun assertIsDisplayed() {
-        uiDevice.findObjectWithTimeout(By.text("Privacy"))
+        uiDevice.findObjectWithTimeout(privacySelector)
     }
 
     fun clickAgreeOnPrivacyDisclaimer() {
-        uiDevice.findObjectWithTimeout(By.text("Agree and continue")).click()
+        uiDevice.findObjectWithTimeout(agreeSelector).click()
     }
 
     fun clickAllowOnNotificationPermissionPromptIfApiLevel33AndAbove(
@@ -23,12 +27,10 @@ class PrivacyPage internal constructor() : Page() {
             return
         }
 
-        val selector = By.text("Allow")
-
-        uiDevice.wait(Until.hasObject(selector), timeout)
+        uiDevice.wait(Until.hasObject(allowSelector), timeout)
 
         try {
-            uiDevice.findObjectWithTimeout(selector).click()
+            uiDevice.findObjectWithTimeout(allowSelector).click()
         } catch (e: IllegalArgumentException) {
             throw IllegalArgumentException(
                 "Failed to allow notification permission within timeout ($timeout)"
diff --git a/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/SelectLocationPage.kt b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/SelectLocationPage.kt
new file mode 100644
index 0000000000..e10a53ec4d
--- /dev/null
+++ b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/SelectLocationPage.kt
@@ -0,0 +1,25 @@
+package net.mullvad.mullvadvpn.test.common.page
+
+import androidx.test.uiautomator.By
+import net.mullvad.mullvadvpn.test.common.extension.findObjectWithTimeout
+
+class SelectLocationPage internal constructor() : Page() {
+    override fun assertIsDisplayed() {
+        uiDevice.findObjectWithTimeout(By.res(SELECT_LOCATION_SCREEN_TEST_TAG))
+    }
+
+    fun clickLocationExpandButton(locationName: String) {
+        val locationCell = uiDevice.findObjectWithTimeout(By.text(locationName)).parent.parent
+        val expandButton = locationCell.findObjectWithTimeout(By.res(EXPAND_BUTTON_TEST_TAG))
+        expandButton.click()
+    }
+
+    fun clickLocationCell(locationName: String) {
+        uiDevice.findObjectWithTimeout(By.text(locationName)).click()
+    }
+
+    companion object {
+        const val SELECT_LOCATION_SCREEN_TEST_TAG = "select_location_screen_test_tag"
+        const val EXPAND_BUTTON_TEST_TAG = "expand_button_test_tag"
+    }
+}
diff --git a/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/SettingsPage.kt b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/SettingsPage.kt
new file mode 100644
index 0000000000..86a317d153
--- /dev/null
+++ b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/SettingsPage.kt
@@ -0,0 +1,20 @@
+package net.mullvad.mullvadvpn.test.common.page
+
+import androidx.test.uiautomator.By
+import net.mullvad.mullvadvpn.test.common.extension.findObjectWithTimeout
+
+class SettingsPage internal constructor() : Page() {
+    private val settingsSelector = By.text("Settings")
+
+    override fun assertIsDisplayed() {
+        uiDevice.findObjectWithTimeout(settingsSelector)
+    }
+
+    fun clickVpnSettings() {
+        uiDevice.findObjectWithTimeout(By.res(VPN_SETTINGS_CELL_TEST_TAG)).click()
+    }
+
+    companion object {
+        const val VPN_SETTINGS_CELL_TEST_TAG = "vpn_settings_cell_test_tag"
+    }
+}
diff --git a/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/SystemVpnConfigurationAlert.kt b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/SystemVpnConfigurationAlert.kt
new file mode 100644
index 0000000000..046d6e8197
--- /dev/null
+++ b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/SystemVpnConfigurationAlert.kt
@@ -0,0 +1,16 @@
+package net.mullvad.mullvadvpn.test.common.page
+
+import androidx.test.uiautomator.By
+import net.mullvad.mullvadvpn.test.common.extension.findObjectWithTimeout
+
+class SystemVpnConfigurationAlert internal constructor() : Page() {
+    private val okSelector = By.text("OK")
+
+    override fun assertIsDisplayed() {
+        uiDevice.findObjectWithTimeout(okSelector)
+    }
+
+    fun clickOk() {
+        uiDevice.findObjectWithTimeout(okSelector).click()
+    }
+}
diff --git a/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/TopBar.kt b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/TopBar.kt
new file mode 100644
index 0000000000..c602e26def
--- /dev/null
+++ b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/TopBar.kt
@@ -0,0 +1,24 @@
+package net.mullvad.mullvadvpn.test.common.page
+
+import androidx.test.uiautomator.By
+import net.mullvad.mullvadvpn.test.common.extension.findObjectWithTimeout
+
+class TopBar internal constructor() : Page() {
+    override fun assertIsDisplayed() {
+        uiDevice.findObjectWithTimeout(By.res(TOP_BAR_TEST_TAG))
+    }
+
+    fun clickSettings() {
+        uiDevice.findObjectWithTimeout(By.res(TOP_BAR_SETTINGS_BUTTON)).click()
+    }
+
+    fun clickAccount() {
+        uiDevice.findObjectWithTimeout(By.res(TOP_BAR_ACCOUNT_BUTTON)).click()
+    }
+
+    companion object {
+        const val TOP_BAR_TEST_TAG = "top_bar_test_tag"
+        const val TOP_BAR_ACCOUNT_BUTTON = "top_bar_account_button"
+        const val TOP_BAR_SETTINGS_BUTTON = "top_bar_settings_button"
+    }
+}
diff --git a/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/VpnSettingsPage.kt b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/VpnSettingsPage.kt
new file mode 100644
index 0000000000..698c84aa6b
--- /dev/null
+++ b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/VpnSettingsPage.kt
@@ -0,0 +1,56 @@
+package net.mullvad.mullvadvpn.test.common.page
+
+import androidx.test.uiautomator.By
+import androidx.test.uiautomator.Direction
+import androidx.test.uiautomator.Until
+import net.mullvad.mullvadvpn.test.common.extension.findObjectWithTimeout
+
+class VpnSettingsPage internal constructor() : Page() {
+    private val vpnSettingsSelector = By.text("VPN settings")
+    private val localNetworkSharingSelector = By.text("Local network sharing")
+
+    override fun assertIsDisplayed() {
+        uiDevice.findObjectWithTimeout(vpnSettingsSelector)
+    }
+
+    fun clickLocalNetworkSharingSwitch() {
+        val localNetworkSharingCell =
+            uiDevice.findObjectWithTimeout(localNetworkSharingSelector).parent
+        val localNetworkSharingSwitch =
+            localNetworkSharingCell.findObjectWithTimeout(By.res(SWITCH_TEST_TAG))
+
+        localNetworkSharingSwitch.click()
+    }
+
+    fun scrollUntilWireguardObfuscationUdpOverTcpCell() {
+        scrollUntilCell(WIREGUARD_OBFUSCATION_UDP_OVER_TCP_CELL_TEST_TAG)
+    }
+
+    fun scrollUntilWireguardObfuscationOffCell() {
+        scrollUntilCell(WIREGUARD_OBFUSCATION_OFF_CELL_TEST_TAG)
+    }
+
+    fun clickWireguardObfuscationUdpOverTcpCell() {
+        uiDevice
+            .findObjectWithTimeout(By.res(WIREGUARD_OBFUSCATION_UDP_OVER_TCP_CELL_TEST_TAG))
+            .click()
+    }
+
+    fun clickWireguardObfuscationOffCell() {
+        uiDevice.findObjectWithTimeout(By.res(WIREGUARD_OBFUSCATION_OFF_CELL_TEST_TAG)).click()
+    }
+
+    private fun scrollUntilCell(testTag: String) {
+        val scrollView2 = uiDevice.findObjectWithTimeout(By.res(SETTINGS_SCROLL_VIEW_TEST_TAG))
+        scrollView2.scrollUntil(Direction.DOWN, Until.hasObject(By.res(testTag)))
+    }
+
+    companion object {
+        const val SETTINGS_SCROLL_VIEW_TEST_TAG = "lazy_list_vpn_settings_test_tag"
+        const val WIREGUARD_OBFUSCATION_UDP_OVER_TCP_CELL_TEST_TAG =
+            "wireguard_obfuscation_udp_over_tcp_cell_test_tag"
+        const val WIREGUARD_OBFUSCATION_OFF_CELL_TEST_TAG =
+            "wireguard_obfuscation_off_cell_test_tag"
+        const val SWITCH_TEST_TAG = "switch_test_tag"
+    }
+}
diff --git a/android/test/e2e/src/main/kotlin/net/mullvad/mullvadvpn/test/e2e/ConnectionTest.kt b/android/test/e2e/src/main/kotlin/net/mullvad/mullvadvpn/test/e2e/ConnectionTest.kt
index 81a9eabfe6..1ba78818a2 100644
--- a/android/test/e2e/src/main/kotlin/net/mullvad/mullvadvpn/test/e2e/ConnectionTest.kt
+++ b/android/test/e2e/src/main/kotlin/net/mullvad/mullvadvpn/test/e2e/ConnectionTest.kt
@@ -1,19 +1,17 @@
 package net.mullvad.mullvadvpn.test.e2e
 
-import androidx.test.uiautomator.By
-import androidx.test.uiautomator.Direction
-import androidx.test.uiautomator.Until
 import kotlin.time.Duration.Companion.milliseconds
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.runBlocking
 import net.mullvad.mullvadvpn.BuildConfig
-import net.mullvad.mullvadvpn.compose.test.EXPAND_BUTTON_TEST_TAG
-import net.mullvad.mullvadvpn.compose.test.SELECT_LOCATION_BUTTON_TEST_TAG
-import net.mullvad.mullvadvpn.compose.test.SWITCH_TEST_TAG
-import net.mullvad.mullvadvpn.compose.test.TOP_BAR_SETTINGS_BUTTON
 import net.mullvad.mullvadvpn.test.common.constant.EXTREMELY_LONG_TIMEOUT
-import net.mullvad.mullvadvpn.test.common.constant.VERY_LONG_TIMEOUT
-import net.mullvad.mullvadvpn.test.common.extension.findObjectWithTimeout
+import net.mullvad.mullvadvpn.test.common.page.ConnectPage
+import net.mullvad.mullvadvpn.test.common.page.SelectLocationPage
+import net.mullvad.mullvadvpn.test.common.page.SettingsPage
+import net.mullvad.mullvadvpn.test.common.page.SystemVpnConfigurationAlert
+import net.mullvad.mullvadvpn.test.common.page.TopBar
+import net.mullvad.mullvadvpn.test.common.page.VpnSettingsPage
+import net.mullvad.mullvadvpn.test.common.page.on
 import net.mullvad.mullvadvpn.test.common.rule.ForgetAllVpnAppsInSettingsTestRule
 import net.mullvad.mullvadvpn.test.e2e.annotations.HasDependencyOnLocalAPI
 import net.mullvad.mullvadvpn.test.e2e.misc.AccountTestRule
@@ -34,19 +32,18 @@ class ConnectionTest : EndToEndTest(BuildConfig.FLAVOR_infrastructure) {
     @JvmField
     val forgetAllVpnAppsInSettingsTestRule = ForgetAllVpnAppsInSettingsTestRule()
 
-    val firewallClient = FirewallClient()
+    private val firewallClient = FirewallClient()
 
     @Test
     fun testConnect() {
         // Given
         app.launchAndEnsureLoggedIn(accountTestRule.validAccountNumber)
 
-        // When
-        device.findObjectWithTimeout(By.text("Connect")).click()
-        device.findObjectWithTimeout(By.text("OK")).click()
+        on<ConnectPage> { clickConnect() }
 
-        // Then
-        device.findObjectWithTimeout(By.text("CONNECTED"), VERY_LONG_TIMEOUT)
+        on<SystemVpnConfigurationAlert> { clickOk() }
+
+        on<ConnectPage> { waitForConnectedLabel() }
     }
 
     @Test
@@ -54,57 +51,109 @@ class ConnectionTest : EndToEndTest(BuildConfig.FLAVOR_infrastructure) {
         // Given
         app.launchAndEnsureLoggedIn(accountTestRule.validAccountNumber)
 
-        // When
-        device.findObjectWithTimeout(By.text("Connect")).click()
-        device.findObjectWithTimeout(By.text("OK")).click()
-        device.findObjectWithTimeout(By.text("CONNECTED"), VERY_LONG_TIMEOUT)
-        val expected = ConnCheckState(true, app.extractOutIpv4Address())
+        on<ConnectPage> { clickConnect() }
+
+        on<SystemVpnConfigurationAlert> { clickOk() }
+
+        var expectedConnectionState: ConnCheckState? = null
+
+        on<ConnectPage> {
+            waitForConnectedLabel()
+            expectedConnectionState = ConnCheckState(true, extractOutIpv4Address())
+        }
 
         // Then
         val result = SimpleMullvadHttpClient(targetContext).runConnectionCheck()
-        assertEquals(expected, result)
+        assertEquals(expectedConnectionState, result)
     }
 
     @Test
     @HasDependencyOnLocalAPI
     @ClearFirewallRules
-    fun testWireGuardObfuscationOff() = runBlocking {
+    fun testWireGuardObfuscationAutomatic() = runBlocking {
         app.launchAndEnsureLoggedIn(accountTestRule.validAccountNumber)
+        enableLocalNetworkSharing()
+
+        on<ConnectPage> { clickSelectLocation() }
+
+        on<SelectLocationPage> {
+            clickLocationExpandButton(DEFAULT_COUNTRY)
+            clickLocationExpandButton(DEFAULT_CITY)
+            clickLocationCell(DEFAULT_RELAY)
+        }
+
+        on<SystemVpnConfigurationAlert> { clickOk() }
+
+        var relayIpAddress: String? = null
 
+        on<ConnectPage> {
+            waitForConnectedLabel()
+            relayIpAddress = extractInIpv4Address()
+            clickDisconnect()
+        }
+
+        // Block UDP traffic to the relay
+        val firewallRule = DropRule.blockUDPTrafficRule(relayIpAddress!!)
+        firewallClient.createRule(firewallRule)
+
+        on<ConnectPage> {
+            clickConnect()
+            // Currently it takes ~45 seconds to connect with wg obfuscation automatic and UDP
+            // traffic blocked so we need to be very forgiving
+            waitForConnectedLabel(timeout = VERY_FORGIVING_WIREGUARD_OFF_CONNECTION_TIMEOUT)
+        }
+    }
+
+    @Test
+    @HasDependencyOnLocalAPI
+    @ClearFirewallRules
+    fun testWireGuardObfuscationOff() = runBlocking {
+        app.launchAndEnsureLoggedIn(accountTestRule.validAccountNumber)
         enableLocalNetworkSharing()
 
-        device.findObjectWithTimeout(By.res(SELECT_LOCATION_BUTTON_TEST_TAG)).click()
-        clickLocationExpandButton(DEFAULT_COUNTRY)
-        clickLocationExpandButton(DEFAULT_CITY)
-        device.findObjectWithTimeout(By.text(DEFAULT_RELAY)).click()
-        device.findObjectWithTimeout(By.text("OK")).click()
-        device.findObjectWithTimeout(By.text("CONNECTED"), VERY_LONG_TIMEOUT)
-        val relayIpAddress = app.extractInIpv4Address()
-        device.findObjectWithTimeout(By.text("Disconnect")).click()
+        on<ConnectPage> { clickSelectLocation() }
 
-        // Disable obfuscation
-        device.findObjectWithTimeout(By.res(TOP_BAR_SETTINGS_BUTTON)).click()
-        device.findObjectWithTimeout(By.text("VPN settings")).click()
-        val scrollView = device.findObjectWithTimeout(By.res(SETTINGS_SCROLL_VIEW_TEST_TAG))
-        scrollView.scrollUntil(
-            Direction.DOWN,
-            Until.hasObject(By.res(WIREGUARD_OBFUSCATION_OFF_CELL_TEST_TAG)),
-        )
-        device.findObjectWithTimeout(By.res(WIREGUARD_OBFUSCATION_OFF_CELL_TEST_TAG)).click()
-        device.pressBack()
-        device.pressBack()
+        on<SelectLocationPage> {
+            clickLocationExpandButton(DEFAULT_COUNTRY)
+            clickLocationExpandButton(DEFAULT_CITY)
+            clickLocationCell(DEFAULT_RELAY)
+        }
+
+        on<SystemVpnConfigurationAlert> { clickOk() }
+
+        var relayIpAddress: String? = null
+
+        on<ConnectPage> {
+            waitForConnectedLabel()
+            relayIpAddress = extractInIpv4Address()
+            clickDisconnect()
+        }
 
         // Block UDP traffic to the relay
-        val firewallRule = DropRule.blockUDPTrafficRule(relayIpAddress)
+        val firewallRule = DropRule.blockUDPTrafficRule(relayIpAddress!!)
         firewallClient.createRule(firewallRule)
 
-        // Ensure it is not possible to connect to relay
-        device.findObjectWithTimeout(By.text("Connect")).click()
-        // Give it some time and then verify still unable to connect. This duration must be long
-        // enough to ensure all retry attempts have been made.
-        delay(UNSUCCESSFUL_CONNECTION_TIMEOUT.milliseconds)
-        device.findObjectWithTimeout(By.text(("CONNECTING...")))
-        device.findObjectWithTimeout(By.text("Cancel")).click()
+        // Enable UDP-over-TCP
+        on<TopBar> { clickSettings() }
+
+        on<SettingsPage> { clickVpnSettings() }
+
+        on<VpnSettingsPage> {
+            scrollUntilWireguardObfuscationOffCell()
+            clickWireguardObfuscationOffCell()
+        }
+
+        device.pressBack()
+        device.pressBack()
+
+        on<ConnectPage> {
+            clickConnect() // Ensure it is not possible to connect to relay
+            // Give it some time and then verify still unable to connect. This duration must be long
+            // enough to ensure all retry attempts have been made.
+            delay(UNSUCCESSFUL_CONNECTION_TIMEOUT.milliseconds)
+            waitForConnectingLabel()
+            clickCancel()
+        }
     }
 
     @Test
@@ -113,68 +162,63 @@ class ConnectionTest : EndToEndTest(BuildConfig.FLAVOR_infrastructure) {
     fun testUDPOverTCP() =
         runBlocking<Unit> {
             app.launchAndEnsureLoggedIn(accountTestRule.validAccountNumber)
-
             enableLocalNetworkSharing()
 
-            device.findObjectWithTimeout(By.res(SELECT_LOCATION_BUTTON_TEST_TAG)).click()
-            clickLocationExpandButton(DEFAULT_COUNTRY)
-            clickLocationExpandButton(DEFAULT_CITY)
-            device.findObjectWithTimeout(By.text(DEFAULT_RELAY)).click()
-            device.findObjectWithTimeout(By.text("OK")).click()
-            device.findObjectWithTimeout(By.text("CONNECTED"), VERY_LONG_TIMEOUT)
-            val relayIpAddress = app.extractInIpv4Address()
-            device.findObjectWithTimeout(By.text("Disconnect")).click()
+            on<ConnectPage> { clickSelectLocation() }
+
+            on<SelectLocationPage> {
+                clickLocationExpandButton(DEFAULT_COUNTRY)
+                clickLocationExpandButton(DEFAULT_CITY)
+                clickLocationCell(DEFAULT_RELAY)
+            }
+
+            on<SystemVpnConfigurationAlert> { clickOk() }
+
+            var relayIpAddress: String? = null
+
+            on<ConnectPage> {
+                waitForConnectedLabel()
+                relayIpAddress = extractInIpv4Address()
+                clickDisconnect()
+            }
 
             // Block UDP traffic to the relay
-            val firewallRule = DropRule.blockUDPTrafficRule(relayIpAddress)
+            val firewallRule = DropRule.blockUDPTrafficRule(relayIpAddress!!)
             firewallClient.createRule(firewallRule)
 
             // Enable UDP-over-TCP
-            device.findObjectWithTimeout(By.res(TOP_BAR_SETTINGS_BUTTON)).click()
-            device.findObjectWithTimeout(By.text("VPN settings")).click()
-            val scrollView2 = device.findObjectWithTimeout(By.res(SETTINGS_SCROLL_VIEW_TEST_TAG))
-            scrollView2.scrollUntil(
-                Direction.DOWN,
-                Until.hasObject(By.res(WIREGUARD_OBFUSCATION_UDP_OVER_TCP_CELL_TEST_TAG)),
-            )
-            device
-                .findObjectWithTimeout(By.res(WIREGUARD_OBFUSCATION_UDP_OVER_TCP_CELL_TEST_TAG))
-                .click()
+            on<TopBar> { clickSettings() }
+
+            on<SettingsPage> { clickVpnSettings() }
+
+            on<VpnSettingsPage> {
+                scrollUntilWireguardObfuscationUdpOverTcpCell()
+                clickWireguardObfuscationUdpOverTcpCell()
+            }
+
             device.pressBack()
             device.pressBack()
 
-            // Ensure it is possible to connect by using UDP-over-TCP
-            device.findObjectWithTimeout(By.text("Connect")).click()
-            device.findObjectWithTimeout(By.text("CONNECTED"), EXTREMELY_LONG_TIMEOUT)
-            device.findObjectWithTimeout(By.text("Disconnect")).click()
+            on<ConnectPage> {
+                clickConnect()
+                waitForConnectedLabel(timeout = EXTREMELY_LONG_TIMEOUT)
+                clickDisconnect()
+            }
         }
 
     private fun enableLocalNetworkSharing() {
-        device.findObjectWithTimeout(By.res(TOP_BAR_SETTINGS_BUTTON)).click()
-        device.findObjectWithTimeout(By.text("VPN settings")).click()
+        on<TopBar> { clickSettings() }
 
-        val localNetworkSharingCell =
-            device.findObjectWithTimeout(By.text("Local network sharing")).parent
-        val localNetworkSharingSwitch =
-            localNetworkSharingCell.findObjectWithTimeout(By.res(SWITCH_TEST_TAG))
+        on<SettingsPage> { clickVpnSettings() }
+
+        on<VpnSettingsPage> { clickLocalNetworkSharingSwitch() }
 
-        localNetworkSharingSwitch.click()
         device.pressBack()
         device.pressBack()
     }
 
-    private fun clickLocationExpandButton(locationName: String) {
-        val locationCell = device.findObjectWithTimeout(By.text(locationName)).parent.parent
-        val expandButton = locationCell.findObjectWithTimeout(By.res(EXPAND_BUTTON_TEST_TAG))
-        expandButton.click()
-    }
-
     companion object {
-        const val SETTINGS_SCROLL_VIEW_TEST_TAG = "lazy_list_vpn_settings_test_tag"
-        const val WIREGUARD_OBFUSCATION_OFF_CELL_TEST_TAG =
-            "wireguard_obfuscation_off_cell_test_tag"
-        const val WIREGUARD_OBFUSCATION_UDP_OVER_TCP_CELL_TEST_TAG =
-            "wireguard_obfuscation_udp_over_tcp_cell_test_tag"
+        const val VERY_FORGIVING_WIREGUARD_OFF_CONNECTION_TIMEOUT = 60000L
         const val UNSUCCESSFUL_CONNECTION_TIMEOUT = 60000L
     }
 }