Convert screen tests to Junit5

17 files changed, 2481 insertions, 2320 deletions

diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/ComposeRuleExtensions.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/ComposeRuleExtensions.kt
index 69cd530b19..7566051c45 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/ComposeRuleExtensions.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/ComposeRuleExtensions.kt
@@ -1,9 +1,9 @@
 package net.mullvad.mullvadvpn.compose
 
 import androidx.compose.runtime.Composable
-import androidx.compose.ui.test.junit4.ComposeContentTestRule
+import de.mannodermaus.junit5.compose.ComposeContext
 import net.mullvad.mullvadvpn.lib.theme.AppTheme
 
-fun ComposeContentTestRule.setContentWithTheme(content: @Composable () -> Unit) {
+fun ComposeContext.setContentWithTheme(content: @Composable () -> Unit) {
     setContent { AppTheme { content() } }
 }
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/CustomPortDialogTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/CustomPortDialogTest.kt
index 43e385b65d..1fd106379d 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/CustomPortDialogTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/CustomPortDialogTest.kt
@@ -2,22 +2,26 @@ package net.mullvad.mullvadvpn.compose.dialog
 
 import android.annotation.SuppressLint
 import androidx.compose.runtime.Composable
-import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.onNodeWithTag
 import androidx.compose.ui.test.performTextInput
+import de.mannodermaus.junit5.compose.createComposeExtension
 import io.mockk.MockKAnnotations
 import net.mullvad.mullvadvpn.compose.setContentWithTheme
 import net.mullvad.mullvadvpn.compose.test.CUSTOM_PORT_DIALOG_INPUT_TEST_TAG
 import net.mullvad.mullvadvpn.model.PortRange
 import net.mullvad.mullvadvpn.onNodeWithTagAndText
-import org.junit.Before
-import org.junit.Rule
-import org.junit.Test
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.RegisterExtension
 
 class CustomPortDialogTest {
-    @get:Rule val composeTestRule = createComposeRule()
+    @OptIn(ExperimentalTestApi::class)
+    @JvmField
+    @RegisterExtension
+    val composeExtension = createComposeExtension()
 
-    @Before
+    @BeforeEach
     fun setup() {
         MockKAnnotations.init(this)
     }
@@ -40,23 +44,22 @@ class CustomPortDialogTest {
     }
 
     @Test
-    fun testShowWireguardCustomPortDialogInvalidInt() {
-        // Input a number to make sure that a too long number does not show and it does not crash
-        // the app
+    fun testShowWireguardCustomPortDialogInvalidInt() =
+        composeExtension.use {
+            // Input a number to make sure that a too long number does not show and it does not
+            // crash
+            // the app
 
-        // Arrange
-        composeTestRule.setContentWithTheme { testWireguardCustomPortDialog() }
+            // Arrange
+            setContentWithTheme { testWireguardCustomPortDialog() }
 
-        // Act
-        composeTestRule
-            .onNodeWithTag(CUSTOM_PORT_DIALOG_INPUT_TEST_TAG)
-            .performTextInput(invalidCustomPort)
+            // Act
+            onNodeWithTag(CUSTOM_PORT_DIALOG_INPUT_TEST_TAG).performTextInput(invalidCustomPort)
 
-        // Assert
-        composeTestRule
-            .onNodeWithTagAndText(CUSTOM_PORT_DIALOG_INPUT_TEST_TAG, invalidCustomPort)
-            .assertDoesNotExist()
-    }
+            // Assert
+            onNodeWithTagAndText(CUSTOM_PORT_DIALOG_INPUT_TEST_TAG, invalidCustomPort)
+                .assertDoesNotExist()
+        }
 
     companion object {
         const val invalidCustomPort = "21474836471"
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/DnsDialogTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/DnsDialogTest.kt
index bc8d87b244..5f9ffaea95 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/DnsDialogTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/DnsDialogTest.kt
@@ -2,16 +2,20 @@ package net.mullvad.mullvadvpn.compose.dialog
 
 import android.annotation.SuppressLint
 import androidx.compose.runtime.Composable
+import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.assertIsNotEnabled
-import androidx.compose.ui.test.junit4.createComposeRule
 import androidx.compose.ui.test.onNodeWithText
+import de.mannodermaus.junit5.compose.createComposeExtension
 import net.mullvad.mullvadvpn.compose.setContentWithTheme
 import net.mullvad.mullvadvpn.viewmodel.DnsDialogViewState
-import org.junit.Rule
-import org.junit.Test
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.RegisterExtension
 
 class DnsDialogTest {
-    @get:Rule val composeTestRule = createComposeRule()
+    @OptIn(ExperimentalTestApi::class)
+    @JvmField
+    @RegisterExtension
+    val composeExtension = createComposeExtension()
 
     private val defaultState =
         DnsDialogViewState(
@@ -35,80 +39,86 @@ class DnsDialogTest {
     }
 
     @Test
-    fun testDnsDialogLanWarningShownWhenLanTrafficDisabledAndLocalAddressUsed() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            testDnsDialog(defaultState.copy(isAllowLanEnabled = false, isLocal = true))
-        }
+    fun testDnsDialogLanWarningShownWhenLanTrafficDisabledAndLocalAddressUsed() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                testDnsDialog(defaultState.copy(isAllowLanEnabled = false, isLocal = true))
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText(LOCAL_DNS_SERVER_WARNING).assertExists()
-    }
+            // Assert
+            onNodeWithText(LOCAL_DNS_SERVER_WARNING).assertExists()
+        }
 
     @Test
-    fun testDnsDialogLanWarningNotShownWhenLanTrafficEnabledAndLocalAddressUsed() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            testDnsDialog(defaultState.copy(isAllowLanEnabled = true, isLocal = true))
-        }
+    fun testDnsDialogLanWarningNotShownWhenLanTrafficEnabledAndLocalAddressUsed() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                testDnsDialog(defaultState.copy(isAllowLanEnabled = true, isLocal = true))
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText(LOCAL_DNS_SERVER_WARNING).assertDoesNotExist()
-    }
+            // Assert
+            onNodeWithText(LOCAL_DNS_SERVER_WARNING).assertDoesNotExist()
+        }
 
     @Test
-    fun testDnsDialogLanWarningNotShownWhenLanTrafficEnabledAndNonLocalAddressUsed() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            testDnsDialog(defaultState.copy(isAllowLanEnabled = true, isLocal = false))
-        }
+    fun testDnsDialogLanWarningNotShownWhenLanTrafficEnabledAndNonLocalAddressUsed() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                testDnsDialog(defaultState.copy(isAllowLanEnabled = true, isLocal = false))
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText(LOCAL_DNS_SERVER_WARNING).assertDoesNotExist()
-    }
+            // Assert
+            onNodeWithText(LOCAL_DNS_SERVER_WARNING).assertDoesNotExist()
+        }
 
     @Test
-    fun testDnsDialogLanWarningNotShownWhenLanTrafficDisabledAndNonLocalAddressUsed() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            testDnsDialog(defaultState.copy(isAllowLanEnabled = false, isLocal = false))
-        }
+    fun testDnsDialogLanWarningNotShownWhenLanTrafficDisabledAndNonLocalAddressUsed() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                testDnsDialog(defaultState.copy(isAllowLanEnabled = false, isLocal = false))
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText(LOCAL_DNS_SERVER_WARNING).assertDoesNotExist()
-    }
+            // Assert
+            onNodeWithText(LOCAL_DNS_SERVER_WARNING).assertDoesNotExist()
+        }
 
     @Test
-    fun testDnsDialogSubmitButtonDisabledOnInvalidDnsAddress() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            testDnsDialog(
-                defaultState.copy(
-                    ipAddress = invalidIpAddress,
-                    validationResult = DnsDialogViewState.ValidationResult.InvalidAddress,
+    fun testDnsDialogSubmitButtonDisabledOnInvalidDnsAddress() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                testDnsDialog(
+                    defaultState.copy(
+                        ipAddress = invalidIpAddress,
+                        validationResult = DnsDialogViewState.ValidationResult.InvalidAddress,
+                    )
                 )
-            )
-        }
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText("Submit").assertIsNotEnabled()
-    }
+            // Assert
+            onNodeWithText("Submit").assertIsNotEnabled()
+        }
 
     @Test
-    fun testDnsDialogSubmitButtonDisabledOnDuplicateDnsAddress() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            testDnsDialog(
-                defaultState.copy(
-                    ipAddress = "192.168.0.1",
-                    validationResult = DnsDialogViewState.ValidationResult.DuplicateAddress,
+    fun testDnsDialogSubmitButtonDisabledOnDuplicateDnsAddress() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                testDnsDialog(
+                    defaultState.copy(
+                        ipAddress = "192.168.0.1",
+                        validationResult = DnsDialogViewState.ValidationResult.DuplicateAddress,
+                    )
                 )
-            )
-        }
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText("Submit").assertIsNotEnabled()
-    }
+            // Assert
+            onNodeWithText("Submit").assertIsNotEnabled()
+        }
 
     companion object {
         private const val LOCAL_DNS_SERVER_WARNING =
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/MtuDialogTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/MtuDialogTest.kt
index 38a3bd170d..28d089cc7e 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/MtuDialogTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/MtuDialogTest.kt
@@ -2,24 +2,28 @@ package net.mullvad.mullvadvpn.compose.dialog
 
 import android.annotation.SuppressLint
 import androidx.compose.runtime.Composable
+import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.assertIsEnabled
 import androidx.compose.ui.test.assertIsNotEnabled
-import androidx.compose.ui.test.junit4.createComposeRule
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
 import androidx.compose.ui.test.performTextInput
+import de.mannodermaus.junit5.compose.createComposeExtension
 import io.mockk.MockKAnnotations
 import io.mockk.mockk
 import io.mockk.verify
 import net.mullvad.mullvadvpn.compose.setContentWithTheme
-import org.junit.Before
-import org.junit.Rule
-import org.junit.Test
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.RegisterExtension
 
 class MtuDialogTest {
-    @get:Rule val composeTestRule = createComposeRule()
+    @OptIn(ExperimentalTestApi::class)
+    @JvmField
+    @RegisterExtension
+    val composeExtension = createComposeExtension()
 
-    @Before
+    @BeforeEach
     fun setup() {
         MockKAnnotations.init(this)
     }
@@ -41,109 +45,114 @@ class MtuDialogTest {
     }
 
     @Test
-    fun testMtuDialogWithDefaultValue() {
-        // Arrange
-        composeTestRule.setContentWithTheme { testMtuDialog() }
+    fun testMtuDialogWithDefaultValue() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme { testMtuDialog() }
 
-        // Assert
-        composeTestRule.onNodeWithText(EMPTY_STRING).assertExists()
-    }
+            // Assert
+            onNodeWithText(EMPTY_STRING).assertExists()
+        }
 
     @Test
-    fun testMtuDialogWithEditValue() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            testMtuDialog(
-                mtuInitial = VALID_DUMMY_MTU_VALUE,
-            )
-        }
+    fun testMtuDialogWithEditValue() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                testMtuDialog(
+                    mtuInitial = VALID_DUMMY_MTU_VALUE,
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText(VALID_DUMMY_MTU_VALUE.toString()).assertExists()
-    }
+            // Assert
+            onNodeWithText(VALID_DUMMY_MTU_VALUE.toString()).assertExists()
+        }
 
     @Test
-    fun testMtuDialogTextInput() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            testMtuDialog(
-                null,
-            )
-        }
+    fun testMtuDialogTextInput() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                testMtuDialog(
+                    null,
+                )
+            }
 
-        // Act
-        composeTestRule
-            .onNodeWithText(EMPTY_STRING)
-            .performTextInput(VALID_DUMMY_MTU_VALUE.toString())
+            // Act
+            onNodeWithText(EMPTY_STRING).performTextInput(VALID_DUMMY_MTU_VALUE.toString())
 
-        // Assert
-        composeTestRule.onNodeWithText(VALID_DUMMY_MTU_VALUE.toString()).assertExists()
-    }
+            // Assert
+            onNodeWithText(VALID_DUMMY_MTU_VALUE.toString()).assertExists()
+        }
 
     @Test
-    fun testMtuDialogSubmitOfValidValue() {
-        // Arrange
-        val mockedSubmitHandler: (Int) -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            testMtuDialog(
-                VALID_DUMMY_MTU_VALUE,
-                onSaveMtu = mockedSubmitHandler,
-            )
-        }
+    fun testMtuDialogSubmitOfValidValue() =
+        composeExtension.use {
+            // Arrange
+            val mockedSubmitHandler: (Int) -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                testMtuDialog(
+                    VALID_DUMMY_MTU_VALUE,
+                    onSaveMtu = mockedSubmitHandler,
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithText("Submit").assertIsEnabled().performClick()
+            // Act
+            onNodeWithText("Submit").assertIsEnabled().performClick()
 
-        // Assert
-        verify { mockedSubmitHandler.invoke(VALID_DUMMY_MTU_VALUE) }
-    }
+            // Assert
+            verify { mockedSubmitHandler.invoke(VALID_DUMMY_MTU_VALUE) }
+        }
 
     @Test
-    fun testMtuDialogSubmitButtonDisabledWhenInvalidInput() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            testMtuDialog(
-                INVALID_DUMMY_MTU_VALUE,
-            )
-        }
+    fun testMtuDialogSubmitButtonDisabledWhenInvalidInput() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                testMtuDialog(
+                    INVALID_DUMMY_MTU_VALUE,
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText("Submit").assertIsNotEnabled()
-    }
+            // Assert
+            onNodeWithText("Submit").assertIsNotEnabled()
+        }
 
     @Test
-    fun testMtuDialogResetClick() {
-        // Arrange
-        val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            testMtuDialog(
-                onResetMtu = mockedClickHandler,
-            )
-        }
+    fun testMtuDialogResetClick() =
+        composeExtension.use {
+            // Arrange
+            val mockedClickHandler: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                testMtuDialog(
+                    onResetMtu = mockedClickHandler,
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithText("Reset to default").performClick()
+            // Act
+            onNodeWithText("Reset to default").performClick()
 
-        // Assert
-        verify { mockedClickHandler.invoke() }
-    }
+            // Assert
+            verify { mockedClickHandler.invoke() }
+        }
 
     @Test
-    fun testMtuDialogCancelClick() {
-        // Arrange
-        val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            testMtuDialog(
-                onDismiss = mockedClickHandler,
-            )
-        }
+    fun testMtuDialogCancelClick() =
+        composeExtension.use {
+            // Arrange
+            val mockedClickHandler: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                testMtuDialog(
+                    onDismiss = mockedClickHandler,
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText("Cancel").performClick()
+            // Assert
+            onNodeWithText("Cancel").performClick()
 
-        // Assert
-        verify { mockedClickHandler.invoke() }
-    }
+            // Assert
+            verify { mockedClickHandler.invoke() }
+        }
 
     companion object {
         private const val EMPTY_STRING = ""
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/PaymentDialogTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/PaymentDialogTest.kt
index 1a626ecf19..9012b3144f 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/PaymentDialogTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/PaymentDialogTest.kt
@@ -1,57 +1,64 @@
 package net.mullvad.mullvadvpn.compose.dialog
 
-import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.onNodeWithText
+import de.mannodermaus.junit5.compose.createComposeExtension
 import net.mullvad.mullvadvpn.compose.dialog.payment.PaymentDialog
 import net.mullvad.mullvadvpn.compose.setContentWithTheme
 import net.mullvad.mullvadvpn.lib.payment.model.ProductId
 import net.mullvad.mullvadvpn.lib.payment.model.PurchaseResult
 import net.mullvad.mullvadvpn.util.toPaymentDialogData
-import org.junit.Rule
-import org.junit.Test
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.RegisterExtension
 
 class PaymentDialogTest {
-    @get:Rule val composeTestRule = createComposeRule()
+    @OptIn(ExperimentalTestApi::class)
+    @JvmField
+    @RegisterExtension
+    val composeExtension = createComposeExtension()
 
     @Test
-    fun testShowPurchaseCompleteDialog() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            PaymentDialog(
-                paymentDialogData = PurchaseResult.Completed.Success.toPaymentDialogData()!!
-            )
-        }
+    fun testShowPurchaseCompleteDialog() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                PaymentDialog(
+                    paymentDialogData = PurchaseResult.Completed.Success.toPaymentDialogData()!!
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText("Time was successfully added").assertExists()
-    }
+            // Assert
+            onNodeWithText("Time was successfully added").assertExists()
+        }
 
     @Test
-    fun testShowVerificationErrorDialog() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            PaymentDialog(
-                paymentDialogData =
-                    PurchaseResult.Error.VerificationError(null).toPaymentDialogData()!!
-            )
-        }
+    fun testShowVerificationErrorDialog() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                PaymentDialog(
+                    paymentDialogData =
+                        PurchaseResult.Error.VerificationError(null).toPaymentDialogData()!!
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText("Verifying purchase").assertExists()
-    }
+            // Assert
+            onNodeWithText("Verifying purchase").assertExists()
+        }
 
     @Test
-    fun testShowFetchProductsErrorDialog() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            PaymentDialog(
-                paymentDialogData =
-                    PurchaseResult.Error.FetchProductsError(ProductId(""), null)
-                        .toPaymentDialogData()!!
-            )
-        }
+    fun testShowFetchProductsErrorDialog() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                PaymentDialog(
+                    paymentDialogData =
+                        PurchaseResult.Error.FetchProductsError(ProductId(""), null)
+                            .toPaymentDialogData()!!
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText("Google Play unavailable").assertExists()
-    }
+            // Assert
+            onNodeWithText("Google Play unavailable").assertExists()
+        }
 }
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/AccountScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/AccountScreenTest.kt
index 3b42cc1c3b..bead0a02e5 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/AccountScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/AccountScreenTest.kt
@@ -1,10 +1,11 @@
 package net.mullvad.mullvadvpn.compose.screen
 
 import androidx.compose.material3.ExperimentalMaterial3Api
-import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.onNodeWithTag
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
+import de.mannodermaus.junit5.compose.createComposeExtension
 import io.mockk.MockKAnnotations
 import io.mockk.every
 import io.mockk.mockk
@@ -20,252 +21,272 @@ import net.mullvad.mullvadvpn.lib.payment.model.ProductId
 import net.mullvad.mullvadvpn.lib.payment.model.ProductPrice
 import net.mullvad.mullvadvpn.viewmodel.AccountUiState
 import net.mullvad.mullvadvpn.viewmodel.AccountViewModel
-import org.junit.Before
-import org.junit.Rule
-import org.junit.Test
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.RegisterExtension
 
+@ExperimentalTestApi
 @OptIn(ExperimentalMaterial3Api::class)
 class AccountScreenTest {
-    @get:Rule val composeTestRule = createComposeRule()
+    @JvmField @RegisterExtension val composeExtension = createComposeExtension()
 
-    @Before
+    @BeforeEach
     fun setup() {
         MockKAnnotations.init(this)
     }
 
     @Test
-    fun testDefaultState() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            AccountScreen(
-                uiState =
-                    AccountUiState(
-                        deviceName = DUMMY_DEVICE_NAME,
-                        accountNumber = DUMMY_ACCOUNT_NUMBER,
-                        accountExpiry = null,
-                        showSitePayment = false
-                    ),
-                uiSideEffect = MutableSharedFlow<AccountViewModel.UiSideEffect>().asSharedFlow(),
-            )
-        }
+    fun testDefaultState() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                AccountScreen(
+                    uiState =
+                        AccountUiState(
+                            deviceName = DUMMY_DEVICE_NAME,
+                            accountNumber = DUMMY_ACCOUNT_NUMBER,
+                            accountExpiry = null,
+                            showSitePayment = false
+                        ),
+                    uiSideEffect =
+                        MutableSharedFlow<AccountViewModel.UiSideEffect>().asSharedFlow(),
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText("Redeem voucher").assertExists()
             onNodeWithText("Log out").assertExists()
         }
-    }
 
     @Test
-    fun testManageAccountClick() {
-        // Arrange
-        val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            AccountScreen(
-                uiState =
-                    AccountUiState(
-                        showSitePayment = true,
-                        deviceName = DUMMY_DEVICE_NAME,
-                        accountNumber = DUMMY_ACCOUNT_NUMBER,
-                        accountExpiry = null,
-                    ),
-                uiSideEffect = MutableSharedFlow<AccountViewModel.UiSideEffect>().asSharedFlow(),
-                onManageAccountClick = mockedClickHandler
-            )
-        }
+    fun testManageAccountClick() =
+        composeExtension.use {
+            // Arrange
+            val mockedClickHandler: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                AccountScreen(
+                    uiState =
+                        AccountUiState(
+                            showSitePayment = true,
+                            deviceName = DUMMY_DEVICE_NAME,
+                            accountNumber = DUMMY_ACCOUNT_NUMBER,
+                            accountExpiry = null,
+                        ),
+                    uiSideEffect =
+                        MutableSharedFlow<AccountViewModel.UiSideEffect>().asSharedFlow(),
+                    onManageAccountClick = mockedClickHandler
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithText("Manage account").performClick()
+            // Act
+            onNodeWithText("Manage account").performClick()
 
-        // Assert
-        verify { mockedClickHandler.invoke() }
-    }
+            // Assert
+            verify(exactly = 1) { mockedClickHandler.invoke() }
+        }
 
     @Test
-    fun testRedeemVoucherClick() {
-        // Arrange
-        val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            AccountScreen(
-                uiState =
-                    AccountUiState(
-                        deviceName = DUMMY_DEVICE_NAME,
-                        accountNumber = DUMMY_ACCOUNT_NUMBER,
-                        accountExpiry = null,
-                        showSitePayment = false
-                    ),
-                uiSideEffect = MutableSharedFlow<AccountViewModel.UiSideEffect>().asSharedFlow(),
-                onRedeemVoucherClick = mockedClickHandler
-            )
-        }
+    fun testRedeemVoucherClick() =
+        composeExtension.use {
+            // Arrange
+            val mockedClickHandler: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                AccountScreen(
+                    uiState =
+                        AccountUiState(
+                            deviceName = DUMMY_DEVICE_NAME,
+                            accountNumber = DUMMY_ACCOUNT_NUMBER,
+                            accountExpiry = null,
+                            showSitePayment = false
+                        ),
+                    uiSideEffect =
+                        MutableSharedFlow<AccountViewModel.UiSideEffect>().asSharedFlow(),
+                    onRedeemVoucherClick = mockedClickHandler
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithText("Redeem voucher").performClick()
+            // Act
+            onNodeWithText("Redeem voucher").performClick()
 
-        // Assert
-        verify { mockedClickHandler.invoke() }
-    }
+            // Assert
+            verify { mockedClickHandler.invoke() }
+        }
 
     @Test
-    fun testLogoutClick() {
-        // Arrange
-        val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            AccountScreen(
-                uiState =
-                    AccountUiState(
-                        deviceName = DUMMY_DEVICE_NAME,
-                        accountNumber = DUMMY_ACCOUNT_NUMBER,
-                        accountExpiry = null,
-                        showSitePayment = false
-                    ),
-                uiSideEffect = MutableSharedFlow<AccountViewModel.UiSideEffect>().asSharedFlow(),
-                onLogoutClick = mockedClickHandler
-            )
-        }
+    fun testLogoutClick() =
+        composeExtension.use {
+            // Arrange
+            val mockedClickHandler: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                AccountScreen(
+                    uiState =
+                        AccountUiState(
+                            deviceName = DUMMY_DEVICE_NAME,
+                            accountNumber = DUMMY_ACCOUNT_NUMBER,
+                            accountExpiry = null,
+                            showSitePayment = false
+                        ),
+                    uiSideEffect =
+                        MutableSharedFlow<AccountViewModel.UiSideEffect>().asSharedFlow(),
+                    onLogoutClick = mockedClickHandler
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithText("Log out").performClick()
+            // Act
+            onNodeWithText("Log out").performClick()
 
-        // Assert
-        verify { mockedClickHandler.invoke() }
-    }
+            // Assert
+            verify { mockedClickHandler.invoke() }
+        }
 
     @Test
-    fun testShowBillingErrorPaymentButton() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            AccountScreen(
-                uiState =
-                    AccountUiState.default().copy(billingPaymentState = PaymentState.Error.Billing),
-                uiSideEffect = MutableSharedFlow<AccountViewModel.UiSideEffect>().asSharedFlow(),
-            )
-        }
+    fun testShowBillingErrorPaymentButton() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                AccountScreen(
+                    uiState =
+                        AccountUiState.default()
+                            .copy(billingPaymentState = PaymentState.Error.Billing),
+                    uiSideEffect =
+                        MutableSharedFlow<AccountViewModel.UiSideEffect>().asSharedFlow(),
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText("Add 30 days time").assertExists()
-    }
+            // Assert
+            onNodeWithText("Add 30 days time").assertExists()
+        }
 
     @Test
-    fun testShowBillingPaymentAvailable() {
-        // Arrange
-        val mockPaymentProduct: PaymentProduct = mockk()
-        every { mockPaymentProduct.price } returns ProductPrice("$10")
-        every { mockPaymentProduct.status } returns null
-        composeTestRule.setContentWithTheme {
-            AccountScreen(
-                uiState =
-                    AccountUiState.default()
-                        .copy(
-                            billingPaymentState =
-                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                        ),
-                uiSideEffect = MutableSharedFlow<AccountViewModel.UiSideEffect>().asSharedFlow(),
-            )
-        }
+    fun testShowBillingPaymentAvailable() =
+        composeExtension.use {
+            // Arrange
+            val mockPaymentProduct: PaymentProduct = mockk()
+            every { mockPaymentProduct.price } returns ProductPrice("$10")
+            every { mockPaymentProduct.status } returns null
+            setContentWithTheme {
+                AccountScreen(
+                    uiState =
+                        AccountUiState.default()
+                            .copy(
+                                billingPaymentState =
+                                    PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+                            ),
+                    uiSideEffect =
+                        MutableSharedFlow<AccountViewModel.UiSideEffect>().asSharedFlow(),
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText("Add 30 days time ($10)").assertExists()
-    }
+            // Assert
+            onNodeWithText("Add 30 days time ($10)").assertExists()
+        }
 
     @Test
-    fun testShowPendingPayment() {
-        // Arrange
-        val mockPaymentProduct: PaymentProduct = mockk()
-        every { mockPaymentProduct.price } returns ProductPrice("$10")
-        every { mockPaymentProduct.status } returns PaymentStatus.PENDING
-        composeTestRule.setContentWithTheme {
-            AccountScreen(
-                uiState =
-                    AccountUiState.default()
-                        .copy(
-                            billingPaymentState =
-                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                        ),
-                uiSideEffect = MutableSharedFlow<AccountViewModel.UiSideEffect>().asSharedFlow(),
-            )
-        }
+    fun testShowPendingPayment() =
+        composeExtension.use {
+            // Arrange
+            val mockPaymentProduct: PaymentProduct = mockk()
+            every { mockPaymentProduct.price } returns ProductPrice("$10")
+            every { mockPaymentProduct.status } returns PaymentStatus.PENDING
+            setContentWithTheme {
+                AccountScreen(
+                    uiState =
+                        AccountUiState.default()
+                            .copy(
+                                billingPaymentState =
+                                    PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+                            ),
+                    uiSideEffect =
+                        MutableSharedFlow<AccountViewModel.UiSideEffect>().asSharedFlow(),
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText("Google Play payment pending").assertExists()
-    }
+            // Assert
+            onNodeWithText("Google Play payment pending").assertExists()
+        }
 
     @Test
-    fun testShowPendingPaymentInfoDialog() {
-        // Arrange
-        val mockPaymentProduct: PaymentProduct = mockk()
-        every { mockPaymentProduct.price } returns ProductPrice("$10")
-        every { mockPaymentProduct.status } returns PaymentStatus.PENDING
-        val mockNavigateToVerificationPending: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            AccountScreen(
-                uiState =
-                    AccountUiState.default()
-                        .copy(
-                            billingPaymentState =
-                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                        ),
-                uiSideEffect = MutableSharedFlow<AccountViewModel.UiSideEffect>().asSharedFlow(),
-                navigateToVerificationPendingDialog = mockNavigateToVerificationPending
-            )
-        }
+    fun testShowPendingPaymentInfoDialog() =
+        composeExtension.use {
+            // Arrange
+            val mockPaymentProduct: PaymentProduct = mockk()
+            every { mockPaymentProduct.price } returns ProductPrice("$10")
+            every { mockPaymentProduct.status } returns PaymentStatus.PENDING
+            val mockNavigateToVerificationPending: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                AccountScreen(
+                    uiState =
+                        AccountUiState.default()
+                            .copy(
+                                billingPaymentState =
+                                    PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+                            ),
+                    uiSideEffect =
+                        MutableSharedFlow<AccountViewModel.UiSideEffect>().asSharedFlow(),
+                    navigateToVerificationPendingDialog = mockNavigateToVerificationPending
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithTag(PLAY_PAYMENT_INFO_ICON_TEST_TAG).performClick()
+            // Act
+            onNodeWithTag(PLAY_PAYMENT_INFO_ICON_TEST_TAG).performClick()
 
-        // Assert
-        verify(exactly = 1) { mockNavigateToVerificationPending.invoke() }
-    }
+            // Assert
+            verify(exactly = 1) { mockNavigateToVerificationPending.invoke() }
+        }
 
     @Test
-    fun testShowVerificationInProgress() {
-        // Arrange
-        val mockPaymentProduct: PaymentProduct = mockk()
-        every { mockPaymentProduct.price } returns ProductPrice("$10")
-        every { mockPaymentProduct.status } returns PaymentStatus.VERIFICATION_IN_PROGRESS
-        composeTestRule.setContentWithTheme {
-            AccountScreen(
-                uiState =
-                    AccountUiState.default()
-                        .copy(
-                            billingPaymentState =
-                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                        ),
-                uiSideEffect = MutableSharedFlow<AccountViewModel.UiSideEffect>().asSharedFlow(),
-            )
-        }
+    fun testShowVerificationInProgress() =
+        composeExtension.use {
+            // Arrange
+            val mockPaymentProduct: PaymentProduct = mockk()
+            every { mockPaymentProduct.price } returns ProductPrice("$10")
+            every { mockPaymentProduct.status } returns PaymentStatus.VERIFICATION_IN_PROGRESS
+            setContentWithTheme {
+                AccountScreen(
+                    uiState =
+                        AccountUiState.default()
+                            .copy(
+                                billingPaymentState =
+                                    PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+                            ),
+                    uiSideEffect =
+                        MutableSharedFlow<AccountViewModel.UiSideEffect>().asSharedFlow(),
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText("Verifying purchase").assertExists()
-    }
+            // Assert
+            onNodeWithText("Verifying purchase").assertExists()
+        }
 
     @Test
-    fun testOnPurchaseBillingProductClick() {
-        // Arrange
-        val clickHandler: (ProductId) -> Unit = mockk(relaxed = true)
-        val mockPaymentProduct: PaymentProduct = mockk()
-        every { mockPaymentProduct.price } returns ProductPrice("$10")
-        every { mockPaymentProduct.productId } returns ProductId("PRODUCT_ID")
-        every { mockPaymentProduct.status } returns null
-        composeTestRule.setContentWithTheme {
-            AccountScreen(
-                uiState =
-                    AccountUiState.default()
-                        .copy(
-                            billingPaymentState =
-                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                        ),
-                onPurchaseBillingProductClick = clickHandler,
-                uiSideEffect = MutableSharedFlow<AccountViewModel.UiSideEffect>().asSharedFlow(),
-            )
-        }
+    fun testOnPurchaseBillingProductClick() =
+        composeExtension.use {
+            // Arrange
+            val clickHandler: (ProductId) -> Unit = mockk(relaxed = true)
+            val mockPaymentProduct: PaymentProduct = mockk()
+            every { mockPaymentProduct.price } returns ProductPrice("$10")
+            every { mockPaymentProduct.productId } returns ProductId("PRODUCT_ID")
+            every { mockPaymentProduct.status } returns null
+            setContentWithTheme {
+                AccountScreen(
+                    uiState =
+                        AccountUiState.default()
+                            .copy(
+                                billingPaymentState =
+                                    PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+                            ),
+                    onPurchaseBillingProductClick = clickHandler,
+                    uiSideEffect =
+                        MutableSharedFlow<AccountViewModel.UiSideEffect>().asSharedFlow(),
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithText("Add 30 days time ($10)").performClick()
+            // Act
+            onNodeWithText("Add 30 days time ($10)").performClick()
 
-        // Assert
-        verify { clickHandler.invoke(ProductId("PRODUCT_ID")) }
-    }
+            // Assert
+            verify { clickHandler.invoke(ProductId("PRODUCT_ID")) }
+        }
 
     companion object {
         private const val DUMMY_DEVICE_NAME = "fake_name"
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ChangelogDialogTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ChangelogDialogTest.kt
index 4e34fe0825..66ed6d25f6 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ChangelogDialogTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ChangelogDialogTest.kt
@@ -1,8 +1,9 @@
 package net.mullvad.mullvadvpn.compose.screen
 
-import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
+import de.mannodermaus.junit5.compose.createComposeExtension
 import io.mockk.MockKAnnotations
 import io.mockk.Runs
 import io.mockk.every
@@ -13,44 +14,47 @@ import net.mullvad.mullvadvpn.compose.dialog.ChangelogDialog
 import net.mullvad.mullvadvpn.compose.setContentWithTheme
 import net.mullvad.mullvadvpn.viewmodel.Changelog
 import net.mullvad.mullvadvpn.viewmodel.ChangelogViewModel
-import org.junit.Before
-import org.junit.Rule
-import org.junit.Test
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.RegisterExtension
 
+@OptIn(ExperimentalTestApi::class)
 class ChangelogDialogTest {
-    @get:Rule val composeTestRule = createComposeRule()
+    @JvmField @RegisterExtension val composeExtension = createComposeExtension()
 
     @MockK lateinit var mockedViewModel: ChangelogViewModel
 
-    @Before
+    @BeforeEach
     fun setup() {
         MockKAnnotations.init(this)
     }
 
     @Test
-    fun testShowChangelogWhenNeeded() {
-        // Arrange
-        every { mockedViewModel.markChangelogAsRead() } just Runs
+    fun testShowChangeLogWhenNeeded() =
+        composeExtension.use {
+            // Arrange
+            // Arrange
+            every { mockedViewModel.markChangelogAsRead() } just Runs
 
-        composeTestRule.setContentWithTheme {
-            ChangelogDialog(
-                Changelog(
-                    changes = listOf(CHANGELOG_ITEM),
-                    version = CHANGELOG_VERSION,
-                ),
-                onDismiss = { mockedViewModel.markChangelogAsRead() }
-            )
-        }
+            setContentWithTheme {
+                ChangelogDialog(
+                    Changelog(
+                        changes = listOf(CHANGELOG_ITEM),
+                        version = CHANGELOG_VERSION,
+                    ),
+                    onDismiss = { mockedViewModel.markChangelogAsRead() }
+                )
+            }
 
-        // Check changelog content showed within dialog
-        composeTestRule.onNodeWithText(CHANGELOG_ITEM).assertExists()
+            // Check changelog content showed within dialog
+            onNodeWithText(CHANGELOG_ITEM).assertExists()
 
-        // perform click on Got It button to check if dismiss occur
-        composeTestRule.onNodeWithText(CHANGELOG_BUTTON_TEXT).performClick()
+            // perform click on Got It button to check if dismiss occur
+            onNodeWithText(CHANGELOG_BUTTON_TEXT).performClick()
 
-        // Assert
-        verify { mockedViewModel.markChangelogAsRead() }
-    }
+            // Assert
+            verify { mockedViewModel.markChangelogAsRead() }
+        }
 
     companion object {
         private const val CHANGELOG_BUTTON_TEXT = "Got it!"
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ConnectScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ConnectScreenTest.kt
index 3838bdc7a0..39e11a3c14 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ConnectScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ConnectScreenTest.kt
@@ -1,9 +1,10 @@
 package net.mullvad.mullvadvpn.compose.screen
 
-import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.onNodeWithTag
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
+import de.mannodermaus.junit5.compose.createComposeExtension
 import io.mockk.MockKAnnotations
 import io.mockk.every
 import io.mockk.mockk
@@ -30,35 +31,38 @@ import net.mullvad.talpid.tunnel.ActionAfterDisconnect
 import net.mullvad.talpid.tunnel.ErrorState
 import net.mullvad.talpid.tunnel.ErrorStateCause
 import org.joda.time.DateTime
-import org.junit.After
-import org.junit.Before
-import org.junit.Rule
-import org.junit.Test
+import org.junit.jupiter.api.AfterEach
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.RegisterExtension
 
 class ConnectScreenTest {
-    @get:Rule val composeTestRule = createComposeRule()
+    @OptIn(ExperimentalTestApi::class)
+    @JvmField
+    @RegisterExtension
+    val composeExtension = createComposeExtension()
 
-    @Before
+    @BeforeEach
     fun setup() {
         MockKAnnotations.init(this)
     }
 
-    @After
+    @AfterEach
     fun teardown() {
         unmockkAll()
     }
 
     @Test
     fun testDefaultState() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                uiState = ConnectUiState.INITIAL,
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                ConnectScreen(
+                    uiState = ConnectUiState.INITIAL,
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithTag(SCROLLABLE_COLUMN_TEST_TAG).assertExists()
             onNodeWithText("UNSECURED CONNECTION").assertExists()
             onNodeWithText("Secure my connection").assertExists()
@@ -67,28 +71,28 @@ class ConnectScreenTest {
 
     @Test
     fun testConnectingState() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                uiState =
-                    ConnectUiState(
-                        location = null,
-                        relayLocation = null,
-                        tunnelUiState = TunnelState.Connecting(null, null),
-                        tunnelRealState = TunnelState.Connecting(null, null),
-                        inAddress = null,
-                        outAddress = "",
-                        showLocation = false,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification = InAppNotification.TunnelStateBlocked,
-                        isPlayBuild = false
-                    ),
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                ConnectScreen(
+                    uiState =
+                        ConnectUiState(
+                            location = null,
+                            relayLocation = null,
+                            tunnelUiState = TunnelState.Connecting(null, null),
+                            tunnelRealState = TunnelState.Connecting(null, null),
+                            inAddress = null,
+                            outAddress = "",
+                            showLocation = false,
+                            deviceName = "",
+                            daysLeftUntilExpiry = null,
+                            inAppNotification = InAppNotification.TunnelStateBlocked,
+                            isPlayBuild = false
+                        ),
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithTag(CIRCULAR_PROGRESS_INDICATOR).assertExists()
             onNodeWithText("CREATING SECURE CONNECTION").assertExists()
             onNodeWithText("Switch location").assertExists()
@@ -99,31 +103,32 @@ class ConnectScreenTest {
 
     @Test
     fun testConnectingStateQuantumSecured() {
-        // Arrange
-        val mockTunnelEndpoint: TunnelEndpoint = mockk(relaxed = true)
-        every { mockTunnelEndpoint.quantumResistant } returns true
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                uiState =
-                    ConnectUiState(
-                        location = null,
-                        relayLocation = null,
-                        tunnelUiState = TunnelState.Connecting(endpoint = mockTunnelEndpoint, null),
-                        tunnelRealState =
-                            TunnelState.Connecting(endpoint = mockTunnelEndpoint, null),
-                        inAddress = null,
-                        outAddress = "",
-                        showLocation = false,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification = InAppNotification.TunnelStateBlocked,
-                        isPlayBuild = false
-                    ),
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            val mockTunnelEndpoint: TunnelEndpoint = mockk(relaxed = true)
+            every { mockTunnelEndpoint.quantumResistant } returns true
+            setContentWithTheme {
+                ConnectScreen(
+                    uiState =
+                        ConnectUiState(
+                            location = null,
+                            relayLocation = null,
+                            tunnelUiState =
+                                TunnelState.Connecting(endpoint = mockTunnelEndpoint, null),
+                            tunnelRealState =
+                                TunnelState.Connecting(endpoint = mockTunnelEndpoint, null),
+                            inAddress = null,
+                            outAddress = "",
+                            showLocation = false,
+                            deviceName = "",
+                            daysLeftUntilExpiry = null,
+                            inAppNotification = InAppNotification.TunnelStateBlocked,
+                            isPlayBuild = false
+                        ),
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithTag(CIRCULAR_PROGRESS_INDICATOR).assertExists()
             onNodeWithText("CREATING QUANTUM SECURE CONNECTION").assertExists()
             onNodeWithText("Switch location").assertExists()
@@ -134,29 +139,29 @@ class ConnectScreenTest {
 
     @Test
     fun testConnectedState() {
-        // Arrange
-        val mockTunnelEndpoint: TunnelEndpoint = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                uiState =
-                    ConnectUiState(
-                        location = null,
-                        relayLocation = null,
-                        tunnelUiState = TunnelState.Connected(mockTunnelEndpoint, null),
-                        tunnelRealState = TunnelState.Connected(mockTunnelEndpoint, null),
-                        inAddress = null,
-                        outAddress = "",
-                        showLocation = false,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification = null,
-                        isPlayBuild = false
-                    ),
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            val mockTunnelEndpoint: TunnelEndpoint = mockk(relaxed = true)
+            setContentWithTheme {
+                ConnectScreen(
+                    uiState =
+                        ConnectUiState(
+                            location = null,
+                            relayLocation = null,
+                            tunnelUiState = TunnelState.Connected(mockTunnelEndpoint, null),
+                            tunnelRealState = TunnelState.Connected(mockTunnelEndpoint, null),
+                            inAddress = null,
+                            outAddress = "",
+                            showLocation = false,
+                            deviceName = "",
+                            daysLeftUntilExpiry = null,
+                            inAppNotification = null,
+                            isPlayBuild = false
+                        ),
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText("SECURE CONNECTION").assertExists()
             onNodeWithText("Switch location").assertExists()
             onNodeWithText("Disconnect").assertExists()
@@ -165,30 +170,30 @@ class ConnectScreenTest {
 
     @Test
     fun testConnectedStateQuantumSecured() {
-        // Arrange
-        val mockTunnelEndpoint: TunnelEndpoint = mockk(relaxed = true)
-        every { mockTunnelEndpoint.quantumResistant } returns true
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                uiState =
-                    ConnectUiState(
-                        location = null,
-                        relayLocation = null,
-                        tunnelUiState = TunnelState.Connected(mockTunnelEndpoint, null),
-                        tunnelRealState = TunnelState.Connected(mockTunnelEndpoint, null),
-                        inAddress = null,
-                        outAddress = "",
-                        showLocation = false,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification = null,
-                        isPlayBuild = false
-                    ),
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            val mockTunnelEndpoint: TunnelEndpoint = mockk(relaxed = true)
+            every { mockTunnelEndpoint.quantumResistant } returns true
+            setContentWithTheme {
+                ConnectScreen(
+                    uiState =
+                        ConnectUiState(
+                            location = null,
+                            relayLocation = null,
+                            tunnelUiState = TunnelState.Connected(mockTunnelEndpoint, null),
+                            tunnelRealState = TunnelState.Connected(mockTunnelEndpoint, null),
+                            inAddress = null,
+                            outAddress = "",
+                            showLocation = false,
+                            deviceName = "",
+                            daysLeftUntilExpiry = null,
+                            inAppNotification = null,
+                            isPlayBuild = false
+                        ),
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText("QUANTUM SECURE CONNECTION").assertExists()
             onNodeWithText("Switch location").assertExists()
             onNodeWithText("Disconnect").assertExists()
@@ -197,31 +202,33 @@ class ConnectScreenTest {
 
     @Test
     fun testDisconnectingState() {
-        // Arrange
-        val mockRelayLocation: RelayItem = mockk(relaxed = true)
-        val mockLocationName = "Home"
-        every { mockRelayLocation.locationName } returns mockLocationName
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                uiState =
-                    ConnectUiState(
-                        location = null,
-                        relayLocation = mockRelayLocation,
-                        tunnelUiState = TunnelState.Disconnecting(ActionAfterDisconnect.Nothing),
-                        tunnelRealState = TunnelState.Disconnecting(ActionAfterDisconnect.Nothing),
-                        inAddress = null,
-                        outAddress = "",
-                        showLocation = true,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification = null,
-                        isPlayBuild = false
-                    ),
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            val mockRelayLocation: RelayItem = mockk(relaxed = true)
+            val mockLocationName = "Home"
+            every { mockRelayLocation.locationName } returns mockLocationName
+            setContentWithTheme {
+                ConnectScreen(
+                    uiState =
+                        ConnectUiState(
+                            location = null,
+                            relayLocation = mockRelayLocation,
+                            tunnelUiState =
+                                TunnelState.Disconnecting(ActionAfterDisconnect.Nothing),
+                            tunnelRealState =
+                                TunnelState.Disconnecting(ActionAfterDisconnect.Nothing),
+                            inAddress = null,
+                            outAddress = "",
+                            showLocation = true,
+                            deviceName = "",
+                            daysLeftUntilExpiry = null,
+                            inAppNotification = null,
+                            isPlayBuild = false
+                        ),
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText("UNSECURED CONNECTION").assertExists()
             onNodeWithText(mockLocationName).assertExists()
             onNodeWithText("Disconnect").assertExists()
@@ -230,31 +237,31 @@ class ConnectScreenTest {
 
     @Test
     fun testDisconnectedState() {
-        // Arrange
-        val mockRelayLocation: RelayItem = mockk(relaxed = true)
-        val mockLocationName = "Home"
-        every { mockRelayLocation.locationName } returns mockLocationName
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                uiState =
-                    ConnectUiState(
-                        location = null,
-                        relayLocation = mockRelayLocation,
-                        tunnelUiState = TunnelState.Disconnected(),
-                        tunnelRealState = TunnelState.Disconnected(),
-                        inAddress = null,
-                        outAddress = "",
-                        showLocation = true,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification = null,
-                        isPlayBuild = false
-                    ),
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            val mockRelayLocation: RelayItem = mockk(relaxed = true)
+            val mockLocationName = "Home"
+            every { mockRelayLocation.locationName } returns mockLocationName
+            setContentWithTheme {
+                ConnectScreen(
+                    uiState =
+                        ConnectUiState(
+                            location = null,
+                            relayLocation = mockRelayLocation,
+                            tunnelUiState = TunnelState.Disconnected(),
+                            tunnelRealState = TunnelState.Disconnected(),
+                            inAddress = null,
+                            outAddress = "",
+                            showLocation = true,
+                            deviceName = "",
+                            daysLeftUntilExpiry = null,
+                            inAppNotification = null,
+                            isPlayBuild = false
+                        ),
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText("UNSECURED CONNECTION").assertExists()
             onNodeWithText(mockLocationName).assertExists()
             onNodeWithText("Secure my connection").assertExists()
@@ -263,36 +270,40 @@ class ConnectScreenTest {
 
     @Test
     fun testErrorStateBlocked() {
-        // Arrange
-        val mockRelayLocation: RelayItem = mockk(relaxed = true)
-        val mockLocationName = "Home"
-        every { mockRelayLocation.locationName } returns mockLocationName
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                uiState =
-                    ConnectUiState(
-                        location = null,
-                        relayLocation = mockRelayLocation,
-                        tunnelUiState =
-                            TunnelState.Error(ErrorState(ErrorStateCause.StartTunnelError, true)),
-                        tunnelRealState =
-                            TunnelState.Error(ErrorState(ErrorStateCause.StartTunnelError, true)),
-                        inAddress = null,
-                        outAddress = "",
-                        showLocation = true,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification =
-                            InAppNotification.TunnelStateError(
-                                ErrorState(ErrorStateCause.StartTunnelError, true)
-                            ),
-                        isPlayBuild = false
-                    ),
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            val mockRelayLocation: RelayItem = mockk(relaxed = true)
+            val mockLocationName = "Home"
+            every { mockRelayLocation.locationName } returns mockLocationName
+            setContentWithTheme {
+                ConnectScreen(
+                    uiState =
+                        ConnectUiState(
+                            location = null,
+                            relayLocation = mockRelayLocation,
+                            tunnelUiState =
+                                TunnelState.Error(
+                                    ErrorState(ErrorStateCause.StartTunnelError, true)
+                                ),
+                            tunnelRealState =
+                                TunnelState.Error(
+                                    ErrorState(ErrorStateCause.StartTunnelError, true)
+                                ),
+                            inAddress = null,
+                            outAddress = "",
+                            showLocation = true,
+                            deviceName = "",
+                            daysLeftUntilExpiry = null,
+                            inAppNotification =
+                                InAppNotification.TunnelStateError(
+                                    ErrorState(ErrorStateCause.StartTunnelError, true)
+                                ),
+                            isPlayBuild = false
+                        ),
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText("BLOCKED CONNECTION").assertExists()
             onNodeWithText(mockLocationName).assertExists()
             onNodeWithText("Disconnect").assertExists()
@@ -302,36 +313,40 @@ class ConnectScreenTest {
 
     @Test
     fun testErrorStateNotBlocked() {
-        // Arrange
-        val mockRelayLocation: RelayItem = mockk(relaxed = true)
-        val mockLocationName = "Home"
-        every { mockRelayLocation.locationName } returns mockLocationName
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                uiState =
-                    ConnectUiState(
-                        location = null,
-                        relayLocation = mockRelayLocation,
-                        tunnelUiState =
-                            TunnelState.Error(ErrorState(ErrorStateCause.StartTunnelError, false)),
-                        tunnelRealState =
-                            TunnelState.Error(ErrorState(ErrorStateCause.StartTunnelError, false)),
-                        inAddress = null,
-                        outAddress = "",
-                        showLocation = true,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification =
-                            InAppNotification.TunnelStateError(
-                                ErrorState(ErrorStateCause.StartTunnelError, false)
-                            ),
-                        isPlayBuild = false
-                    ),
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            val mockRelayLocation: RelayItem = mockk(relaxed = true)
+            val mockLocationName = "Home"
+            every { mockRelayLocation.locationName } returns mockLocationName
+            setContentWithTheme {
+                ConnectScreen(
+                    uiState =
+                        ConnectUiState(
+                            location = null,
+                            relayLocation = mockRelayLocation,
+                            tunnelUiState =
+                                TunnelState.Error(
+                                    ErrorState(ErrorStateCause.StartTunnelError, false)
+                                ),
+                            tunnelRealState =
+                                TunnelState.Error(
+                                    ErrorState(ErrorStateCause.StartTunnelError, false)
+                                ),
+                            inAddress = null,
+                            outAddress = "",
+                            showLocation = true,
+                            deviceName = "",
+                            daysLeftUntilExpiry = null,
+                            inAppNotification =
+                                InAppNotification.TunnelStateError(
+                                    ErrorState(ErrorStateCause.StartTunnelError, false)
+                                ),
+                            isPlayBuild = false
+                        ),
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText("FAILED TO SECURE CONNECTION").assertExists()
             onNodeWithText(mockLocationName).assertExists()
             onNodeWithText("Dismiss").assertExists()
@@ -342,29 +357,30 @@ class ConnectScreenTest {
 
     @Test
     fun testReconnectingState() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                uiState =
-                    ConnectUiState(
-                        location = null,
-                        relayLocation = null,
-                        tunnelUiState = TunnelState.Disconnecting(ActionAfterDisconnect.Reconnect),
-                        tunnelRealState =
-                            TunnelState.Disconnecting(ActionAfterDisconnect.Reconnect),
-                        inAddress = null,
-                        outAddress = "",
-                        showLocation = false,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification = InAppNotification.TunnelStateBlocked,
-                        isPlayBuild = false
-                    ),
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                ConnectScreen(
+                    uiState =
+                        ConnectUiState(
+                            location = null,
+                            relayLocation = null,
+                            tunnelUiState =
+                                TunnelState.Disconnecting(ActionAfterDisconnect.Reconnect),
+                            tunnelRealState =
+                                TunnelState.Disconnecting(ActionAfterDisconnect.Reconnect),
+                            inAddress = null,
+                            outAddress = "",
+                            showLocation = false,
+                            deviceName = "",
+                            daysLeftUntilExpiry = null,
+                            inAppNotification = InAppNotification.TunnelStateBlocked,
+                            isPlayBuild = false
+                        ),
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithTag(CIRCULAR_PROGRESS_INDICATOR).assertExists()
             onNodeWithText("CREATING SECURE CONNECTION").assertExists()
             onNodeWithText("Switch location").assertExists()
@@ -375,31 +391,32 @@ class ConnectScreenTest {
 
     @Test
     fun testDisconnectingBlockState() {
-        // Arrange
-        val mockRelayLocation: RelayItem = mockk(relaxed = true)
-        val mockLocationName = "Home"
-        every { mockRelayLocation.locationName } returns mockLocationName
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                uiState =
-                    ConnectUiState(
-                        location = null,
-                        relayLocation = mockRelayLocation,
-                        tunnelUiState = TunnelState.Disconnecting(ActionAfterDisconnect.Block),
-                        tunnelRealState = TunnelState.Disconnecting(ActionAfterDisconnect.Block),
-                        inAddress = null,
-                        outAddress = "",
-                        showLocation = true,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification = InAppNotification.TunnelStateBlocked,
-                        isPlayBuild = false
-                    ),
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            val mockRelayLocation: RelayItem = mockk(relaxed = true)
+            val mockLocationName = "Home"
+            every { mockRelayLocation.locationName } returns mockLocationName
+            setContentWithTheme {
+                ConnectScreen(
+                    uiState =
+                        ConnectUiState(
+                            location = null,
+                            relayLocation = mockRelayLocation,
+                            tunnelUiState = TunnelState.Disconnecting(ActionAfterDisconnect.Block),
+                            tunnelRealState =
+                                TunnelState.Disconnecting(ActionAfterDisconnect.Block),
+                            inAddress = null,
+                            outAddress = "",
+                            showLocation = true,
+                            deviceName = "",
+                            daysLeftUntilExpiry = null,
+                            inAppNotification = InAppNotification.TunnelStateBlocked,
+                            isPlayBuild = false
+                        ),
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText("SECURE CONNECTION").assertExists()
             onNodeWithText(mockLocationName).assertExists()
             onNodeWithText("Disconnect").assertExists()
@@ -409,199 +426,210 @@ class ConnectScreenTest {
 
     @Test
     fun testClickSelectLocationButton() {
-        // Arrange
-        val mockRelayLocation: RelayItem = mockk(relaxed = true)
-        val mockLocationName = "Home"
-        every { mockRelayLocation.locationName } returns mockLocationName
-        val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                uiState =
-                    ConnectUiState(
-                        location = null,
-                        relayLocation = mockRelayLocation,
-                        tunnelUiState = TunnelState.Disconnected(),
-                        tunnelRealState = TunnelState.Disconnected(),
-                        inAddress = null,
-                        outAddress = "",
-                        showLocation = false,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification = null,
-                        isPlayBuild = false
-                    ),
-                onSwitchLocationClick = mockedClickHandler
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            val mockRelayLocation: RelayItem = mockk(relaxed = true)
+            val mockLocationName = "Home"
+            every { mockRelayLocation.locationName } returns mockLocationName
+            val mockedClickHandler: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                ConnectScreen(
+                    uiState =
+                        ConnectUiState(
+                            location = null,
+                            relayLocation = mockRelayLocation,
+                            tunnelUiState = TunnelState.Disconnected(),
+                            tunnelRealState = TunnelState.Disconnected(),
+                            inAddress = null,
+                            outAddress = "",
+                            showLocation = false,
+                            deviceName = "",
+                            daysLeftUntilExpiry = null,
+                            inAppNotification = null,
+                            isPlayBuild = false
+                        ),
+                    onSwitchLocationClick = mockedClickHandler
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithTag(SELECT_LOCATION_BUTTON_TEST_TAG).performClick()
+            // Act
+            onNodeWithTag(SELECT_LOCATION_BUTTON_TEST_TAG).performClick()
 
-        // Assert
-        verify { mockedClickHandler.invoke() }
+            // Assert
+            verify { mockedClickHandler.invoke() }
+        }
     }
 
     @Test
     fun testOnDisconnectClick() {
-        // Arrange
-        val mockTunnelEndpoint: TunnelEndpoint = mockk(relaxed = true)
-        val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                uiState =
-                    ConnectUiState(
-                        location = null,
-                        relayLocation = null,
-                        tunnelUiState = TunnelState.Connected(mockTunnelEndpoint, null),
-                        tunnelRealState = TunnelState.Connected(mockTunnelEndpoint, null),
-                        inAddress = null,
-                        outAddress = "",
-                        showLocation = false,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification = null,
-                        isPlayBuild = false
-                    ),
-                onDisconnectClick = mockedClickHandler
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            val mockTunnelEndpoint: TunnelEndpoint = mockk(relaxed = true)
+            val mockedClickHandler: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                ConnectScreen(
+                    uiState =
+                        ConnectUiState(
+                            location = null,
+                            relayLocation = null,
+                            tunnelUiState = TunnelState.Connected(mockTunnelEndpoint, null),
+                            tunnelRealState = TunnelState.Connected(mockTunnelEndpoint, null),
+                            inAddress = null,
+                            outAddress = "",
+                            showLocation = false,
+                            deviceName = "",
+                            daysLeftUntilExpiry = null,
+                            inAppNotification = null,
+                            isPlayBuild = false
+                        ),
+                    onDisconnectClick = mockedClickHandler
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithTag(CONNECT_BUTTON_TEST_TAG).performClick()
+            // Act
+            onNodeWithTag(CONNECT_BUTTON_TEST_TAG).performClick()
 
-        // Assert
-        verify { mockedClickHandler.invoke() }
+            // Assert
+            verify { mockedClickHandler.invoke() }
+        }
     }
 
     @Test
     fun testOnReconnectClick() {
-        // Arrange
-        val mockTunnelEndpoint: TunnelEndpoint = mockk(relaxed = true)
-        val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                uiState =
-                    ConnectUiState(
-                        location = null,
-                        relayLocation = null,
-                        tunnelUiState = TunnelState.Connected(mockTunnelEndpoint, null),
-                        tunnelRealState = TunnelState.Connected(mockTunnelEndpoint, null),
-                        inAddress = null,
-                        outAddress = "",
-                        showLocation = false,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification = null,
-                        isPlayBuild = false
-                    ),
-                onReconnectClick = mockedClickHandler
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            val mockTunnelEndpoint: TunnelEndpoint = mockk(relaxed = true)
+            val mockedClickHandler: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                ConnectScreen(
+                    uiState =
+                        ConnectUiState(
+                            location = null,
+                            relayLocation = null,
+                            tunnelUiState = TunnelState.Connected(mockTunnelEndpoint, null),
+                            tunnelRealState = TunnelState.Connected(mockTunnelEndpoint, null),
+                            inAddress = null,
+                            outAddress = "",
+                            showLocation = false,
+                            deviceName = "",
+                            daysLeftUntilExpiry = null,
+                            inAppNotification = null,
+                            isPlayBuild = false
+                        ),
+                    onReconnectClick = mockedClickHandler
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithTag(RECONNECT_BUTTON_TEST_TAG).performClick()
+            // Act
+            onNodeWithTag(RECONNECT_BUTTON_TEST_TAG).performClick()
 
-        // Assert
-        verify { mockedClickHandler.invoke() }
+            // Assert
+            verify { mockedClickHandler.invoke() }
+        }
     }
 
     @Test
     fun testOnConnectClick() {
-        // Arrange
-        val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                uiState =
-                    ConnectUiState(
-                        location = null,
-                        relayLocation = null,
-                        tunnelUiState = TunnelState.Disconnected(),
-                        tunnelRealState = TunnelState.Disconnected(),
-                        inAddress = null,
-                        outAddress = "",
-                        showLocation = false,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification = null,
-                        isPlayBuild = false
-                    ),
-                onConnectClick = mockedClickHandler
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            val mockedClickHandler: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                ConnectScreen(
+                    uiState =
+                        ConnectUiState(
+                            location = null,
+                            relayLocation = null,
+                            tunnelUiState = TunnelState.Disconnected(),
+                            tunnelRealState = TunnelState.Disconnected(),
+                            inAddress = null,
+                            outAddress = "",
+                            showLocation = false,
+                            deviceName = "",
+                            daysLeftUntilExpiry = null,
+                            inAppNotification = null,
+                            isPlayBuild = false
+                        ),
+                    onConnectClick = mockedClickHandler
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithTag(CONNECT_BUTTON_TEST_TAG).performClick()
+            // Act
+            onNodeWithTag(CONNECT_BUTTON_TEST_TAG).performClick()
 
-        // Assert
-        verify { mockedClickHandler.invoke() }
+            // Assert
+            verify { mockedClickHandler.invoke() }
+        }
     }
 
     @Test
     fun testOnCancelClick() {
-        // Arrange
-        val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                uiState =
-                    ConnectUiState(
-                        location = null,
-                        relayLocation = null,
-                        tunnelUiState = TunnelState.Connecting(null, null),
-                        tunnelRealState = TunnelState.Connecting(null, null),
-                        inAddress = null,
-                        outAddress = "",
-                        showLocation = false,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification = null,
-                        isPlayBuild = false
-                    ),
-                onCancelClick = mockedClickHandler
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            val mockedClickHandler: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                ConnectScreen(
+                    uiState =
+                        ConnectUiState(
+                            location = null,
+                            relayLocation = null,
+                            tunnelUiState = TunnelState.Connecting(null, null),
+                            tunnelRealState = TunnelState.Connecting(null, null),
+                            inAddress = null,
+                            outAddress = "",
+                            showLocation = false,
+                            deviceName = "",
+                            daysLeftUntilExpiry = null,
+                            inAppNotification = null,
+                            isPlayBuild = false
+                        ),
+                    onCancelClick = mockedClickHandler
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithTag(CONNECT_BUTTON_TEST_TAG).performClick()
+            // Act
+            onNodeWithTag(CONNECT_BUTTON_TEST_TAG).performClick()
 
-        // Assert
-        verify { mockedClickHandler.invoke() }
+            // Assert
+            verify { mockedClickHandler.invoke() }
+        }
     }
 
     @Test
     fun showLocationInfo() {
-        // Arrange
-        val mockLocation: GeoIpLocation = mockk(relaxed = true)
-        val mockTunnelEndpoint: TunnelEndpoint = mockk(relaxed = true)
-        val mockHostName = "Host-Name"
-        val mockPort = 99
-        val mockHost = "Host"
-        val mockProtocol = TransportProtocol.Udp
-        val mockInAddress = Triple(mockHost, mockPort, mockProtocol)
-        val mockOutAddress = "HostAddressV4 / HostAddressV4"
-        every { mockLocation.hostname } returns mockHostName
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                uiState =
-                    ConnectUiState(
-                        location = mockLocation,
-                        relayLocation = null,
-                        tunnelUiState = TunnelState.Connected(mockTunnelEndpoint, null),
-                        tunnelRealState = TunnelState.Connected(mockTunnelEndpoint, null),
-                        inAddress = mockInAddress,
-                        outAddress = mockOutAddress,
-                        showLocation = false,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification = null,
-                        isPlayBuild = false
-                    ),
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            val mockLocation: GeoIpLocation = mockk(relaxed = true)
+            val mockTunnelEndpoint: TunnelEndpoint = mockk(relaxed = true)
+            val mockHostName = "Host-Name"
+            val mockPort = 99
+            val mockHost = "Host"
+            val mockProtocol = TransportProtocol.Udp
+            val mockInAddress = Triple(mockHost, mockPort, mockProtocol)
+            val mockOutAddress = "HostAddressV4 / HostAddressV4"
+            every { mockLocation.hostname } returns mockHostName
+            setContentWithTheme {
+                ConnectScreen(
+                    uiState =
+                        ConnectUiState(
+                            location = mockLocation,
+                            relayLocation = null,
+                            tunnelUiState = TunnelState.Connected(mockTunnelEndpoint, null),
+                            tunnelRealState = TunnelState.Connected(mockTunnelEndpoint, null),
+                            inAddress = mockInAddress,
+                            outAddress = mockOutAddress,
+                            showLocation = false,
+                            deviceName = "",
+                            daysLeftUntilExpiry = null,
+                            inAppNotification = null,
+                            isPlayBuild = false
+                        ),
+                )
+            }
 
-        composeTestRule.onNodeWithTag(LOCATION_INFO_TEST_TAG).performClick()
+            // Act
+            onNodeWithTag(LOCATION_INFO_TEST_TAG).performClick()
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText(mockHostName).assertExists()
             onNodeWithText("WireGuard").assertExists()
             onNodeWithText("In $mockHost:$mockPort UDP").assertExists()
@@ -611,35 +639,35 @@ class ConnectScreenTest {
 
     @Test
     fun testOutdatedVersionNotification() {
-        // Arrange
-        val versionInfo =
-            VersionInfo(
-                currentVersion = "1.0",
-                upgradeVersion = "1.1",
-                isOutdated = true,
-                isSupported = true
-            )
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                uiState =
-                    ConnectUiState(
-                        location = null,
-                        relayLocation = null,
-                        tunnelUiState = TunnelState.Connecting(null, null),
-                        tunnelRealState = TunnelState.Connecting(null, null),
-                        inAddress = null,
-                        outAddress = "",
-                        showLocation = false,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification = InAppNotification.UpdateAvailable(versionInfo),
-                        isPlayBuild = false
-                    ),
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            val versionInfo =
+                VersionInfo(
+                    currentVersion = "1.0",
+                    upgradeVersion = "1.1",
+                    isOutdated = true,
+                    isSupported = true
+                )
+            setContentWithTheme {
+                ConnectScreen(
+                    uiState =
+                        ConnectUiState(
+                            location = null,
+                            relayLocation = null,
+                            tunnelUiState = TunnelState.Connecting(null, null),
+                            tunnelRealState = TunnelState.Connecting(null, null),
+                            inAddress = null,
+                            outAddress = "",
+                            showLocation = false,
+                            deviceName = "",
+                            daysLeftUntilExpiry = null,
+                            inAppNotification = InAppNotification.UpdateAvailable(versionInfo),
+                            isPlayBuild = false
+                        ),
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText("UPDATE AVAILABLE").assertExists()
             onNodeWithText("Install Mullvad VPN (1.1) to stay up to date").assertExists()
         }
@@ -647,35 +675,35 @@ class ConnectScreenTest {
 
     @Test
     fun testUnsupportedVersionNotification() {
-        // Arrange
-        val versionInfo =
-            VersionInfo(
-                currentVersion = "1.0",
-                upgradeVersion = "1.1",
-                isOutdated = true,
-                isSupported = false
-            )
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                uiState =
-                    ConnectUiState(
-                        location = null,
-                        relayLocation = null,
-                        tunnelUiState = TunnelState.Connecting(null, null),
-                        tunnelRealState = TunnelState.Connecting(null, null),
-                        inAddress = null,
-                        outAddress = "",
-                        showLocation = false,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification = InAppNotification.UnsupportedVersion(versionInfo),
-                        isPlayBuild = false
-                    ),
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            val versionInfo =
+                VersionInfo(
+                    currentVersion = "1.0",
+                    upgradeVersion = "1.1",
+                    isOutdated = true,
+                    isSupported = false
+                )
+            setContentWithTheme {
+                ConnectScreen(
+                    uiState =
+                        ConnectUiState(
+                            location = null,
+                            relayLocation = null,
+                            tunnelUiState = TunnelState.Connecting(null, null),
+                            tunnelRealState = TunnelState.Connecting(null, null),
+                            inAddress = null,
+                            outAddress = "",
+                            showLocation = false,
+                            deviceName = "",
+                            daysLeftUntilExpiry = null,
+                            inAppNotification = InAppNotification.UnsupportedVersion(versionInfo),
+                            isPlayBuild = false
+                        ),
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText("UNSUPPORTED VERSION").assertExists()
             onNodeWithText(
                     "Your privacy might be at risk with this unsupported app version. Please update now."
@@ -686,29 +714,29 @@ class ConnectScreenTest {
 
     @Test
     fun testAccountExpiredNotification() {
-        // Arrange
-        val expiryDate = DateTime(2020, 11, 11, 10, 10)
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                uiState =
-                    ConnectUiState(
-                        location = null,
-                        relayLocation = null,
-                        tunnelUiState = TunnelState.Connecting(null, null),
-                        tunnelRealState = TunnelState.Connecting(null, null),
-                        inAddress = null,
-                        outAddress = "",
-                        showLocation = false,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification = InAppNotification.AccountExpiry(expiryDate),
-                        isPlayBuild = false
-                    ),
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            val expiryDate = DateTime(2020, 11, 11, 10, 10)
+            setContentWithTheme {
+                ConnectScreen(
+                    uiState =
+                        ConnectUiState(
+                            location = null,
+                            relayLocation = null,
+                            tunnelUiState = TunnelState.Connecting(null, null),
+                            tunnelRealState = TunnelState.Connecting(null, null),
+                            inAddress = null,
+                            outAddress = "",
+                            showLocation = false,
+                            deviceName = "",
+                            daysLeftUntilExpiry = null,
+                            inAppNotification = InAppNotification.AccountExpiry(expiryDate),
+                            isPlayBuild = false
+                        ),
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText("ACCOUNT CREDIT EXPIRES SOON").assertExists()
             onNodeWithText("Out of time").assertExists()
         }
@@ -716,85 +744,94 @@ class ConnectScreenTest {
 
     @Test
     fun testOnUpdateVersionClick() {
-        // Arrange
-        val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-        val versionInfo =
-            VersionInfo(
-                currentVersion = "1.0",
-                upgradeVersion = "1.1",
-                isOutdated = true,
-                isSupported = false
-            )
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                onUpdateVersionClick = mockedClickHandler,
-                uiState =
-                    ConnectUiState(
-                        location = null,
-                        relayLocation = null,
-                        tunnelUiState = TunnelState.Connecting(null, null),
-                        tunnelRealState = TunnelState.Connecting(null, null),
-                        inAddress = null,
-                        outAddress = "",
-                        showLocation = false,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification = InAppNotification.UnsupportedVersion(versionInfo),
-                        isPlayBuild = false
-                    ),
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            val mockedClickHandler: () -> Unit = mockk(relaxed = true)
+            val versionInfo =
+                VersionInfo(
+                    currentVersion = "1.0",
+                    upgradeVersion = "1.1",
+                    isOutdated = true,
+                    isSupported = false
+                )
+            setContentWithTheme {
+                ConnectScreen(
+                    onUpdateVersionClick = mockedClickHandler,
+                    uiState =
+                        ConnectUiState(
+                            location = null,
+                            relayLocation = null,
+                            tunnelUiState = TunnelState.Connecting(null, null),
+                            tunnelRealState = TunnelState.Connecting(null, null),
+                            inAddress = null,
+                            outAddress = "",
+                            showLocation = false,
+                            deviceName = "",
+                            daysLeftUntilExpiry = null,
+                            inAppNotification = InAppNotification.UnsupportedVersion(versionInfo),
+                            isPlayBuild = false
+                        ),
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithTag(NOTIFICATION_BANNER_ACTION).performClick()
+            // Act
+            onNodeWithTag(NOTIFICATION_BANNER_ACTION).performClick()
 
-        // Assert
-        verify { mockedClickHandler.invoke() }
+            // Assert
+            verify { mockedClickHandler.invoke() }
+        }
     }
 
     @Test
     fun testOnShowAccountClick() {
-        // Arrange
-        val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-        val expiryDate = DateTime(2020, 11, 11, 10, 10)
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                onManageAccountClick = mockedClickHandler,
-                uiState =
-                    ConnectUiState(
-                        location = null,
-                        relayLocation = null,
-                        tunnelUiState = TunnelState.Connecting(null, null),
-                        tunnelRealState = TunnelState.Connecting(null, null),
-                        inAddress = null,
-                        outAddress = "",
-                        showLocation = false,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification = InAppNotification.AccountExpiry(expiryDate),
-                        isPlayBuild = false
-                    ),
-            )
-        }
+        composeExtension.use {
+            // Arrange
+            val mockedClickHandler: () -> Unit = mockk(relaxed = true)
+            val expiryDate = DateTime(2020, 11, 11, 10, 10)
+            setContentWithTheme {
+                ConnectScreen(
+                    onManageAccountClick = mockedClickHandler,
+                    uiState =
+                        ConnectUiState(
+                            location = null,
+                            relayLocation = null,
+                            tunnelUiState = TunnelState.Connecting(null, null),
+                            tunnelRealState = TunnelState.Connecting(null, null),
+                            inAddress = null,
+                            outAddress = "",
+                            showLocation = false,
+                            deviceName = "",
+                            daysLeftUntilExpiry = null,
+                            inAppNotification = InAppNotification.AccountExpiry(expiryDate),
+                            isPlayBuild = false
+                        ),
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithTag(NOTIFICATION_BANNER_ACTION).performClick()
+            // Act
+            onNodeWithTag(NOTIFICATION_BANNER_ACTION).performClick()
 
-        // Assert
-        verify { mockedClickHandler.invoke() }
+            // Assert
+            verify { mockedClickHandler.invoke() }
+        }
     }
 
     @Test
     fun testOpenAccountView() {
-        // Arrange
-        val onAccountClickMockk: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(uiState = ConnectUiState.INITIAL, onAccountClick = onAccountClickMockk)
-        }
+        composeExtension.use {
+            // Arrange
+            val onAccountClickMockk: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                ConnectScreen(
+                    uiState = ConnectUiState.INITIAL,
+                    onAccountClick = onAccountClickMockk
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithTag(TOP_BAR_ACCOUNT_BUTTON).performClick()
+            // Assert
+            onNodeWithTag(TOP_BAR_ACCOUNT_BUTTON).performClick()
 
-        verify(exactly = 1) { onAccountClickMockk() }
+            verify(exactly = 1) { onAccountClickMockk() }
+        }
     }
 }
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/DeviceRevokedScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/DeviceRevokedScreenTest.kt
index 56e24b9a08..f65a23fa9c 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/DeviceRevokedScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/DeviceRevokedScreenTest.kt
@@ -1,64 +1,69 @@
 package net.mullvad.mullvadvpn.compose.screen
 
-import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
+import de.mannodermaus.junit5.compose.createComposeExtension
 import io.mockk.MockKAnnotations
 import io.mockk.mockk
 import io.mockk.verify
 import net.mullvad.mullvadvpn.compose.setContentWithTheme
 import net.mullvad.mullvadvpn.compose.state.DeviceRevokedUiState
-import org.junit.Before
-import org.junit.Rule
-import org.junit.Test
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.RegisterExtension
 
+@OptIn(ExperimentalTestApi::class)
 class DeviceRevokedScreenTest {
-    @get:Rule val composeTestRule = createComposeRule()
+    @JvmField @RegisterExtension val composeExtension = createComposeExtension()
 
-    @Before
+    @BeforeEach
     fun setup() {
         MockKAnnotations.init(this)
     }
 
     @Test
-    fun testUnblockWarningShowingWhenSecured() {
-        // Arrange
-        val state = DeviceRevokedUiState.SECURED
+    fun testUnblockWarningShowingWhenSecured() =
+        composeExtension.use {
+            // Arrange
+            val state = DeviceRevokedUiState.SECURED
 
-        // Act
-        composeTestRule.setContentWithTheme { DeviceRevokedScreen(state) }
+            // Act
+            setContentWithTheme { DeviceRevokedScreen(state) }
 
-        // Assert
-        composeTestRule.onNodeWithText(UNBLOCK_WARNING).assertExists()
-    }
+            // Assert
+            onNodeWithText(UNBLOCK_WARNING).assertExists()
+        }
 
     @Test
-    fun testUnblockWarningNotShowingWhenNotSecured() {
-        // Arrange
-        val state = DeviceRevokedUiState.UNSECURED
+    fun testUnblockWarningNotShowingWhenNotSecured() =
+        composeExtension.use {
+            // Arrange
+            val state = DeviceRevokedUiState.UNSECURED
 
-        // Act
-        composeTestRule.setContentWithTheme { DeviceRevokedScreen(state) }
+            // Act
+            setContentWithTheme { DeviceRevokedScreen(state) }
 
-        // Assert
-        composeTestRule.onNodeWithText(UNBLOCK_WARNING).assertDoesNotExist()
-    }
+            // Assert
+            onNodeWithText(UNBLOCK_WARNING).assertDoesNotExist()
+        }
 
     @Test
-    fun testGoToLogin() {
-        // Arrange
-        val state = DeviceRevokedUiState.UNSECURED
-        val mockOnGoToLoginClicked: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            DeviceRevokedScreen(state = state, onGoToLoginClicked = mockOnGoToLoginClicked)
-        }
+    fun testGoToLogin() =
+        composeExtension.use {
+            // Arrange
+            val state = DeviceRevokedUiState.UNSECURED
+            val mockOnGoToLoginClicked: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                DeviceRevokedScreen(state = state, onGoToLoginClicked = mockOnGoToLoginClicked)
+            }
 
-        // Act
-        composeTestRule.onNodeWithText(GO_TO_LOGIN_BUTTON_TEXT).performClick()
+            // Act
+            onNodeWithText(GO_TO_LOGIN_BUTTON_TEXT).performClick()
 
-        // Assert
-        verify { mockOnGoToLoginClicked.invoke() }
-    }
+            // Assert
+            verify { mockOnGoToLoginClicked.invoke() }
+        }
 
     companion object {
         private const val GO_TO_LOGIN_BUTTON_TEXT = "Go to login"
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/FilterScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/FilterScreenTest.kt
index 32fd727329..11f15b7823 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/FilterScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/FilterScreenTest.kt
@@ -1,8 +1,9 @@
 package net.mullvad.mullvadvpn.compose.screen
 
-import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
+import de.mannodermaus.junit5.compose.createComposeExtension
 import io.mockk.MockKAnnotations
 import io.mockk.mockk
 import io.mockk.verify
@@ -10,132 +11,128 @@ import net.mullvad.mullvadvpn.compose.setContentWithTheme
 import net.mullvad.mullvadvpn.compose.state.RelayFilterState
 import net.mullvad.mullvadvpn.model.Ownership
 import net.mullvad.mullvadvpn.relaylist.Provider
-import org.junit.Rule
-import org.junit.Test
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.RegisterExtension
 
+@OptIn(ExperimentalTestApi::class)
 class FilterScreenTest {
-    @get:Rule val composeTestRule = createComposeRule()
+    @JvmField @RegisterExtension val composeExtension = createComposeExtension()
 
     fun setup() {
         MockKAnnotations.init(this)
     }
 
     @Test
-    fun testDefaultState() {
-        composeTestRule.setContentWithTheme {
-            FilterScreen(
-                uiState =
-                    RelayFilterState(
-                        allProviders = DUMMY_RELAY_ALL_PROVIDERS,
-                        selectedOwnership = null,
-                        selectedProviders = DUMMY_SELECTED_PROVIDERS,
-                    ),
-                onSelectedProvider = { _, _ -> }
-            )
-        }
-        composeTestRule.apply {
+    fun testDefaultState() =
+        composeExtension.use {
+            setContentWithTheme {
+                FilterScreen(
+                    uiState =
+                        RelayFilterState(
+                            allProviders = DUMMY_RELAY_ALL_PROVIDERS,
+                            selectedOwnership = null,
+                            selectedProviders = DUMMY_SELECTED_PROVIDERS,
+                        ),
+                    onSelectedProvider = { _, _ -> }
+                )
+            }
             onNodeWithText("Ownership").assertExists()
             onNodeWithText("Providers").assertExists()
         }
-    }
 
     @Test
-    fun testIsAnyCellShowing() {
-        composeTestRule.setContentWithTheme {
-            FilterScreen(
-                uiState =
-                    RelayFilterState(
-                        allProviders = DUMMY_RELAY_ALL_PROVIDERS,
-                        selectedOwnership = null,
-                        selectedProviders = DUMMY_SELECTED_PROVIDERS
-                    ),
-                onSelectedProvider = { _, _ -> }
-            )
-        }
-        composeTestRule.apply {
+    fun testIsAnyCellShowing() =
+        composeExtension.use {
+            setContentWithTheme {
+                FilterScreen(
+                    uiState =
+                        RelayFilterState(
+                            allProviders = DUMMY_RELAY_ALL_PROVIDERS,
+                            selectedOwnership = null,
+                            selectedProviders = DUMMY_SELECTED_PROVIDERS
+                        ),
+                    onSelectedProvider = { _, _ -> }
+                )
+            }
             onNodeWithText("Ownership").performClick()
             onNodeWithText("Any").assertExists()
         }
-    }
 
     @Test
-    fun testIsMullvadCellShowing() {
-        composeTestRule.setContentWithTheme {
-            FilterScreen(
-                uiState =
-                    RelayFilterState(
-                        allProviders = DUMMY_RELAY_ALL_PROVIDERS,
-                        selectedOwnership = Ownership.MullvadOwned,
-                        selectedProviders = DUMMY_SELECTED_PROVIDERS
-                    ),
-                onSelectedProvider = { _, _ -> }
-            )
-        }
-        composeTestRule.apply {
+    fun testIsMullvadCellShowing() =
+        composeExtension.use {
+            setContentWithTheme {
+                FilterScreen(
+                    uiState =
+                        RelayFilterState(
+                            allProviders = DUMMY_RELAY_ALL_PROVIDERS,
+                            selectedOwnership = Ownership.MullvadOwned,
+                            selectedProviders = DUMMY_SELECTED_PROVIDERS
+                        ),
+                    onSelectedProvider = { _, _ -> }
+                )
+            }
             onNodeWithText("Ownership").performClick()
             onNodeWithText("Mullvad owned only").assertExists()
         }
-    }
 
     @Test
-    fun testIsRentedCellShowing() {
-        composeTestRule.setContentWithTheme {
-            FilterScreen(
-                uiState =
-                    RelayFilterState(
-                        allProviders = DUMMY_RELAY_ALL_PROVIDERS,
-                        selectedOwnership = Ownership.Rented,
-                        selectedProviders = DUMMY_SELECTED_PROVIDERS
-                    ),
-                onSelectedProvider = { _, _ -> }
-            )
-        }
-        composeTestRule.apply {
+    fun testIsRentedCellShowing() =
+        composeExtension.use {
+            setContentWithTheme {
+                FilterScreen(
+                    uiState =
+                        RelayFilterState(
+                            allProviders = DUMMY_RELAY_ALL_PROVIDERS,
+                            selectedOwnership = Ownership.Rented,
+                            selectedProviders = DUMMY_SELECTED_PROVIDERS
+                        ),
+                    onSelectedProvider = { _, _ -> }
+                )
+            }
             onNodeWithText("Ownership").performClick()
             onNodeWithText("Rented only").assertExists()
         }
-    }
 
     @Test
-    fun testShowProviders() {
-        composeTestRule.setContentWithTheme {
-            FilterScreen(
-                uiState =
-                    RelayFilterState(
-                        allProviders = DUMMY_RELAY_ALL_PROVIDERS,
-                        selectedOwnership = null,
-                        selectedProviders = DUMMY_SELECTED_PROVIDERS
-                    ),
-                onSelectedProvider = { _, _ -> }
-            )
-        }
+    fun testShowProviders() =
+        composeExtension.use {
+            setContentWithTheme {
+                FilterScreen(
+                    uiState =
+                        RelayFilterState(
+                            allProviders = DUMMY_RELAY_ALL_PROVIDERS,
+                            selectedOwnership = null,
+                            selectedProviders = DUMMY_SELECTED_PROVIDERS
+                        ),
+                    onSelectedProvider = { _, _ -> }
+                )
+            }
 
-        composeTestRule.apply {
             onNodeWithText("Providers").performClick()
             onNodeWithText("Creanova").assertExists()
-            onNodeWithText("Creanova").assertExists()
             onNodeWithText("100TB").assertExists()
         }
-    }
 
     @Test
-    fun testApplyButtonClick() {
-        val mockClickListener: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            FilterScreen(
-                uiState =
-                    RelayFilterState(
-                        allProviders = listOf(),
-                        selectedOwnership = null,
-                        selectedProviders = listOf(Provider("31173", true))
-                    ),
-                onSelectedProvider = { _, _ -> },
-                onApplyClick = mockClickListener
-            )
+    fun testApplyButtonClick() =
+        composeExtension.use {
+            val mockClickListener: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                FilterScreen(
+                    uiState =
+                        RelayFilterState(
+                            allProviders = listOf(),
+                            selectedOwnership = null,
+                            selectedProviders = listOf(Provider("31173", true))
+                        ),
+                    onSelectedProvider = { _, _ -> },
+                    onApplyClick = mockClickListener
+                )
+            }
+            onNodeWithText("Apply").performClick()
+            verify { mockClickListener() }
         }
-        composeTestRule.onNodeWithText("Apply").performClick()
-        verify { mockClickListener() }
-    }
 
     companion object {
 
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/OutOfTimeScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/OutOfTimeScreenTest.kt
index d43a0931a1..f3f1d0cba5 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/OutOfTimeScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/OutOfTimeScreenTest.kt
@@ -1,9 +1,10 @@
 package net.mullvad.mullvadvpn.compose.screen
 
-import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.onNodeWithTag
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
+import de.mannodermaus.junit5.compose.createComposeExtension
 import io.mockk.MockKAnnotations
 import io.mockk.every
 import io.mockk.mockk
@@ -17,34 +18,35 @@ import net.mullvad.mullvadvpn.lib.payment.model.PaymentStatus
 import net.mullvad.mullvadvpn.lib.payment.model.ProductId
 import net.mullvad.mullvadvpn.lib.payment.model.ProductPrice
 import net.mullvad.mullvadvpn.model.TunnelState
-import org.junit.Before
-import org.junit.Rule
-import org.junit.Test
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.RegisterExtension
 
+@OptIn(ExperimentalTestApi::class)
 class OutOfTimeScreenTest {
-    @get:Rule val composeTestRule = createComposeRule()
+    @JvmField @RegisterExtension val composeExtension = createComposeExtension()
 
-    @Before
+    @BeforeEach
     fun setup() {
         MockKAnnotations.init(this)
     }
 
     @Test
-    fun testDisableSitePayment() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            OutOfTimeScreen(
-                uiState = OutOfTimeUiState(deviceName = ""),
-                onSitePaymentClick = {},
-                onRedeemVoucherClick = {},
-                onSettingsClick = {},
-                onAccountClick = {},
-                onDisconnectClick = {}
-            )
-        }
+    fun testDisableSitePayment() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                OutOfTimeScreen(
+                    uiState = OutOfTimeUiState(deviceName = ""),
+                    onSitePaymentClick = {},
+                    onRedeemVoucherClick = {},
+                    onSettingsClick = {},
+                    onAccountClick = {},
+                    onDisconnectClick = {}
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText(
                     "Either buy credit on our website or redeem a voucher.",
                     substring = true
@@ -52,241 +54,251 @@ class OutOfTimeScreenTest {
                 .assertDoesNotExist()
             onNodeWithText("Buy credit").assertDoesNotExist()
         }
-    }
 
     @Test
-    fun testOpenAccountView() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            OutOfTimeScreen(
-                uiState = OutOfTimeUiState(deviceName = "", showSitePayment = true),
-                onSitePaymentClick = {},
-                onRedeemVoucherClick = {},
-                onSettingsClick = {},
-                onAccountClick = {},
-                onDisconnectClick = {}
-            )
-        }
+    fun testOpenAccountView() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                OutOfTimeScreen(
+                    uiState = OutOfTimeUiState(deviceName = "", showSitePayment = true),
+                    onSitePaymentClick = {},
+                    onRedeemVoucherClick = {},
+                    onSettingsClick = {},
+                    onAccountClick = {},
+                    onDisconnectClick = {}
+                )
+            }
 
-        // Assert
-        composeTestRule.apply { onNodeWithText("Congrats!").assertDoesNotExist() }
-    }
+            // Assert
+            onNodeWithText("Congrats!").assertDoesNotExist()
+        }
 
     @Test
-    fun testClickSitePaymentButton() {
-        // Arrange
-        val mockClickListener: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            OutOfTimeScreen(
-                uiState = OutOfTimeUiState(deviceName = "", showSitePayment = true),
-                onSitePaymentClick = mockClickListener,
-                onRedeemVoucherClick = {},
-                onSettingsClick = {},
-                onAccountClick = {},
-                onDisconnectClick = {}
-            )
-        }
+    fun testClickSitePaymentButton() =
+        composeExtension.use {
+            // Arrange
+            val mockClickListener: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                OutOfTimeScreen(
+                    uiState = OutOfTimeUiState(deviceName = "", showSitePayment = true),
+                    onSitePaymentClick = mockClickListener,
+                    onRedeemVoucherClick = {},
+                    onSettingsClick = {},
+                    onAccountClick = {},
+                    onDisconnectClick = {}
+                )
+            }
 
-        // Act
-        composeTestRule.apply { onNodeWithText("Buy credit").performClick() }
+            // Act
+            onNodeWithText("Buy credit").performClick()
 
-        // Assert
-        verify(exactly = 1) { mockClickListener.invoke() }
-    }
+            // Assert
+            verify(exactly = 1) { mockClickListener.invoke() }
+        }
 
     @Test
-    fun testClickRedeemVoucher() {
-        // Arrange
-        val mockClickListener: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            OutOfTimeScreen(
-                uiState = OutOfTimeUiState(deviceName = "", showSitePayment = true),
-                onSitePaymentClick = {},
-                onRedeemVoucherClick = mockClickListener,
-                onSettingsClick = {},
-                onAccountClick = {},
-                onDisconnectClick = {}
-            )
-        }
+    fun testClickRedeemVoucher() =
+        composeExtension.use {
+            // Arrange
+            val mockClickListener: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                OutOfTimeScreen(
+                    uiState = OutOfTimeUiState(deviceName = "", showSitePayment = true),
+                    onSitePaymentClick = {},
+                    onRedeemVoucherClick = mockClickListener,
+                    onSettingsClick = {},
+                    onAccountClick = {},
+                    onDisconnectClick = {}
+                )
+            }
 
-        // Act
-        composeTestRule.apply { onNodeWithText("Redeem voucher").performClick() }
+            // Act
+            onNodeWithText("Redeem voucher").performClick()
 
-        // Assert
-        verify(exactly = 1) { mockClickListener.invoke() }
-    }
+            // Assert
+            verify(exactly = 1) { mockClickListener.invoke() }
+        }
 
     @Test
-    fun testClickDisconnect() {
-        // Arrange
-        val mockClickListener: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            OutOfTimeScreen(
-                uiState =
-                    OutOfTimeUiState(
-                        tunnelState = TunnelState.Connecting(null, null),
-                        deviceName = "",
-                        showSitePayment = true
-                    ),
-                onSitePaymentClick = {},
-                onRedeemVoucherClick = {},
-                onSettingsClick = {},
-                onAccountClick = {},
-                onDisconnectClick = mockClickListener
-            )
-        }
+    fun testClickDisconnect() =
+        composeExtension.use {
+            // Arrange
+            val mockClickListener: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                OutOfTimeScreen(
+                    uiState =
+                        OutOfTimeUiState(
+                            tunnelState = TunnelState.Connecting(null, null),
+                            deviceName = "",
+                            showSitePayment = true
+                        ),
+                    onSitePaymentClick = {},
+                    onRedeemVoucherClick = {},
+                    onSettingsClick = {},
+                    onAccountClick = {},
+                    onDisconnectClick = mockClickListener
+                )
+            }
 
-        // Act
-        composeTestRule.apply { onNodeWithText("Disconnect").performClick() }
+            // Act
+            onNodeWithText("Disconnect").performClick()
 
-        // Assert
-        verify(exactly = 1) { mockClickListener.invoke() }
-    }
+            // Assert
+            verify(exactly = 1) { mockClickListener.invoke() }
+        }
 
     @Test
-    fun testShowBillingErrorPaymentButton() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            OutOfTimeScreen(
-                uiState =
-                    OutOfTimeUiState(
-                        showSitePayment = true,
-                        billingPaymentState = PaymentState.Error.Billing
-                    ),
-                onSitePaymentClick = {},
-                onRedeemVoucherClick = {},
-                onSettingsClick = {},
-                onAccountClick = {},
-                onPurchaseBillingProductClick = { _ -> }
-            )
-        }
+    fun testShowBillingErrorPaymentButton() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                OutOfTimeScreen(
+                    uiState =
+                        OutOfTimeUiState(
+                            showSitePayment = true,
+                            billingPaymentState = PaymentState.Error.Billing
+                        ),
+                    onSitePaymentClick = {},
+                    onRedeemVoucherClick = {},
+                    onSettingsClick = {},
+                    onAccountClick = {},
+                    onPurchaseBillingProductClick = { _ -> }
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText("Add 30 days time").assertExists()
-    }
+            // Assert
+            onNodeWithText("Add 30 days time").assertExists()
+        }
 
     @Test
-    fun testShowBillingPaymentAvailable() {
-        // Arrange
-        val mockPaymentProduct: PaymentProduct = mockk()
-        every { mockPaymentProduct.price } returns ProductPrice("$10")
-        every { mockPaymentProduct.status } returns null
-        composeTestRule.setContentWithTheme {
-            OutOfTimeScreen(
-                uiState =
-                    OutOfTimeUiState(
-                        showSitePayment = true,
-                        billingPaymentState =
-                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                    ),
-                onSitePaymentClick = {},
-                onRedeemVoucherClick = {},
-                onSettingsClick = {},
-                onAccountClick = {},
-                onPurchaseBillingProductClick = { _ -> }
-            )
-        }
+    fun testShowBillingPaymentAvailable() =
+        composeExtension.use {
+            // Arrange
+            val mockPaymentProduct: PaymentProduct = mockk()
+            every { mockPaymentProduct.price } returns ProductPrice("$10")
+            every { mockPaymentProduct.status } returns null
+            setContentWithTheme {
+                OutOfTimeScreen(
+                    uiState =
+                        OutOfTimeUiState(
+                            showSitePayment = true,
+                            billingPaymentState =
+                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+                        ),
+                    onSitePaymentClick = {},
+                    onRedeemVoucherClick = {},
+                    onSettingsClick = {},
+                    onAccountClick = {},
+                    onPurchaseBillingProductClick = { _ -> }
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText("Add 30 days time ($10)").assertExists()
-    }
+            // Assert
+            onNodeWithText("Add 30 days time ($10)").assertExists()
+        }
 
     @Test
-    fun testShowPendingPayment() {
-        // Arrange
-        val mockPaymentProduct: PaymentProduct = mockk()
-        every { mockPaymentProduct.price } returns ProductPrice("$10")
-        every { mockPaymentProduct.status } returns PaymentStatus.PENDING
-        composeTestRule.setContentWithTheme {
-            OutOfTimeScreen(
-                uiState =
-                    OutOfTimeUiState(
-                        showSitePayment = true,
-                        billingPaymentState =
-                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                    ),
-            )
-        }
+    fun testShowPendingPayment() =
+        composeExtension.use {
+            // Arrange
+            val mockPaymentProduct: PaymentProduct = mockk()
+            every { mockPaymentProduct.price } returns ProductPrice("$10")
+            every { mockPaymentProduct.status } returns PaymentStatus.PENDING
+            setContentWithTheme {
+                OutOfTimeScreen(
+                    uiState =
+                        OutOfTimeUiState(
+                            showSitePayment = true,
+                            billingPaymentState =
+                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+                        ),
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText("Google Play payment pending").assertExists()
-    }
+            // Assert
+            onNodeWithText("Google Play payment pending").assertExists()
+        }
 
     @Test
-    fun testShowPendingPaymentInfoDialog() {
-        // Arrange
-        val mockPaymentProduct: PaymentProduct = mockk()
-        every { mockPaymentProduct.price } returns ProductPrice("$10")
-        every { mockPaymentProduct.status } returns PaymentStatus.PENDING
-        val mockNavigateToVerificationPending: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            OutOfTimeScreen(
-                uiState =
-                    OutOfTimeUiState(
-                        showSitePayment = true,
-                        billingPaymentState =
-                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                    ),
-                navigateToVerificationPendingDialog = mockNavigateToVerificationPending
-            )
-        }
+    fun testShowPendingPaymentInfoDialog() =
+        composeExtension.use {
+            // Arrange
+            val mockPaymentProduct: PaymentProduct = mockk()
+            every { mockPaymentProduct.price } returns ProductPrice("$10")
+            every { mockPaymentProduct.status } returns PaymentStatus.PENDING
+            val mockNavigateToVerificationPending: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                OutOfTimeScreen(
+                    uiState =
+                        OutOfTimeUiState(
+                            showSitePayment = true,
+                            billingPaymentState =
+                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+                        ),
+                    navigateToVerificationPendingDialog = mockNavigateToVerificationPending
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithTag(PLAY_PAYMENT_INFO_ICON_TEST_TAG).performClick()
-        composeTestRule.onNodeWithTag(PLAY_PAYMENT_INFO_ICON_TEST_TAG).assertExists()
+            // Act
+            onNodeWithTag(PLAY_PAYMENT_INFO_ICON_TEST_TAG).performClick()
+            onNodeWithTag(PLAY_PAYMENT_INFO_ICON_TEST_TAG).assertExists()
 
-        verify(exactly = 1) { mockNavigateToVerificationPending.invoke() }
-    }
+            // Assert
+            verify(exactly = 1) { mockNavigateToVerificationPending.invoke() }
+        }
 
     @Test
-    fun testShowVerificationInProgress() {
-        // Arrange
-        val mockPaymentProduct: PaymentProduct = mockk()
-        every { mockPaymentProduct.price } returns ProductPrice("$10")
-        every { mockPaymentProduct.status } returns PaymentStatus.VERIFICATION_IN_PROGRESS
-        composeTestRule.setContentWithTheme {
-            OutOfTimeScreen(
-                uiState =
-                    OutOfTimeUiState(
-                        billingPaymentState =
-                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
-                        showSitePayment = true,
-                    )
-            )
-        }
+    fun testShowVerificationInProgress() =
+        composeExtension.use {
+            // Arrange
+            val mockPaymentProduct: PaymentProduct = mockk()
+            every { mockPaymentProduct.price } returns ProductPrice("$10")
+            every { mockPaymentProduct.status } returns PaymentStatus.VERIFICATION_IN_PROGRESS
+            setContentWithTheme {
+                OutOfTimeScreen(
+                    uiState =
+                        OutOfTimeUiState(
+                            billingPaymentState =
+                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
+                            showSitePayment = true,
+                        )
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText("Verifying purchase").assertExists()
-    }
+            // Assert
+            onNodeWithText("Verifying purchase").assertExists()
+        }
 
     @Test
-    fun testOnPurchaseBillingProductClick() {
-        // Arrange
-        val clickHandler: (ProductId) -> Unit = mockk(relaxed = true)
-        val mockPaymentProduct: PaymentProduct = mockk()
-        every { mockPaymentProduct.price } returns ProductPrice("$10")
-        every { mockPaymentProduct.productId } returns ProductId("PRODUCT_ID")
-        every { mockPaymentProduct.status } returns null
-        composeTestRule.setContentWithTheme {
-            OutOfTimeScreen(
-                uiState =
-                    OutOfTimeUiState(
-                        billingPaymentState =
-                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
-                        showSitePayment = true,
-                    ),
-                onSitePaymentClick = {},
-                onRedeemVoucherClick = {},
-                onSettingsClick = {},
-                onAccountClick = {},
-                onPurchaseBillingProductClick = clickHandler
-            )
-        }
+    fun testOnPurchaseBillingProductClick() =
+        composeExtension.use {
+            // Arrange
+            val clickHandler: (ProductId) -> Unit = mockk(relaxed = true)
+            val mockPaymentProduct: PaymentProduct = mockk()
+            every { mockPaymentProduct.price } returns ProductPrice("$10")
+            every { mockPaymentProduct.productId } returns ProductId("PRODUCT_ID")
+            every { mockPaymentProduct.status } returns null
+            setContentWithTheme {
+                OutOfTimeScreen(
+                    uiState =
+                        OutOfTimeUiState(
+                            billingPaymentState =
+                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
+                            showSitePayment = true,
+                        ),
+                    onSitePaymentClick = {},
+                    onRedeemVoucherClick = {},
+                    onSettingsClick = {},
+                    onAccountClick = {},
+                    onPurchaseBillingProductClick = clickHandler
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithText("Add 30 days time ($10)").performClick()
+            // Act
+            onNodeWithText("Add 30 days time ($10)").performClick()
 
-        // Assert
-        verify { clickHandler(ProductId("PRODUCT_ID")) }
-    }
+            // Assert
+            verify { clickHandler(ProductId("PRODUCT_ID")) }
+        }
 }
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/RedeemVoucherDialogTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/RedeemVoucherDialogTest.kt
index 5a51a8f885..9a2a9ce86d 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/RedeemVoucherDialogTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/RedeemVoucherDialogTest.kt
@@ -1,10 +1,11 @@
 package net.mullvad.mullvadvpn.compose.screen
 
-import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.onNodeWithTag
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
 import androidx.compose.ui.test.performTextInput
+import de.mannodermaus.junit5.compose.createComposeExtension
 import io.mockk.mockk
 import io.mockk.mockkObject
 import io.mockk.verify
@@ -14,131 +15,138 @@ import net.mullvad.mullvadvpn.compose.state.VoucherDialogState
 import net.mullvad.mullvadvpn.compose.state.VoucherDialogUiState
 import net.mullvad.mullvadvpn.compose.test.VOUCHER_INPUT_TEST_TAG
 import net.mullvad.mullvadvpn.util.VoucherRegexHelper
-import org.junit.Before
-import org.junit.Rule
-import org.junit.Test
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.RegisterExtension
 
+@OptIn(ExperimentalTestApi::class)
 class RedeemVoucherDialogTest {
-    @get:Rule val composeTestRule = createComposeRule()
+    @JvmField @RegisterExtension val composeExtension = createComposeExtension()
 
-    @Before
+    @BeforeEach
     fun setup() {
         mockkObject(VoucherRegexHelper)
     }
 
     @Test
-    fun testDismissDialog() {
-        // Arrange
-        val mockedClickHandler: (Boolean) -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            RedeemVoucherDialog(
-                uiState = VoucherDialogUiState.INITIAL,
-                onVoucherInputChange = {},
-                onRedeem = {},
-                onDismiss = mockedClickHandler
-            )
-        }
+    fun testDismissDialog() =
+        composeExtension.use {
+            // Arrange
+            val mockedClickHandler: (Boolean) -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                RedeemVoucherDialog(
+                    uiState = VoucherDialogUiState.INITIAL,
+                    onVoucherInputChange = {},
+                    onRedeem = {},
+                    onDismiss = mockedClickHandler
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithText(CANCEL_BUTTON_TEXT).performClick()
+            // Act
+            onNodeWithText(CANCEL_BUTTON_TEXT).performClick()
 
-        // Assert
-        verify { mockedClickHandler.invoke(false) }
-    }
+            // Assert
+            verify { mockedClickHandler.invoke(false) }
+        }
 
     @Test
-    fun testDismissDialogAfterSuccessfulRedeem() {
-        // Arrange
-        val mockedClickHandler: (Boolean) -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            RedeemVoucherDialog(
-                uiState =
-                    VoucherDialogUiState(voucherViewModelState = VoucherDialogState.Success(0)),
-                onVoucherInputChange = {},
-                onRedeem = {},
-                onDismiss = mockedClickHandler
-            )
-        }
+    fun testDismissDialogAfterSuccessfulRedeem() =
+        composeExtension.use {
+            // Arrange
+            val mockedClickHandler: (Boolean) -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                RedeemVoucherDialog(
+                    uiState =
+                        VoucherDialogUiState(voucherViewModelState = VoucherDialogState.Success(0)),
+                    onVoucherInputChange = {},
+                    onRedeem = {},
+                    onDismiss = mockedClickHandler
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithText(GOT_IT_BUTTON_TEXT).performClick()
+            // Act
+            onNodeWithText(GOT_IT_BUTTON_TEXT).performClick()
 
-        // Assert
-        verify { mockedClickHandler.invoke(true) }
-    }
+            // Assert
+            verify { mockedClickHandler.invoke(true) }
+        }
 
     @Test
-    fun testInsertVoucher() {
-        // Arrange
-        val mockedClickHandler: (String) -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            RedeemVoucherDialog(
-                uiState = VoucherDialogUiState(),
-                onVoucherInputChange = mockedClickHandler,
-                onRedeem = {},
-                onDismiss = {}
-            )
-        }
+    fun testInsertVoucher() =
+        composeExtension.use {
+            // Arrange
+            val mockedClickHandler: (String) -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                RedeemVoucherDialog(
+                    uiState = VoucherDialogUiState(),
+                    onVoucherInputChange = mockedClickHandler,
+                    onRedeem = {},
+                    onDismiss = {}
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithTag(VOUCHER_INPUT_TEST_TAG).performTextInput(DUMMY_VOUCHER)
+            // Act
+            onNodeWithTag(VOUCHER_INPUT_TEST_TAG).performTextInput(DUMMY_VOUCHER)
 
-        // Assert
-        verify { mockedClickHandler.invoke(DUMMY_VOUCHER) }
-    }
+            // Assert
+            verify { mockedClickHandler.invoke(DUMMY_VOUCHER) }
+        }
 
     @Test
-    fun testVerifyingState() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            RedeemVoucherDialog(
-                uiState =
-                    VoucherDialogUiState(voucherViewModelState = VoucherDialogState.Verifying),
-                onVoucherInputChange = {},
-                onRedeem = {},
-                onDismiss = {}
-            )
-        }
+    fun testVerifyingState() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                RedeemVoucherDialog(
+                    uiState =
+                        VoucherDialogUiState(voucherViewModelState = VoucherDialogState.Verifying),
+                    onVoucherInputChange = {},
+                    onRedeem = {},
+                    onDismiss = {}
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText("Verifying voucher…").assertExists()
-    }
+            // Assert
+            onNodeWithText("Verifying voucher…").assertExists()
+        }
 
     @Test
-    fun testSuccessState() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            RedeemVoucherDialog(
-                uiState =
-                    VoucherDialogUiState(voucherViewModelState = VoucherDialogState.Success(0)),
-                onVoucherInputChange = {},
-                onRedeem = {},
-                onDismiss = {}
-            )
-        }
+    fun testSuccessState() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                RedeemVoucherDialog(
+                    uiState =
+                        VoucherDialogUiState(voucherViewModelState = VoucherDialogState.Success(0)),
+                    onVoucherInputChange = {},
+                    onRedeem = {},
+                    onDismiss = {}
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText("Voucher was successfully redeemed.").assertExists()
-    }
+            // Assert
+            onNodeWithText("Voucher was successfully redeemed.").assertExists()
+        }
 
     @Test
-    fun testErrorState() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            RedeemVoucherDialog(
-                uiState =
-                    VoucherDialogUiState(
-                        voucherViewModelState = VoucherDialogState.Error(ERROR_MESSAGE)
-                    ),
-                onVoucherInputChange = {},
-                onRedeem = {},
-                onDismiss = {}
-            )
-        }
+    fun testErrorState() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                RedeemVoucherDialog(
+                    uiState =
+                        VoucherDialogUiState(
+                            voucherViewModelState = VoucherDialogState.Error(ERROR_MESSAGE)
+                        ),
+                    onVoucherInputChange = {},
+                    onRedeem = {},
+                    onDismiss = {}
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText(ERROR_MESSAGE).assertExists()
-    }
+            // Assert
+            onNodeWithText(ERROR_MESSAGE).assertExists()
+        }
 
     companion object {
         private const val REDEEM_BUTTON_TEXT = "Redeem"
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/SelectLocationScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/SelectLocationScreenTest.kt
index ea1d261689..fbc8b046fd 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/SelectLocationScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/SelectLocationScreenTest.kt
@@ -1,9 +1,10 @@
 package net.mullvad.mullvadvpn.compose.screen
 
-import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.onNodeWithTag
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performTextInput
+import de.mannodermaus.junit5.compose.createComposeExtension
 import io.mockk.MockKAnnotations
 import io.mockk.mockk
 import io.mockk.verify
@@ -19,49 +20,51 @@ import net.mullvad.mullvadvpn.model.RelayListCountry
 import net.mullvad.mullvadvpn.model.WireguardEndpointData
 import net.mullvad.mullvadvpn.model.WireguardRelayEndpointData
 import net.mullvad.mullvadvpn.relaylist.toRelayCountries
-import org.junit.Before
-import org.junit.Rule
-import org.junit.Test
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.RegisterExtension
 
+@OptIn(ExperimentalTestApi::class)
 class SelectLocationScreenTest {
-    @get:Rule val composeTestRule = createComposeRule()
+    @JvmField @RegisterExtension val composeExtension = createComposeExtension()
 
-    @Before
+    @BeforeEach
     fun setup() {
         MockKAnnotations.init(this)
     }
 
     @Test
-    fun testDefaultState() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            SelectLocationScreen(
-                uiState = SelectLocationUiState.Loading,
-            )
-        }
+    fun testDefaultState() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                SelectLocationScreen(
+                    uiState = SelectLocationUiState.Loading,
+                )
+            }
 
-        // Assert
-        composeTestRule.apply { onNodeWithTag(CIRCULAR_PROGRESS_INDICATOR).assertExists() }
-    }
+            // Assert
+            onNodeWithTag(CIRCULAR_PROGRESS_INDICATOR).assertExists()
+        }
 
     @Test
-    fun testShowRelayListState() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            SelectLocationScreen(
-                uiState =
-                    SelectLocationUiState.ShowData(
-                        countries = DUMMY_RELAY_COUNTRIES,
-                        selectedRelay = null,
-                        selectedOwnership = null,
-                        selectedProvidersCount = 0,
-                        searchTerm = ""
-                    ),
-            )
-        }
+    fun testShowRelayListState() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                SelectLocationScreen(
+                    uiState =
+                        SelectLocationUiState.ShowData(
+                            countries = DUMMY_RELAY_COUNTRIES,
+                            selectedRelay = null,
+                            selectedOwnership = null,
+                            selectedProvidersCount = 0,
+                            searchTerm = ""
+                        ),
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText("Relay Country 1").assertExists()
             onNodeWithText("Relay City 1").assertDoesNotExist()
             onNodeWithText("Relay host 1").assertDoesNotExist()
@@ -69,37 +72,36 @@ class SelectLocationScreenTest {
             onNodeWithText("Relay City 2").assertDoesNotExist()
             onNodeWithText("Relay host 2").assertDoesNotExist()
         }
-    }
 
     @Test
-    fun testShowRelayListStateSelected() {
-        val updatedDummyList =
-            DUMMY_RELAY_COUNTRIES.let {
-                val cities = it[0].cities.toMutableList()
-                val city = cities.removeAt(0)
-                cities.add(0, city.copy(expanded = true))
+    fun testShowRelayListStateSelected() =
+        composeExtension.use {
+            val updatedDummyList =
+                DUMMY_RELAY_COUNTRIES.let {
+                    val cities = it[0].cities.toMutableList()
+                    val city = cities.removeAt(0)
+                    cities.add(0, city.copy(expanded = true))
 
-                val mutableRelayList = it.toMutableList()
-                mutableRelayList[0] = it[0].copy(expanded = true, cities = cities.toList())
-                mutableRelayList
-            }
+                    val mutableRelayList = it.toMutableList()
+                    mutableRelayList[0] = it[0].copy(expanded = true, cities = cities.toList())
+                    mutableRelayList
+                }
 
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            SelectLocationScreen(
-                uiState =
-                    SelectLocationUiState.ShowData(
-                        countries = updatedDummyList,
-                        selectedRelay = updatedDummyList[0].cities[0].relays[0],
-                        selectedOwnership = null,
-                        selectedProvidersCount = 0,
-                        searchTerm = ""
-                    ),
-            )
-        }
+            // Arrange
+            setContentWithTheme {
+                SelectLocationScreen(
+                    uiState =
+                        SelectLocationUiState.ShowData(
+                            countries = updatedDummyList,
+                            selectedRelay = updatedDummyList[0].cities[0].relays[0],
+                            selectedOwnership = null,
+                            selectedProvidersCount = 0,
+                            searchTerm = ""
+                        ),
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText("Relay Country 1").assertExists()
             onNodeWithText("Relay City 1").assertExists()
             onNodeWithText("Relay host 1").assertExists()
@@ -107,59 +109,58 @@ class SelectLocationScreenTest {
             onNodeWithText("Relay City 2").assertDoesNotExist()
             onNodeWithText("Relay host 2").assertDoesNotExist()
         }
-    }
 
     @Test
-    fun testSearchInput() {
-        // Arrange
-        val mockedSearchTermInput: (String) -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            SelectLocationScreen(
-                uiState =
-                    SelectLocationUiState.ShowData(
-                        countries = emptyList(),
-                        selectedRelay = null,
-                        selectedOwnership = null,
-                        selectedProvidersCount = 0,
-                        searchTerm = ""
-                    ),
-                onSearchTermInput = mockedSearchTermInput
-            )
-        }
-        val mockSearchString = "SEARCH"
+    fun testSearchInput() =
+        composeExtension.use {
+            // Arrange
+            val mockedSearchTermInput: (String) -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                SelectLocationScreen(
+                    uiState =
+                        SelectLocationUiState.ShowData(
+                            countries = emptyList(),
+                            selectedRelay = null,
+                            selectedOwnership = null,
+                            selectedProvidersCount = 0,
+                            searchTerm = ""
+                        ),
+                    onSearchTermInput = mockedSearchTermInput
+                )
+            }
+            val mockSearchString = "SEARCH"
 
-        // Act
-        composeTestRule.apply { onNodeWithText("Search for...").performTextInput(mockSearchString) }
+            // Act
+            onNodeWithText("Search for...").performTextInput(mockSearchString)
 
-        // Assert
-        verify { mockedSearchTermInput.invoke(mockSearchString) }
-    }
+            // Assert
+            verify { mockedSearchTermInput.invoke(mockSearchString) }
+        }
 
     @Test
-    fun testSearchTermNotFound() {
-        // Arrange
-        val mockedSearchTermInput: (String) -> Unit = mockk(relaxed = true)
-        val mockSearchString = "SEARCH"
-        composeTestRule.setContentWithTheme {
-            SelectLocationScreen(
-                uiState =
-                    SelectLocationUiState.ShowData(
-                        countries = emptyList(),
-                        selectedRelay = null,
-                        selectedOwnership = null,
-                        selectedProvidersCount = 0,
-                        searchTerm = mockSearchString
-                    ),
-                onSearchTermInput = mockedSearchTermInput
-            )
-        }
+    fun testSearchTermNotFound() =
+        composeExtension.use {
+            // Arrange
+            val mockedSearchTermInput: (String) -> Unit = mockk(relaxed = true)
+            val mockSearchString = "SEARCH"
+            setContentWithTheme {
+                SelectLocationScreen(
+                    uiState =
+                        SelectLocationUiState.ShowData(
+                            countries = emptyList(),
+                            selectedRelay = null,
+                            selectedOwnership = null,
+                            selectedProvidersCount = 0,
+                            searchTerm = mockSearchString
+                        ),
+                    onSearchTermInput = mockedSearchTermInput
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText("No result for $mockSearchString.", substring = true).assertExists()
             onNodeWithText("Try a different search", substring = true).assertExists()
         }
-    }
 
     companion object {
         private val DUMMY_RELAY_1 =
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/SettingsScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/SettingsScreenTest.kt
index f5daf29bb2..1c47082b14 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/SettingsScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/SettingsScreenTest.kt
@@ -1,66 +1,66 @@
 package net.mullvad.mullvadvpn.compose.screen
 
 import androidx.compose.material3.ExperimentalMaterial3Api
-import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.onNodeWithText
+import de.mannodermaus.junit5.compose.createComposeExtension
 import io.mockk.MockKAnnotations
 import net.mullvad.mullvadvpn.compose.setContentWithTheme
 import net.mullvad.mullvadvpn.compose.state.SettingsUiState
-import org.junit.Before
-import org.junit.Rule
-import org.junit.Test
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.RegisterExtension
 
+@OptIn(ExperimentalTestApi::class)
 class SettingsScreenTest {
-    @get:Rule val composeTestRule = createComposeRule()
+    @JvmField @RegisterExtension val composeExtension = createComposeExtension()
 
-    @Before
+    @BeforeEach
     fun setup() {
         MockKAnnotations.init(this)
     }
 
     @Test
     @OptIn(ExperimentalMaterial3Api::class)
-    fun testLoggedInState() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            SettingsScreen(
-                uiState =
-                    SettingsUiState(
-                        appVersion = "",
-                        isLoggedIn = true,
-                        isUpdateAvailable = true,
-                        isPlayBuild = false
-                    ),
-            )
-        }
-        // Assert
-        composeTestRule.apply {
+    fun testLoggedInState() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                SettingsScreen(
+                    uiState =
+                        SettingsUiState(
+                            appVersion = "",
+                            isLoggedIn = true,
+                            isUpdateAvailable = true,
+                            isPlayBuild = false
+                        ),
+                )
+            }
+            // Assert
             onNodeWithText("VPN settings").assertExists()
             onNodeWithText("Split tunneling").assertExists()
             onNodeWithText("App version").assertExists()
         }
-    }
 
     @Test
     @OptIn(ExperimentalMaterial3Api::class)
-    fun testLoggedOutState() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            SettingsScreen(
-                uiState =
-                    SettingsUiState(
-                        appVersion = "",
-                        isLoggedIn = false,
-                        isUpdateAvailable = true,
-                        isPlayBuild = false
-                    ),
-            )
-        }
-        // Assert
-        composeTestRule.apply {
+    fun testLoggedOutState() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                SettingsScreen(
+                    uiState =
+                        SettingsUiState(
+                            appVersion = "",
+                            isLoggedIn = false,
+                            isUpdateAvailable = true,
+                            isPlayBuild = false
+                        ),
+                )
+            }
+            // Assert
             onNodeWithText("VPN settings").assertDoesNotExist()
             onNodeWithText("Split tunneling").assertDoesNotExist()
             onNodeWithText("App version").assertExists()
         }
-    }
 }
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/SplitTunnelingScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/SplitTunnelingScreenTest.kt
index 7f4ee36d47..dc1362ad37 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/SplitTunnelingScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/SplitTunnelingScreenTest.kt
@@ -1,9 +1,10 @@
 package net.mullvad.mullvadvpn.compose.screen
 
 import android.graphics.Bitmap
-import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
+import de.mannodermaus.junit5.compose.createComposeExtension
 import io.mockk.MockKAnnotations
 import io.mockk.every
 import io.mockk.mockk
@@ -14,15 +15,16 @@ import net.mullvad.mullvadvpn.applist.ApplicationsIconManager
 import net.mullvad.mullvadvpn.compose.setContentWithTheme
 import net.mullvad.mullvadvpn.compose.state.SplitTunnelingUiState
 import org.junit.After
-import org.junit.Before
-import org.junit.Rule
-import org.junit.Test
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.RegisterExtension
 import org.koin.core.context.loadKoinModules
 import org.koin.core.context.unloadKoinModules
 import org.koin.dsl.module
 
+@OptIn(ExperimentalTestApi::class)
 class SplitTunnelingScreenTest {
-    @get:Rule val composeTestRule = createComposeRule()
+    @JvmField @RegisterExtension val composeExtension = createComposeExtension()
 
     private val mockBitmap: Bitmap = Bitmap.createBitmap(10, 10, Bitmap.Config.ARGB_8888)
     private val testModule = module {
@@ -33,7 +35,7 @@ class SplitTunnelingScreenTest {
         }
     }
 
-    @Before
+    @BeforeEach
     fun setup() {
         MockKAnnotations.init(this)
         loadKoinModules(testModule)
@@ -46,42 +48,47 @@ class SplitTunnelingScreenTest {
     }
 
     @Test
-    fun testLoadingState() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            SplitTunnelingScreen(uiState = SplitTunnelingUiState.Loading)
-        }
+    fun testLoadingState() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme { SplitTunnelingScreen(uiState = SplitTunnelingUiState.Loading) }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText(TITLE).assertExists()
             onNodeWithText(DESCRIPTION).assertExists()
             onNodeWithText(EXCLUDED_APPLICATIONS).assertDoesNotExist()
             onNodeWithText(SHOW_SYSTEM_APPS).assertDoesNotExist()
             onNodeWithText(ALL_APPLICATIONS).assertDoesNotExist()
         }
-    }
 
     @Test
-    fun testListDisplayed() {
-        // Arrange
-        val excludedApp =
-            AppData(packageName = EXCLUDED_APP_PACKAGE_NAME, iconRes = 0, name = EXCLUDED_APP_NAME)
-        val includedApp =
-            AppData(packageName = INCLUDED_APP_PACKAGE_NAME, iconRes = 0, name = INCLUDED_APP_NAME)
-        composeTestRule.setContentWithTheme {
-            SplitTunnelingScreen(
-                uiState =
-                    SplitTunnelingUiState.ShowAppList(
-                        excludedApps = listOf(excludedApp),
-                        includedApps = listOf(includedApp),
-                        showSystemApps = false
-                    )
-            )
-        }
+    fun testListDisplayed() =
+        composeExtension.use {
+            // Arrange
+            val excludedApp =
+                AppData(
+                    packageName = EXCLUDED_APP_PACKAGE_NAME,
+                    iconRes = 0,
+                    name = EXCLUDED_APP_NAME
+                )
+            val includedApp =
+                AppData(
+                    packageName = INCLUDED_APP_PACKAGE_NAME,
+                    iconRes = 0,
+                    name = INCLUDED_APP_NAME
+                )
+            setContentWithTheme {
+                SplitTunnelingScreen(
+                    uiState =
+                        SplitTunnelingUiState.ShowAppList(
+                            excludedApps = listOf(excludedApp),
+                            includedApps = listOf(includedApp),
+                            showSystemApps = false
+                        )
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText(TITLE).assertExists()
             onNodeWithText(DESCRIPTION).assertExists()
             onNodeWithText(EXCLUDED_APPLICATIONS).assertExists()
@@ -90,26 +97,29 @@ class SplitTunnelingScreenTest {
             onNodeWithText(ALL_APPLICATIONS).assertExists()
             onNodeWithText(INCLUDED_APP_NAME).assertExists()
         }
-    }
 
     @Test
-    fun testNoExcludedApps() {
-        // Arrange
-        val includedApp =
-            AppData(packageName = INCLUDED_APP_PACKAGE_NAME, iconRes = 0, name = INCLUDED_APP_NAME)
-        composeTestRule.setContentWithTheme {
-            SplitTunnelingScreen(
-                uiState =
-                    SplitTunnelingUiState.ShowAppList(
-                        excludedApps = emptyList(),
-                        includedApps = listOf(includedApp),
-                        showSystemApps = false
-                    )
-            )
-        }
+    fun testNoExcludedApps() =
+        composeExtension.use {
+            // Arrange
+            val includedApp =
+                AppData(
+                    packageName = INCLUDED_APP_PACKAGE_NAME,
+                    iconRes = 0,
+                    name = INCLUDED_APP_NAME
+                )
+            setContentWithTheme {
+                SplitTunnelingScreen(
+                    uiState =
+                        SplitTunnelingUiState.ShowAppList(
+                            excludedApps = emptyList(),
+                            includedApps = listOf(includedApp),
+                            showSystemApps = false
+                        )
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText(TITLE).assertExists()
             onNodeWithText(DESCRIPTION).assertExists()
             onNodeWithText(EXCLUDED_APPLICATIONS).assertDoesNotExist()
@@ -118,88 +128,114 @@ class SplitTunnelingScreenTest {
             onNodeWithText(ALL_APPLICATIONS).assertExists()
             onNodeWithText(INCLUDED_APP_NAME).assertExists()
         }
-    }
 
     @Test
-    fun testClickIncludedItem() {
-        // Arrange
-        val excludedApp =
-            AppData(packageName = EXCLUDED_APP_PACKAGE_NAME, iconRes = 0, name = EXCLUDED_APP_NAME)
-        val includedApp =
-            AppData(packageName = INCLUDED_APP_PACKAGE_NAME, iconRes = 0, name = INCLUDED_APP_NAME)
-        val mockedClickHandler: (String) -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            SplitTunnelingScreen(
-                uiState =
-                    SplitTunnelingUiState.ShowAppList(
-                        excludedApps = listOf(excludedApp),
-                        includedApps = listOf(includedApp),
-                        showSystemApps = false
-                    ),
-                onExcludeAppClick = mockedClickHandler
-            )
-        }
+    fun testClickIncludedItem() =
+        composeExtension.use {
+            // Arrange
+            val excludedApp =
+                AppData(
+                    packageName = EXCLUDED_APP_PACKAGE_NAME,
+                    iconRes = 0,
+                    name = EXCLUDED_APP_NAME
+                )
+            val includedApp =
+                AppData(
+                    packageName = INCLUDED_APP_PACKAGE_NAME,
+                    iconRes = 0,
+                    name = INCLUDED_APP_NAME
+                )
+            val mockedClickHandler: (String) -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                SplitTunnelingScreen(
+                    uiState =
+                        SplitTunnelingUiState.ShowAppList(
+                            excludedApps = listOf(excludedApp),
+                            includedApps = listOf(includedApp),
+                            showSystemApps = false
+                        ),
+                    onExcludeAppClick = mockedClickHandler
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithText(INCLUDED_APP_NAME).performClick()
+            // Act
+            onNodeWithText(INCLUDED_APP_NAME).performClick()
 
-        // Assert
-        verify { mockedClickHandler.invoke(INCLUDED_APP_PACKAGE_NAME) }
-    }
+            // Assert
+            verify { mockedClickHandler.invoke(INCLUDED_APP_PACKAGE_NAME) }
+        }
 
     @Test
-    fun testClickExcludedItem() {
-        // Arrange
-        val excludedApp =
-            AppData(packageName = EXCLUDED_APP_PACKAGE_NAME, iconRes = 0, name = EXCLUDED_APP_NAME)
-        val includedApp =
-            AppData(packageName = INCLUDED_APP_PACKAGE_NAME, iconRes = 0, name = INCLUDED_APP_NAME)
-        val mockedClickHandler: (String) -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            SplitTunnelingScreen(
-                uiState =
-                    SplitTunnelingUiState.ShowAppList(
-                        excludedApps = listOf(excludedApp),
-                        includedApps = listOf(includedApp),
-                        showSystemApps = false
-                    ),
-                onIncludeAppClick = mockedClickHandler
-            )
-        }
+    fun testClickExcludedItem() =
+        composeExtension.use {
+            // Arrange
+            val excludedApp =
+                AppData(
+                    packageName = EXCLUDED_APP_PACKAGE_NAME,
+                    iconRes = 0,
+                    name = EXCLUDED_APP_NAME
+                )
+            val includedApp =
+                AppData(
+                    packageName = INCLUDED_APP_PACKAGE_NAME,
+                    iconRes = 0,
+                    name = INCLUDED_APP_NAME
+                )
+            val mockedClickHandler: (String) -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                SplitTunnelingScreen(
+                    uiState =
+                        SplitTunnelingUiState.ShowAppList(
+                            excludedApps = listOf(excludedApp),
+                            includedApps = listOf(includedApp),
+                            showSystemApps = false
+                        ),
+                    onIncludeAppClick = mockedClickHandler
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithText(EXCLUDED_APP_NAME).performClick()
+            // Act
+            onNodeWithText(EXCLUDED_APP_NAME).performClick()
 
-        // Assert
-        verify { mockedClickHandler.invoke(EXCLUDED_APP_PACKAGE_NAME) }
-    }
+            // Assert
+            verify { mockedClickHandler.invoke(EXCLUDED_APP_PACKAGE_NAME) }
+        }
 
     @Test
-    fun testClickShowSystemApps() {
-        // Arrange
-        val excludedApp =
-            AppData(packageName = EXCLUDED_APP_PACKAGE_NAME, iconRes = 0, name = EXCLUDED_APP_NAME)
-        val includedApp =
-            AppData(packageName = INCLUDED_APP_PACKAGE_NAME, iconRes = 0, name = INCLUDED_APP_NAME)
-        val mockedClickHandler: (Boolean) -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            SplitTunnelingScreen(
-                uiState =
-                    SplitTunnelingUiState.ShowAppList(
-                        excludedApps = listOf(excludedApp),
-                        includedApps = listOf(includedApp),
-                        showSystemApps = false
-                    ),
-                onShowSystemAppsClick = mockedClickHandler
-            )
-        }
+    fun testClickShowSystemApps() =
+        composeExtension.use {
+            // Arrange
+            val excludedApp =
+                AppData(
+                    packageName = EXCLUDED_APP_PACKAGE_NAME,
+                    iconRes = 0,
+                    name = EXCLUDED_APP_NAME
+                )
+            val includedApp =
+                AppData(
+                    packageName = INCLUDED_APP_PACKAGE_NAME,
+                    iconRes = 0,
+                    name = INCLUDED_APP_NAME
+                )
+            val mockedClickHandler: (Boolean) -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                SplitTunnelingScreen(
+                    uiState =
+                        SplitTunnelingUiState.ShowAppList(
+                            excludedApps = listOf(excludedApp),
+                            includedApps = listOf(includedApp),
+                            showSystemApps = false
+                        ),
+                    onShowSystemAppsClick = mockedClickHandler
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithText(SHOW_SYSTEM_APPS).performClick()
+            // Act
+            onNodeWithText(SHOW_SYSTEM_APPS).performClick()
 
-        // Assert
-        verify { mockedClickHandler.invoke(true) }
-    }
+            // Assert
+            verify { mockedClickHandler.invoke(true) }
+        }
 
     companion object {
         private const val EXCLUDED_APP_PACKAGE_NAME = "excluded-pkg"
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/VpnSettingsScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/VpnSettingsScreenTest.kt
index 1ca2b3e1f7..74c73553b2 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/VpnSettingsScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/VpnSettingsScreenTest.kt
@@ -1,12 +1,13 @@
 package net.mullvad.mullvadvpn.compose.screen
 
+import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.hasTestTag
-import androidx.compose.ui.test.junit4.createComposeRule
 import androidx.compose.ui.test.onNodeWithContentDescription
 import androidx.compose.ui.test.onNodeWithTag
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
 import androidx.compose.ui.test.performScrollToNode
+import de.mannodermaus.junit5.compose.createComposeExtension
 import io.mockk.MockKAnnotations
 import io.mockk.mockk
 import io.mockk.verify
@@ -25,483 +26,483 @@ import net.mullvad.mullvadvpn.model.PortRange
 import net.mullvad.mullvadvpn.model.QuantumResistantState
 import net.mullvad.mullvadvpn.onNodeWithTagAndText
 import net.mullvad.mullvadvpn.viewmodel.CustomDnsItem
-import org.junit.Before
-import org.junit.Rule
-import org.junit.Test
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.RegisterExtension
 
+@OptIn(ExperimentalTestApi::class)
 class VpnSettingsScreenTest {
-    @get:Rule val composeTestRule = createComposeRule()
+    @JvmField @RegisterExtension val composeExtension = createComposeExtension()
 
-    @Before
+    @BeforeEach
     fun setup() {
         MockKAnnotations.init(this)
     }
 
     @Test
-    fun testDefaultState() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            VpnSettingsScreen(
-                uiState = VpnSettingsUiState.createDefault(),
-            )
-        }
+    fun testDefaultState() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                VpnSettingsScreen(
+                    uiState = VpnSettingsUiState.createDefault(),
+                )
+            }
 
-        composeTestRule.apply { onNodeWithText("Auto-connect").assertExists() }
+            apply { onNodeWithText("Auto-connect").assertExists() }
 
-        composeTestRule
-            .onNodeWithTag(LAZY_LIST_TEST_TAG)
-            .performScrollToNode(hasTestTag(LAZY_LIST_LAST_ITEM_TEST_TAG))
+            onNodeWithTag(LAZY_LIST_TEST_TAG)
+                .performScrollToNode(hasTestTag(LAZY_LIST_LAST_ITEM_TEST_TAG))
 
-        // Assert
-        composeTestRule.apply {
-            onNodeWithText("WireGuard MTU").assertExists()
-            onNodeWithText("Default").assertExists()
+            // Assert
+            apply {
+                onNodeWithText("WireGuard MTU").assertExists()
+                onNodeWithText("Default").assertExists()
+            }
         }
-    }
 
     @Test
-    fun testMtuCustomValue() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            VpnSettingsScreen(
-                uiState = VpnSettingsUiState.createDefault(mtu = VALID_DUMMY_MTU_VALUE),
-            )
-        }
+    fun testMtuCustomValue() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                VpnSettingsScreen(
+                    uiState = VpnSettingsUiState.createDefault(mtu = VALID_DUMMY_MTU_VALUE),
+                )
+            }
 
-        composeTestRule
-            .onNodeWithTag(LAZY_LIST_TEST_TAG)
-            .performScrollToNode(hasTestTag(LAZY_LIST_LAST_ITEM_TEST_TAG))
+            onNodeWithTag(LAZY_LIST_TEST_TAG)
+                .performScrollToNode(hasTestTag(LAZY_LIST_LAST_ITEM_TEST_TAG))
 
-        // Assert
-        composeTestRule.onNodeWithText(VALID_DUMMY_MTU_VALUE).assertExists()
-    }
+            // Assert
+            onNodeWithText(VALID_DUMMY_MTU_VALUE).assertExists()
+        }
 
     @Test
-    fun testCustomDnsAddressesAndAddButtonVisibleWhenCustomDnsEnabled() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            VpnSettingsScreen(
-                uiState =
-                    VpnSettingsUiState.createDefault(
-                        isCustomDnsEnabled = true,
-                        customDnsItems =
-                            listOf(
-                                CustomDnsItem(address = DUMMY_DNS_ADDRESS, false),
-                                CustomDnsItem(address = DUMMY_DNS_ADDRESS_2, false),
-                                CustomDnsItem(address = DUMMY_DNS_ADDRESS_3, false)
-                            )
-                    ),
-            )
-        }
+    fun testCustomDnsAddressesAndAddButtonVisibleWhenCustomDnsEnabled() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                VpnSettingsScreen(
+                    uiState =
+                        VpnSettingsUiState.createDefault(
+                            isCustomDnsEnabled = true,
+                            customDnsItems =
+                                listOf(
+                                    CustomDnsItem(address = DUMMY_DNS_ADDRESS, false),
+                                    CustomDnsItem(address = DUMMY_DNS_ADDRESS_2, false),
+                                    CustomDnsItem(address = DUMMY_DNS_ADDRESS_3, false)
+                                )
+                        ),
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText(DUMMY_DNS_ADDRESS).assertExists()
             onNodeWithText(DUMMY_DNS_ADDRESS_2).assertExists()
             onNodeWithText(DUMMY_DNS_ADDRESS_3).assertExists()
             onNodeWithText("Add a server").assertExists()
         }
-    }
 
     @Test
-    fun testCustomDnsAddressesAndAddButtonNotVisibleWhenCustomDnsDisabled() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            VpnSettingsScreen(
-                uiState =
-                    VpnSettingsUiState.createDefault(
-                        isCustomDnsEnabled = false,
-                        customDnsItems = listOf(CustomDnsItem(address = DUMMY_DNS_ADDRESS, false))
-                    ),
-            )
+    fun testCustomDnsAddressesAndAddButtonNotVisibleWhenCustomDnsDisabled() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                VpnSettingsScreen(
+                    uiState =
+                        VpnSettingsUiState.createDefault(
+                            isCustomDnsEnabled = false,
+                            customDnsItems =
+                                listOf(CustomDnsItem(address = DUMMY_DNS_ADDRESS, false))
+                        ),
+                )
+            }
+            onNodeWithTag(LAZY_LIST_TEST_TAG)
+                .performScrollToNode(hasTestTag(LAZY_LIST_LAST_ITEM_TEST_TAG))
+            // Assert
+            onNodeWithText(DUMMY_DNS_ADDRESS).assertDoesNotExist()
+            onNodeWithText("Add a server").assertDoesNotExist()
         }
-        composeTestRule
-            .onNodeWithTag(LAZY_LIST_TEST_TAG)
-            .performScrollToNode(hasTestTag(LAZY_LIST_LAST_ITEM_TEST_TAG))
-        // Assert
-        composeTestRule.onNodeWithText(DUMMY_DNS_ADDRESS).assertDoesNotExist()
-        composeTestRule.onNodeWithText("Add a server").assertDoesNotExist()
-    }
 
     @Test
-    fun testLanWarningNotShownWhenLanTrafficEnabledAndLocalAddressIsUsed() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            VpnSettingsScreen(
-                uiState =
-                    VpnSettingsUiState.createDefault(
-                        isCustomDnsEnabled = true,
-                        isLocalNetworkSharingEnabled = true,
-                        customDnsItems =
-                            listOf(CustomDnsItem(address = DUMMY_DNS_ADDRESS, isLocal = true))
-                    ),
-            )
-        }
+    fun testLanWarningNotShownWhenLanTrafficEnabledAndLocalAddressIsUsed() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                VpnSettingsScreen(
+                    uiState =
+                        VpnSettingsUiState.createDefault(
+                            isCustomDnsEnabled = true,
+                            isLocalNetworkSharingEnabled = true,
+                            customDnsItems =
+                                listOf(CustomDnsItem(address = DUMMY_DNS_ADDRESS, isLocal = true))
+                        ),
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithContentDescription(LOCAL_DNS_SERVER_WARNING).assertDoesNotExist()
-    }
+            // Assert
+            onNodeWithContentDescription(LOCAL_DNS_SERVER_WARNING).assertDoesNotExist()
+        }
 
     @Test
-    fun testLanWarningNotShowedWhenLanTrafficDisabledAndLocalAddressIsNotUsed() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            VpnSettingsScreen(
-                uiState =
-                    VpnSettingsUiState.createDefault(
-                        isCustomDnsEnabled = true,
-                        customDnsItems =
-                            listOf(CustomDnsItem(address = DUMMY_DNS_ADDRESS, isLocal = false))
-                    ),
-            )
-        }
+    fun testLanWarningNotShowedWhenLanTrafficDisabledAndLocalAddressIsNotUsed() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                VpnSettingsScreen(
+                    uiState =
+                        VpnSettingsUiState.createDefault(
+                            isCustomDnsEnabled = true,
+                            customDnsItems =
+                                listOf(CustomDnsItem(address = DUMMY_DNS_ADDRESS, isLocal = false))
+                        ),
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithContentDescription(LOCAL_DNS_SERVER_WARNING).assertDoesNotExist()
-    }
+            // Assert
+            onNodeWithContentDescription(LOCAL_DNS_SERVER_WARNING).assertDoesNotExist()
+        }
 
     @Test
-    fun testLanWarningNotShowedWhenLanTrafficEnabledAndLocalAddressIsNotUsed() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            VpnSettingsScreen(
-                uiState =
-                    VpnSettingsUiState.createDefault(
-                        isCustomDnsEnabled = true,
-                        customDnsItems =
-                            listOf(CustomDnsItem(address = DUMMY_DNS_ADDRESS, isLocal = false))
-                    ),
-            )
-        }
+    fun testLanWarningNotShowedWhenLanTrafficEnabledAndLocalAddressIsNotUsed() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                VpnSettingsScreen(
+                    uiState =
+                        VpnSettingsUiState.createDefault(
+                            isCustomDnsEnabled = true,
+                            customDnsItems =
+                                listOf(CustomDnsItem(address = DUMMY_DNS_ADDRESS, isLocal = false))
+                        ),
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithContentDescription(LOCAL_DNS_SERVER_WARNING).assertDoesNotExist()
-    }
+            // Assert
+            onNodeWithContentDescription(LOCAL_DNS_SERVER_WARNING).assertDoesNotExist()
+        }
 
     @Test
-    fun testLanWarningShowedWhenAllowLanEnabledAndLocalDnsAddressIsUsed() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            VpnSettingsScreen(
-                uiState =
-                    VpnSettingsUiState.createDefault(
-                        isCustomDnsEnabled = true,
-                        customDnsItems =
-                            listOf(CustomDnsItem(address = DUMMY_DNS_ADDRESS, isLocal = true))
-                    ),
-            )
-        }
+    fun testLanWarningShowedWhenAllowLanEnabledAndLocalDnsAddressIsUsed() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                VpnSettingsScreen(
+                    uiState =
+                        VpnSettingsUiState.createDefault(
+                            isCustomDnsEnabled = true,
+                            customDnsItems =
+                                listOf(CustomDnsItem(address = DUMMY_DNS_ADDRESS, isLocal = true))
+                        ),
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithContentDescription(LOCAL_DNS_SERVER_WARNING).assertExists()
         }
-    }
 
     @Test
-    fun testShowSelectedTunnelQuantumOption() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            VpnSettingsScreen(
-                uiState =
-                    VpnSettingsUiState.createDefault(quantumResistant = QuantumResistantState.On),
-            )
-        }
-        composeTestRule
-            .onNodeWithTag(LAZY_LIST_TEST_TAG)
-            .performScrollToNode(hasTestTag(LAZY_LIST_QUANTUM_ITEM_OFF_TEST_TAG))
+    fun testShowSelectedTunnelQuantumOption() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                VpnSettingsScreen(
+                    uiState =
+                        VpnSettingsUiState.createDefault(
+                            quantumResistant = QuantumResistantState.On
+                        ),
+                )
+            }
+            onNodeWithTag(LAZY_LIST_TEST_TAG)
+                .performScrollToNode(hasTestTag(LAZY_LIST_QUANTUM_ITEM_OFF_TEST_TAG))
 
-        // Assert
-        composeTestRule
-            .onNodeWithTagAndText(testTag = LAZY_LIST_QUANTUM_ITEM_ON_TEST_TAG, text = "On")
-            .assertExists()
-    }
+            // Assert
+            onNodeWithTagAndText(testTag = LAZY_LIST_QUANTUM_ITEM_ON_TEST_TAG, text = "On")
+                .assertExists()
+        }
 
     @Test
-    fun testSelectTunnelQuantumOption() {
-        // Arrange
-        val mockSelectQuantumResistantSettingListener: (QuantumResistantState) -> Unit =
-            mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            VpnSettingsScreen(
-                uiState =
-                    VpnSettingsUiState.createDefault(
-                        quantumResistant = QuantumResistantState.Auto,
-                    ),
-                onSelectQuantumResistanceSetting = mockSelectQuantumResistantSettingListener
-            )
-        }
-        composeTestRule
-            .onNodeWithTag(LAZY_LIST_TEST_TAG)
-            .performScrollToNode(hasTestTag(LAZY_LIST_QUANTUM_ITEM_OFF_TEST_TAG))
+    fun testSelectTunnelQuantumOption() =
+        composeExtension.use {
+            // Arrange
+            val mockSelectQuantumResistantSettingListener: (QuantumResistantState) -> Unit =
+                mockk(relaxed = true)
+            setContentWithTheme {
+                VpnSettingsScreen(
+                    uiState =
+                        VpnSettingsUiState.createDefault(
+                            quantumResistant = QuantumResistantState.Auto,
+                        ),
+                    onSelectQuantumResistanceSetting = mockSelectQuantumResistantSettingListener,
+                )
+            }
+            onNodeWithTag(LAZY_LIST_TEST_TAG)
+                .performScrollToNode(hasTestTag(LAZY_LIST_QUANTUM_ITEM_OFF_TEST_TAG))
 
-        // Assert
-        composeTestRule
-            .onNodeWithTagAndText(testTag = LAZY_LIST_QUANTUM_ITEM_ON_TEST_TAG, text = "On")
-            .performClick()
-        verify(exactly = 1) {
-            mockSelectQuantumResistantSettingListener.invoke(QuantumResistantState.On)
+            // Assert
+            onNodeWithTagAndText(testTag = LAZY_LIST_QUANTUM_ITEM_ON_TEST_TAG, text = "On")
+                .performClick()
+            verify(exactly = 1) {
+                mockSelectQuantumResistantSettingListener.invoke(QuantumResistantState.On)
+            }
         }
-    }
 
     @Test
-    fun testShowWireguardPortOptions() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            VpnSettingsScreen(
-                uiState =
-                    VpnSettingsUiState.createDefault(
-                        selectedWireguardPort = Constraint.Only(Port(53))
-                    ),
-            )
-        }
+    fun testShowWireguardPortOptions() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                VpnSettingsScreen(
+                    uiState =
+                        VpnSettingsUiState.createDefault(
+                            selectedWireguardPort = Constraint.Only(Port(53))
+                        ),
+                )
+            }
 
-        // Act
-        composeTestRule
-            .onNodeWithTag(LAZY_LIST_TEST_TAG)
-            .performScrollToNode(
-                hasTestTag(String.format(LAZY_LIST_WIREGUARD_PORT_ITEM_X_TEST_TAG, 53))
-            )
+            // Act
+            onNodeWithTag(LAZY_LIST_TEST_TAG)
+                .performScrollToNode(
+                    hasTestTag(String.format(LAZY_LIST_WIREGUARD_PORT_ITEM_X_TEST_TAG, 53))
+                )
 
-        // Assert
-        composeTestRule
-            .onNodeWithTagAndText(
-                testTag = String.format(LAZY_LIST_WIREGUARD_PORT_ITEM_X_TEST_TAG, 51820),
-                text = "51820"
-            )
-            .assertExists()
-    }
+            // Assert
+            onNodeWithTagAndText(
+                    testTag = String.format(LAZY_LIST_WIREGUARD_PORT_ITEM_X_TEST_TAG, 51820),
+                    text = "51820"
+                )
+                .assertExists()
+        }
 
     @Test
-    fun testSelectWireguardPortOption() {
-        // Arrange
-        val mockSelectWireguardPortSelectionListener: (Constraint<Port>) -> Unit =
-            mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            VpnSettingsScreen(
-                uiState =
-                    VpnSettingsUiState.createDefault(
-                        selectedWireguardPort = Constraint.Only(Port(53))
-                    ),
-                onWireguardPortSelected = mockSelectWireguardPortSelectionListener
-            )
-        }
+    fun testSelectWireguardPortOption() =
+        composeExtension.use {
+            // Arrange
+            val mockSelectWireguardPortSelectionListener: (Constraint<Port>) -> Unit =
+                mockk(relaxed = true)
+            setContentWithTheme {
+                VpnSettingsScreen(
+                    uiState =
+                        VpnSettingsUiState.createDefault(
+                            selectedWireguardPort = Constraint.Only(Port(53))
+                        ),
+                    onWireguardPortSelected = mockSelectWireguardPortSelectionListener,
+                )
+            }
 
-        // Act
-        composeTestRule
-            .onNodeWithTag(LAZY_LIST_TEST_TAG)
-            .performScrollToNode(
-                hasTestTag(String.format(LAZY_LIST_WIREGUARD_PORT_ITEM_X_TEST_TAG, 53))
-            )
-        composeTestRule
-            .onNodeWithTagAndText(
-                testTag = String.format(LAZY_LIST_WIREGUARD_PORT_ITEM_X_TEST_TAG, 51820),
-                text = "51820"
-            )
-            .performClick()
+            // Act
+            onNodeWithTag(LAZY_LIST_TEST_TAG)
+                .performScrollToNode(
+                    hasTestTag(String.format(LAZY_LIST_WIREGUARD_PORT_ITEM_X_TEST_TAG, 53))
+                )
+            onNodeWithTagAndText(
+                    testTag = String.format(LAZY_LIST_WIREGUARD_PORT_ITEM_X_TEST_TAG, 51820),
+                    text = "51820"
+                )
+                .performClick()
 
-        // Assert
-        verify(exactly = 1) {
-            mockSelectWireguardPortSelectionListener.invoke(Constraint.Only(Port(51820)))
+            // Assert
+            verify(exactly = 1) {
+                mockSelectWireguardPortSelectionListener.invoke(Constraint.Only(Port(51820)))
+            }
         }
-    }
 
     @Test
-    fun testShowWireguardCustomPort() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            VpnSettingsScreen(
-                uiState =
-                    VpnSettingsUiState.createDefault(
-                        customWireguardPort = Constraint.Only(Port(4000))
-                    ),
-            )
-        }
+    fun testShowWireguardCustomPort() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                VpnSettingsScreen(
+                    uiState =
+                        VpnSettingsUiState.createDefault(
+                            customWireguardPort = Constraint.Only(Port(4000))
+                        ),
+                )
+            }
 
-        // Act
-        composeTestRule
-            .onNodeWithTag(LAZY_LIST_TEST_TAG)
-            .performScrollToNode(hasTestTag(LAZY_LIST_WIREGUARD_CUSTOM_PORT_TEXT_TEST_TAG))
+            // Act
+            onNodeWithTag(LAZY_LIST_TEST_TAG)
+                .performScrollToNode(hasTestTag(LAZY_LIST_WIREGUARD_CUSTOM_PORT_TEXT_TEST_TAG))
 
-        // Assert
-        composeTestRule.onNodeWithText("4000").assertExists()
-    }
+            // Assert
+            onNodeWithText("4000").assertExists()
+        }
 
     @Test
-    fun testSelectWireguardCustomPort() {
-        // Arrange
-        val onWireguardPortSelected: (Constraint<Port>) -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            VpnSettingsScreen(
-                uiState =
-                    VpnSettingsUiState.createDefault(
-                        selectedWireguardPort = Constraint.Only(Port(4000)),
-                        customWireguardPort = Constraint.Only(Port(4000))
-                    ),
-                onWireguardPortSelected = onWireguardPortSelected
-            )
-        }
+    fun testSelectWireguardCustomPort() =
+        composeExtension.use {
+            // Arrange
+            val onWireguardPortSelected: (Constraint<Port>) -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                VpnSettingsScreen(
+                    uiState =
+                        VpnSettingsUiState.createDefault(
+                            selectedWireguardPort = Constraint.Only(Port(4000)),
+                            customWireguardPort = Constraint.Only(Port(4000))
+                        ),
+                    onWireguardPortSelected = onWireguardPortSelected
+                )
+            }
 
-        // Act
-        composeTestRule
-            .onNodeWithTag(LAZY_LIST_TEST_TAG)
-            .performScrollToNode(hasTestTag(LAZY_LIST_WIREGUARD_CUSTOM_PORT_TEXT_TEST_TAG))
-        composeTestRule
-            .onNodeWithTag(testTag = LAZY_LIST_WIREGUARD_CUSTOM_PORT_TEXT_TEST_TAG)
-            .performClick()
+            // Act
+            onNodeWithTag(LAZY_LIST_TEST_TAG)
+                .performScrollToNode(hasTestTag(LAZY_LIST_WIREGUARD_CUSTOM_PORT_TEXT_TEST_TAG))
+            onNodeWithTag(testTag = LAZY_LIST_WIREGUARD_CUSTOM_PORT_TEXT_TEST_TAG).performClick()
 
-        // Assert
-        verify { onWireguardPortSelected.invoke(Constraint.Only(Port(4000))) }
-    }
+            // Assert
+            verify { onWireguardPortSelected.invoke(Constraint.Only(Port(4000))) }
+        }
 
     // Navigation Tests
 
     @Test
-    fun testMtuClick() {
-        // Arrange
-        val mockedClickHandler: (Int?) -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            VpnSettingsScreen(
-                uiState = VpnSettingsUiState.createDefault(),
-                navigateToMtuDialog = mockedClickHandler
-            )
-        }
+    fun testMtuClick() =
+        composeExtension.use {
+            // Arrange
+            val mockedClickHandler: (Int?) -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                VpnSettingsScreen(
+                    uiState = VpnSettingsUiState.createDefault(),
+                    navigateToMtuDialog = mockedClickHandler
+                )
+            }
 
-        composeTestRule
-            .onNodeWithTag(LAZY_LIST_TEST_TAG)
-            .performScrollToNode(hasTestTag(LAZY_LIST_LAST_ITEM_TEST_TAG))
+            onNodeWithTag(LAZY_LIST_TEST_TAG)
+                .performScrollToNode(hasTestTag(LAZY_LIST_LAST_ITEM_TEST_TAG))
 
-        // Act
-        composeTestRule.onNodeWithText("WireGuard MTU").performClick()
+            // Act
+            onNodeWithText("WireGuard MTU").performClick()
 
-        // Assert
-        verify { mockedClickHandler.invoke(null) }
-    }
+            // Assert
+            verify { mockedClickHandler.invoke(null) }
+        }
 
     @Test
-    fun testClickAddDns() {
-        // Arrange
-        val mockedClickHandler: (Int?, String?) -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            VpnSettingsScreen(
-                uiState = VpnSettingsUiState.createDefault(isCustomDnsEnabled = true),
-                navigateToDns = mockedClickHandler
-            )
-        }
+    fun testClickAddDns() =
+        composeExtension.use {
+            // Arrange
+            val mockedClickHandler: (Int?, String?) -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                VpnSettingsScreen(
+                    uiState = VpnSettingsUiState.createDefault(isCustomDnsEnabled = true),
+                    navigateToDns = mockedClickHandler
+                )
+            }
 
-        // Act
-        composeTestRule.onNodeWithText("Add a server").performClick()
+            // Act
+            onNodeWithText("Add a server").performClick()
 
-        // Assert
-        verify { mockedClickHandler.invoke(null, null) }
-    }
+            // Assert
+            verify { mockedClickHandler.invoke(null, null) }
+        }
 
     @Test
-    fun testShowTunnelQuantumInfo() {
-        val mockedShowTunnelQuantumInfoClick: () -> Unit = mockk(relaxed = true)
+    fun testShowTunnelQuantumInfo() =
+        composeExtension.use {
+            val mockedShowTunnelQuantumInfoClick: () -> Unit = mockk(relaxed = true)
 
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            VpnSettingsScreen(
-                uiState = VpnSettingsUiState.createDefault(),
-                navigateToQuantumResistanceInfo = mockedShowTunnelQuantumInfoClick
-            )
-        }
+            // Arrange
+            setContentWithTheme {
+                VpnSettingsScreen(
+                    uiState = VpnSettingsUiState.createDefault(),
+                    navigateToQuantumResistanceInfo = mockedShowTunnelQuantumInfoClick
+                )
+            }
 
-        // Act
-        composeTestRule
-            .onNodeWithTag(LAZY_LIST_TEST_TAG)
-            .performScrollToNode(hasTestTag(LAZY_LIST_QUANTUM_ITEM_ON_TEST_TAG))
-        composeTestRule.onNodeWithText("Quantum-resistant tunnel").performClick()
+            // Act
 
-        // Assert
-        verify(exactly = 1) { mockedShowTunnelQuantumInfoClick() }
-    }
+            onNodeWithTag(LAZY_LIST_TEST_TAG)
+                .performScrollToNode(hasTestTag(LAZY_LIST_QUANTUM_ITEM_ON_TEST_TAG))
+            onNodeWithText("Quantum-resistant tunnel").performClick()
+
+            // Assert
+            verify(exactly = 1) { mockedShowTunnelQuantumInfoClick() }
+        }
 
     @Test
-    fun testShowWireguardPortInfo() {
-        val mockedClickHandler: (List<PortRange>) -> Unit = mockk(relaxed = true)
+    fun testShowWireguardPortInfo() =
+        composeExtension.use {
+            val mockedClickHandler: (List<PortRange>) -> Unit = mockk(relaxed = true)
 
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            VpnSettingsScreen(
-                uiState = VpnSettingsUiState.createDefault(),
-                navigateToWireguardPortInfo = mockedClickHandler
-            )
-        }
+            // Arrange
+            setContentWithTheme {
+                VpnSettingsScreen(
+                    uiState = VpnSettingsUiState.createDefault(),
+                    navigateToWireguardPortInfo = mockedClickHandler
+                )
+            }
 
-        composeTestRule.onNodeWithText("WireGuard port").performClick()
+            onNodeWithText("WireGuard port").performClick()
 
-        verify(exactly = 1) { mockedClickHandler.invoke(any()) }
-    }
+            verify(exactly = 1) { mockedClickHandler.invoke(any()) }
+        }
 
     @Test
-    fun testShowWireguardCustomPortDialog() {
-        val mockedClickHandler: () -> Unit = mockk(relaxed = true)
+    fun testShowWireguardCustomPortDialog() =
+        composeExtension.use {
+            val mockedClickHandler: () -> Unit = mockk(relaxed = true)
 
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            VpnSettingsScreen(
-                uiState = VpnSettingsUiState.createDefault(),
-                navigateToWireguardPortDialog = mockedClickHandler
-            )
-        }
+            // Arrange
+            setContentWithTheme {
+                VpnSettingsScreen(
+                    uiState = VpnSettingsUiState.createDefault(),
+                    navigateToWireguardPortDialog = mockedClickHandler
+                )
+            }
 
-        composeTestRule
-            .onNodeWithTag(LAZY_LIST_TEST_TAG)
-            .performScrollToNode(hasTestTag(LAZY_LIST_WIREGUARD_CUSTOM_PORT_NUMBER_TEST_TAG))
-        composeTestRule.onNodeWithText("Custom").performClick()
+            onNodeWithTag(LAZY_LIST_TEST_TAG)
+                .performScrollToNode(hasTestTag(LAZY_LIST_WIREGUARD_CUSTOM_PORT_NUMBER_TEST_TAG))
+            onNodeWithText("Custom").performClick()
 
-        // Assert
-        verify(exactly = 1) { mockedClickHandler.invoke() }
-    }
+            // Assert
+            verify(exactly = 1) { mockedClickHandler.invoke() }
+        }
 
     @Test
-    fun testClickWireguardCustomPortMainCell() {
-        // Arrange
-        val mockOnShowCustomPortDialog: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            VpnSettingsScreen(
-                uiState = VpnSettingsUiState.createDefault(),
-                navigateToWireguardPortDialog = mockOnShowCustomPortDialog
-            )
-        }
+    fun testClickWireguardCustomPortMainCell() =
+        composeExtension.use {
+            // Arrange
+            val mockOnShowCustomPortDialog: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                VpnSettingsScreen(
+                    uiState = VpnSettingsUiState.createDefault(),
+                    navigateToWireguardPortDialog = mockOnShowCustomPortDialog
+                )
+            }
 
-        // Act
-        composeTestRule
-            .onNodeWithTag(LAZY_LIST_TEST_TAG)
-            .performScrollToNode(hasTestTag(LAZY_LIST_WIREGUARD_CUSTOM_PORT_TEXT_TEST_TAG))
-        composeTestRule.onNodeWithTag(LAZY_LIST_WIREGUARD_CUSTOM_PORT_TEXT_TEST_TAG).performClick()
+            // Act
+            onNodeWithTag(LAZY_LIST_TEST_TAG)
+                .performScrollToNode(hasTestTag(LAZY_LIST_WIREGUARD_CUSTOM_PORT_TEXT_TEST_TAG))
+            onNodeWithTag(LAZY_LIST_WIREGUARD_CUSTOM_PORT_TEXT_TEST_TAG).performClick()
 
-        // Assert
-        verify { mockOnShowCustomPortDialog.invoke() }
-    }
+            // Assert
+            verify { mockOnShowCustomPortDialog.invoke() }
+        }
 
     @Test
-    fun testClickWireguardCustomPortNumberCell() {
-        // Arrange
-        val mockOnShowCustomPortDialog: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            VpnSettingsScreen(
-                uiState =
-                    VpnSettingsUiState.createDefault(
-                        selectedWireguardPort = Constraint.Only(Port(4000))
-                    ),
-                navigateToWireguardPortDialog = mockOnShowCustomPortDialog
-            )
-        }
+    fun testClickWireguardCustomPortNumberCell() =
+        composeExtension.use {
+            // Arrange
+            val mockOnShowCustomPortDialog: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                VpnSettingsScreen(
+                    uiState =
+                        VpnSettingsUiState.createDefault(
+                            selectedWireguardPort = Constraint.Only(Port(4000))
+                        ),
+                    navigateToWireguardPortDialog = mockOnShowCustomPortDialog
+                )
+            }
 
-        // Act
-        composeTestRule
-            .onNodeWithTag(LAZY_LIST_TEST_TAG)
-            .performScrollToNode(hasTestTag(LAZY_LIST_WIREGUARD_CUSTOM_PORT_TEXT_TEST_TAG))
-        composeTestRule
-            .onNodeWithTag(testTag = LAZY_LIST_WIREGUARD_CUSTOM_PORT_NUMBER_TEST_TAG)
-            .performClick()
+            // Act
+            onNodeWithTag(LAZY_LIST_TEST_TAG)
+                .performScrollToNode(hasTestTag(LAZY_LIST_WIREGUARD_CUSTOM_PORT_TEXT_TEST_TAG))
+            onNodeWithTag(testTag = LAZY_LIST_WIREGUARD_CUSTOM_PORT_NUMBER_TEST_TAG).performClick()
 
-        // Assert
-        verify { mockOnShowCustomPortDialog.invoke() }
-    }
+            // Assert
+            verify { mockOnShowCustomPortDialog.invoke() }
+        }
 
     companion object {
         private const val LOCAL_DNS_SERVER_WARNING =
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/WelcomeScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/WelcomeScreenTest.kt
index e62b1a399b..b1c9e5f817 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/WelcomeScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/WelcomeScreenTest.kt
@@ -1,9 +1,10 @@
 package net.mullvad.mullvadvpn.compose.screen
 
-import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.onNodeWithTag
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
+import de.mannodermaus.junit5.compose.createComposeExtension
 import io.mockk.MockKAnnotations
 import io.mockk.every
 import io.mockk.mockk
@@ -16,59 +17,59 @@ import net.mullvad.mullvadvpn.lib.payment.model.PaymentProduct
 import net.mullvad.mullvadvpn.lib.payment.model.PaymentStatus
 import net.mullvad.mullvadvpn.lib.payment.model.ProductId
 import net.mullvad.mullvadvpn.lib.payment.model.ProductPrice
-import org.junit.Before
-import org.junit.Rule
-import org.junit.Test
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.RegisterExtension
 
+@OptIn(ExperimentalTestApi::class)
 class WelcomeScreenTest {
-    @get:Rule val composeTestRule = createComposeRule()
+    @JvmField @RegisterExtension val composeExtension = createComposeExtension()
 
-    @Before
+    @BeforeEach
     fun setup() {
         MockKAnnotations.init(this)
     }
 
     @Test
-    fun testDefaultState() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            WelcomeScreen(
-                uiState = WelcomeUiState(),
-                onSitePaymentClick = {},
-                onRedeemVoucherClick = {},
-                onSettingsClick = {},
-                onAccountClick = {},
-                navigateToDeviceInfoDialog = {},
-                onPurchaseBillingProductClick = { _ -> },
-                navigateToVerificationPendingDialog = {}
-            )
-        }
+    fun testDefaultState() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                WelcomeScreen(
+                    uiState = WelcomeUiState(),
+                    onSitePaymentClick = {},
+                    onRedeemVoucherClick = {},
+                    onSettingsClick = {},
+                    onAccountClick = {},
+                    onPurchaseBillingProductClick = { _ -> },
+                    navigateToDeviceInfoDialog = {},
+                    navigateToVerificationPendingDialog = {}
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText("Congrats!").assertExists()
             onNodeWithText("Here’s your account number. Save it!").assertExists()
         }
-    }
 
     @Test
-    fun testDisableSitePayment() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            WelcomeScreen(
-                uiState = WelcomeUiState(),
-                onSitePaymentClick = {},
-                onRedeemVoucherClick = {},
-                onSettingsClick = {},
-                onAccountClick = {},
-                navigateToDeviceInfoDialog = {},
-                onPurchaseBillingProductClick = { _ -> },
-                navigateToVerificationPendingDialog = {}
-            )
-        }
+    fun testDisableSitePayment() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                WelcomeScreen(
+                    uiState = WelcomeUiState(),
+                    onSitePaymentClick = {},
+                    onRedeemVoucherClick = {},
+                    onSettingsClick = {},
+                    onAccountClick = {},
+                    onPurchaseBillingProductClick = { _ -> },
+                    navigateToDeviceInfoDialog = {},
+                    navigateToVerificationPendingDialog = {}
+                )
+            }
 
-        // Assert
-        composeTestRule.apply {
+            // Assert
             onNodeWithText(
                     "Either buy credit on our website or redeem a voucher.",
                     substring = true
@@ -76,242 +77,251 @@ class WelcomeScreenTest {
                 .assertDoesNotExist()
             onNodeWithText("Buy credit").assertDoesNotExist()
         }
-    }
 
     @Test
-    fun testShowAccountNumber() {
-        // Arrange
-        val rawAccountNumber = "1111222233334444"
-        val expectedAccountNumber = "1111 2222 3333 4444"
-        composeTestRule.setContentWithTheme {
-            WelcomeScreen(
-                uiState = WelcomeUiState(accountNumber = rawAccountNumber),
-                onSitePaymentClick = {},
-                onRedeemVoucherClick = {},
-                onSettingsClick = {},
-                onAccountClick = {},
-                onPurchaseBillingProductClick = { _ -> },
-                navigateToDeviceInfoDialog = {},
-                navigateToVerificationPendingDialog = {}
-            )
-        }
+    fun testShowAccountNumber() =
+        composeExtension.use {
+            // Arrange
+            val rawAccountNumber = "1111222233334444"
+            val expectedAccountNumber = "1111 2222 3333 4444"
+            setContentWithTheme {
+                WelcomeScreen(
+                    uiState = WelcomeUiState(accountNumber = rawAccountNumber),
+                    onSitePaymentClick = {},
+                    onRedeemVoucherClick = {},
+                    onSettingsClick = {},
+                    onAccountClick = {},
+                    onPurchaseBillingProductClick = { _ -> },
+                    navigateToDeviceInfoDialog = {},
+                    navigateToVerificationPendingDialog = {}
+                )
+            }
 
-        // Assert
-        composeTestRule.apply { onNodeWithText(expectedAccountNumber).assertExists() }
-    }
+            // Assert
+            onNodeWithText(expectedAccountNumber).assertExists()
+        }
 
     @Test
-    fun testClickSitePaymentButton() {
-        // Arrange
-        val mockClickListener: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            WelcomeScreen(
-                uiState = WelcomeUiState(showSitePayment = true),
-                onSitePaymentClick = mockClickListener,
-                onRedeemVoucherClick = {},
-                onSettingsClick = {},
-                onAccountClick = {},
-                onPurchaseBillingProductClick = { _ -> },
-                navigateToVerificationPendingDialog = {},
-                navigateToDeviceInfoDialog = {}
-            )
-        }
+    fun testClickSitePaymentButton() =
+        composeExtension.use {
+            // Arrange
+            val mockClickListener: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                WelcomeScreen(
+                    uiState = WelcomeUiState(showSitePayment = true),
+                    onSitePaymentClick = mockClickListener,
+                    onRedeemVoucherClick = {},
+                    onSettingsClick = {},
+                    onAccountClick = {},
+                    onPurchaseBillingProductClick = { _ -> },
+                    navigateToDeviceInfoDialog = {},
+                    navigateToVerificationPendingDialog = {}
+                )
+            }
 
-        // Act
-        composeTestRule.apply { onNodeWithText("Buy credit").performClick() }
+            // Act
+            onNodeWithText("Buy credit").performClick()
 
-        // Assert
-        verify(exactly = 1) { mockClickListener.invoke() }
-    }
-
-    @Test
-    fun testClickRedeemVoucher() {
-        // Arrange
-        val mockClickListener: () -> Unit = mockk(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            WelcomeScreen(
-                uiState = WelcomeUiState(),
-                onSitePaymentClick = {},
-                onRedeemVoucherClick = mockClickListener,
-                onSettingsClick = {},
-                onAccountClick = {},
-                onPurchaseBillingProductClick = { _ -> },
-                navigateToVerificationPendingDialog = {},
-                navigateToDeviceInfoDialog = {}
-            )
+            // Assert
+            verify(exactly = 1) { mockClickListener.invoke() }
         }
 
-        // Act
-        composeTestRule.apply { onNodeWithText("Redeem voucher").performClick() }
+    @Test
+    fun testClickRedeemVoucher() =
+        composeExtension.use {
+            // Arrange
+            val mockClickListener: () -> Unit = mockk(relaxed = true)
+            setContentWithTheme {
+                WelcomeScreen(
+                    uiState = WelcomeUiState(),
+                    onSitePaymentClick = {},
+                    onRedeemVoucherClick = mockClickListener,
+                    onSettingsClick = {},
+                    onAccountClick = {},
+                    onPurchaseBillingProductClick = { _ -> },
+                    navigateToDeviceInfoDialog = {},
+                    navigateToVerificationPendingDialog = {}
+                )
+            }
 
-        // Assert
-        verify(exactly = 1) { mockClickListener.invoke() }
-    }
+            // Act
+            onNodeWithText("Redeem voucher").performClick()
 
-    @Test
-    fun testShowBillingErrorPaymentButton() {
-        // Arrange
-        composeTestRule.setContentWithTheme {
-            WelcomeScreen(
-                uiState = WelcomeUiState().copy(billingPaymentState = PaymentState.Error.Billing),
-                onSitePaymentClick = {},
-                onRedeemVoucherClick = {},
-                onSettingsClick = {},
-                onAccountClick = {},
-                onPurchaseBillingProductClick = { _ -> },
-                navigateToVerificationPendingDialog = {},
-                navigateToDeviceInfoDialog = {}
-            )
+            // Assert
+            verify(exactly = 1) { mockClickListener.invoke() }
         }
 
-        // Assert
-        composeTestRule.onNodeWithText("Add 30 days time").assertExists()
-    }
-
     @Test
-    fun testShowBillingPaymentAvailable() {
-        // Arrange
-        val mockPaymentProduct: PaymentProduct = mockk()
-        every { mockPaymentProduct.price } returns ProductPrice("$10")
-        every { mockPaymentProduct.status } returns null
-        composeTestRule.setContentWithTheme {
-            WelcomeScreen(
-                uiState =
-                    WelcomeUiState(
-                        billingPaymentState =
-                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                    ),
-                onSitePaymentClick = {},
-                onRedeemVoucherClick = {},
-                onSettingsClick = {},
-                onAccountClick = {},
-                onPurchaseBillingProductClick = { _ -> },
-                navigateToVerificationPendingDialog = {},
-                navigateToDeviceInfoDialog = {}
-            )
-        }
+    fun testShowBillingErrorPaymentButton() =
+        composeExtension.use {
+            // Arrange
+            setContentWithTheme {
+                WelcomeScreen(
+                    uiState =
+                        WelcomeUiState().copy(billingPaymentState = PaymentState.Error.Billing),
+                    onSitePaymentClick = {},
+                    onRedeemVoucherClick = {},
+                    onSettingsClick = {},
+                    onAccountClick = {},
+                    onPurchaseBillingProductClick = { _ -> },
+                    navigateToDeviceInfoDialog = {},
+                    navigateToVerificationPendingDialog = {}
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText("Add 30 days time ($10)").assertExists()
-    }
+            // Assert
+            onNodeWithText("Add 30 days time").assertExists()
+        }
 
     @Test
-    fun testShowPendingPayment() {
-        // Arrange
-        val mockPaymentProduct: PaymentProduct = mockk()
-        every { mockPaymentProduct.price } returns ProductPrice("$10")
-        every { mockPaymentProduct.status } returns PaymentStatus.PENDING
-        composeTestRule.setContentWithTheme {
-            WelcomeScreen(
-                uiState =
-                    WelcomeUiState()
-                        .copy(
+    fun testShowBillingPaymentAvailable() =
+        composeExtension.use {
+            // Arrange
+            val mockPaymentProduct: PaymentProduct = mockk()
+            every { mockPaymentProduct.price } returns ProductPrice("$10")
+            every { mockPaymentProduct.status } returns null
+            setContentWithTheme {
+                WelcomeScreen(
+                    uiState =
+                        WelcomeUiState(
                             billingPaymentState =
                                 PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
                         ),
-                onSitePaymentClick = {},
-                onRedeemVoucherClick = {},
-                onSettingsClick = {},
-                onAccountClick = {},
-                onPurchaseBillingProductClick = { _ -> },
-                navigateToVerificationPendingDialog = {},
-                navigateToDeviceInfoDialog = {}
-            )
+                    onSitePaymentClick = {},
+                    onRedeemVoucherClick = {},
+                    onSettingsClick = {},
+                    onAccountClick = {},
+                    onPurchaseBillingProductClick = { _ -> },
+                    navigateToDeviceInfoDialog = {},
+                    navigateToVerificationPendingDialog = {}
+                )
+            }
+
+            // Assert
+            onNodeWithText("Add 30 days time ($10)").assertExists()
         }
 
-        // Assert
-        composeTestRule.onNodeWithText("Google Play payment pending").assertExists()
-    }
+    @Test
+    fun testShowPendingPayment() =
+        composeExtension.use {
+            // Arrange
+            val mockPaymentProduct: PaymentProduct = mockk()
+            every { mockPaymentProduct.price } returns ProductPrice("$10")
+            every { mockPaymentProduct.status } returns PaymentStatus.PENDING
+            setContentWithTheme {
+                WelcomeScreen(
+                    uiState =
+                        WelcomeUiState()
+                            .copy(
+                                billingPaymentState =
+                                    PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+                            ),
+                    onSitePaymentClick = {},
+                    onRedeemVoucherClick = {},
+                    onSettingsClick = {},
+                    onAccountClick = {},
+                    onPurchaseBillingProductClick = { _ -> },
+                    navigateToDeviceInfoDialog = {},
+                    navigateToVerificationPendingDialog = {}
+                )
+            }
+
+            // Assert
+            onNodeWithText("Google Play payment pending").assertExists()
+        }
 
     @Test
-    fun testShowPendingPaymentInfoDialog() {
-        // Arrange
-        val mockPaymentProduct: PaymentProduct = mockk()
-        every { mockPaymentProduct.price } returns ProductPrice("$10")
-        every { mockPaymentProduct.status } returns PaymentStatus.PENDING
-        val mockShowPendingInfo = mockk<() -> Unit>(relaxed = true)
-        composeTestRule.setContentWithTheme {
-            WelcomeScreen(
-                uiState =
-                    WelcomeUiState()
-                        .copy(
-                            billingPaymentState =
-                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                        ),
-                onSitePaymentClick = {},
-                onRedeemVoucherClick = {},
-                onSettingsClick = {},
-                onAccountClick = {},
-                onPurchaseBillingProductClick = { _ -> },
-                navigateToVerificationPendingDialog = mockShowPendingInfo,
-                navigateToDeviceInfoDialog = {}
-            )
+    fun testShowPendingPaymentInfoDialog() =
+        composeExtension.use {
+            // Arrange
+            val mockPaymentProduct: PaymentProduct = mockk()
+            every { mockPaymentProduct.price } returns ProductPrice("$10")
+            every { mockPaymentProduct.status } returns PaymentStatus.PENDING
+            val mockShowPendingInfo = mockk<() -> Unit>(relaxed = true)
+            setContentWithTheme {
+                WelcomeScreen(
+                    uiState =
+                        WelcomeUiState()
+                            .copy(
+                                billingPaymentState =
+                                    PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+                            ),
+                    onSitePaymentClick = {},
+                    onRedeemVoucherClick = {},
+                    onSettingsClick = {},
+                    onAccountClick = {},
+                    onPurchaseBillingProductClick = { _ -> },
+                    navigateToVerificationPendingDialog = mockShowPendingInfo,
+                    navigateToDeviceInfoDialog = {}
+                )
+            }
+
+            // Act
+            onNodeWithTag(PLAY_PAYMENT_INFO_ICON_TEST_TAG).performClick()
+
+            // Assert
+            verify(exactly = 1) { mockShowPendingInfo() }
         }
 
-        // Act
-        composeTestRule.onNodeWithTag(PLAY_PAYMENT_INFO_ICON_TEST_TAG).performClick()
+    @Test
+    fun testShowVerificationInProgress() =
+        composeExtension.use {
+            // Arrange
+            val mockPaymentProduct: PaymentProduct = mockk()
+            every { mockPaymentProduct.price } returns ProductPrice("$10")
+            every { mockPaymentProduct.status } returns PaymentStatus.VERIFICATION_IN_PROGRESS
+            setContentWithTheme {
+                WelcomeScreen(
+                    uiState =
+                        WelcomeUiState()
+                            .copy(
+                                billingPaymentState =
+                                    PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+                            ),
+                    onSitePaymentClick = {},
+                    onRedeemVoucherClick = {},
+                    onSettingsClick = {},
+                    onAccountClick = {},
+                    onPurchaseBillingProductClick = { _ -> },
+                    navigateToDeviceInfoDialog = {},
+                    navigateToVerificationPendingDialog = {}
+                )
+            }
 
-        // Assert
-        verify(exactly = 1) { mockShowPendingInfo() }
-    }
+            // Assert
+            onNodeWithText("Verifying purchase").assertExists()
+        }
 
     @Test
-    fun testShowVerificationInProgress() {
-        // Arrange
-        val mockPaymentProduct: PaymentProduct = mockk()
-        every { mockPaymentProduct.price } returns ProductPrice("$10")
-        every { mockPaymentProduct.status } returns PaymentStatus.VERIFICATION_IN_PROGRESS
-        composeTestRule.setContentWithTheme {
-            WelcomeScreen(
-                uiState =
-                    WelcomeUiState()
-                        .copy(
+    fun testOnPurchaseBillingProductClick() =
+        composeExtension.use {
+            // Arrange
+            val clickHandler: (ProductId) -> Unit = mockk(relaxed = true)
+            val mockPaymentProduct: PaymentProduct = mockk()
+            every { mockPaymentProduct.price } returns ProductPrice("$10")
+            every { mockPaymentProduct.productId } returns ProductId("PRODUCT_ID")
+            every { mockPaymentProduct.status } returns null
+            setContentWithTheme {
+                WelcomeScreen(
+                    uiState =
+                        WelcomeUiState(
                             billingPaymentState =
                                 PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
                         ),
-                onSitePaymentClick = {},
-                onRedeemVoucherClick = {},
-                onSettingsClick = {},
-                onAccountClick = {},
-                onPurchaseBillingProductClick = { _ -> },
-                navigateToVerificationPendingDialog = {},
-                navigateToDeviceInfoDialog = {}
-            )
-        }
+                    onSitePaymentClick = {},
+                    onRedeemVoucherClick = {},
+                    onSettingsClick = {},
+                    onAccountClick = {},
+                    onPurchaseBillingProductClick = clickHandler,
+                    navigateToDeviceInfoDialog = {},
+                    navigateToVerificationPendingDialog = {}
+                )
+            }
 
-        // Assert
-        composeTestRule.onNodeWithText("Verifying purchase").assertExists()
-    }
+            // Act
+            onNodeWithText("Add 30 days time ($10)").performClick()
 
-    @Test
-    fun testOnPurchaseBillingProductClick() {
-        // Arrange
-        val clickHandler: (ProductId) -> Unit = mockk(relaxed = true)
-        val mockPaymentProduct: PaymentProduct = mockk()
-        every { mockPaymentProduct.price } returns ProductPrice("$10")
-        every { mockPaymentProduct.productId } returns ProductId("PRODUCT_ID")
-        every { mockPaymentProduct.status } returns null
-        composeTestRule.setContentWithTheme {
-            WelcomeScreen(
-                uiState =
-                    WelcomeUiState(
-                        billingPaymentState =
-                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                    ),
-                onSitePaymentClick = {},
-                onRedeemVoucherClick = {},
-                onSettingsClick = {},
-                onAccountClick = {},
-                onPurchaseBillingProductClick = clickHandler,
-                navigateToVerificationPendingDialog = {},
-                navigateToDeviceInfoDialog = {}
-            )
+            // Assert
+            verify { clickHandler(ProductId("PRODUCT_ID")) }
         }
-
-        // Act
-        composeTestRule.onNodeWithText("Add 30 days time ($10)").performClick()
-
-        // Assert
-        verify { clickHandler(ProductId("PRODUCT_ID")) }
-    }
 }