Remove bad default arguments

31 files changed, 1998 insertions, 1943 deletions

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/dialog/ChangelogDialogTest.kt
index b8165f80fc..b811209d1c 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ChangelogDialogTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/ChangelogDialogTest.kt
@@ -1,12 +1,12 @@
-package net.mullvad.mullvadvpn.compose.screen
+package net.mullvad.mullvadvpn.compose.dialog
 
 import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
+import de.mannodermaus.junit5.compose.ComposeContext
 import io.mockk.MockKAnnotations
 import io.mockk.impl.annotations.MockK
 import net.mullvad.mullvadvpn.compose.createEdgeToEdgeComposeExtension
-import net.mullvad.mullvadvpn.compose.dialog.ChangelogDialog
 import net.mullvad.mullvadvpn.compose.setContentWithTheme
 import net.mullvad.mullvadvpn.viewmodel.AppInfoViewModel
 import net.mullvad.mullvadvpn.viewmodel.ChangelogUiState
@@ -25,16 +25,19 @@ class ChangelogDialogTest {
         MockKAnnotations.init(this)
     }
 
+    private fun ComposeContext.initDialog(state: ChangelogUiState, onDismiss: () -> Unit = {}) {
+        setContentWithTheme { ChangelogDialog(state = state, onDismiss = onDismiss) }
+    }
+
     @Test
     fun testShowChangeLogWhenNeeded() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                ChangelogDialog(
+            initDialog(
+                state =
                     ChangelogUiState(changes = listOf(CHANGELOG_ITEM), version = CHANGELOG_VERSION),
-                    onDismiss = {},
-                )
-            }
+                onDismiss = {},
+            )
 
             // Check changelog content showed within dialog
             onNodeWithText(CHANGELOG_ITEM).assertExists()
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/CreateCustomListDialogTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/CreateCustomListDialogTest.kt
index 915db82438..83dd0c3fdb 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/CreateCustomListDialogTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/CreateCustomListDialogTest.kt
@@ -5,6 +5,7 @@ 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.ComposeContext
 import io.mockk.MockKAnnotations
 import io.mockk.mockk
 import io.mockk.verify
@@ -30,12 +31,28 @@ class CreateCustomListDialogTest {
         MockKAnnotations.init(this)
     }
 
+    private fun ComposeContext.initDialog(
+        state: CreateCustomListUiState = CreateCustomListUiState(),
+        createCustomList: (String) -> Unit = {},
+        onInputChanged: () -> Unit = {},
+        onDismiss: () -> Unit = {},
+    ) {
+        setContentWithTheme {
+            CreateCustomListDialog(
+                state = state,
+                createCustomList = createCustomList,
+                onInputChanged = onInputChanged,
+                onDismiss = onDismiss,
+            )
+        }
+    }
+
     @Test
     fun givenNoErrorShouldShowNoErrorMessage() =
         composeExtension.use {
             // Arrange
             val state = CreateCustomListUiState(error = null)
-            setContentWithTheme { CreateCustomListDialog(state = state) }
+            initDialog(state)
 
             // Assert
             onNodeWithText(NAME_EXIST_ERROR_TEXT).assertDoesNotExist()
@@ -50,7 +67,7 @@ class CreateCustomListDialogTest {
                 CreateCustomListUiState(
                     error = CreateWithLocationsError.Create(CustomListAlreadyExists)
                 )
-            setContentWithTheme { CreateCustomListDialog(state = state) }
+            initDialog(state)
 
             // Assert
             onNodeWithText(NAME_EXIST_ERROR_TEXT).assertExists()
@@ -65,7 +82,7 @@ class CreateCustomListDialogTest {
                 CreateCustomListUiState(
                     error = CreateWithLocationsError.Create(UnknownCustomListError(Throwable()))
                 )
-            setContentWithTheme { CreateCustomListDialog(state = state) }
+            initDialog(state)
 
             // Assert
             onNodeWithText(NAME_EXIST_ERROR_TEXT).assertDoesNotExist()
@@ -78,9 +95,7 @@ class CreateCustomListDialogTest {
             // Arrange
             val mockedOnDismiss: () -> Unit = mockk(relaxed = true)
             val state = CreateCustomListUiState()
-            setContentWithTheme {
-                CreateCustomListDialog(state = state, onDismiss = mockedOnDismiss)
-            }
+            initDialog(state, onDismiss = mockedOnDismiss)
 
             // Act
             onNodeWithText(CANCEL_BUTTON_TEXT).performClick()
@@ -95,9 +110,7 @@ class CreateCustomListDialogTest {
             // Arrange
             val mockedCreateCustomList: (String) -> Unit = mockk(relaxed = true)
             val state = CreateCustomListUiState()
-            setContentWithTheme {
-                CreateCustomListDialog(state = state, createCustomList = mockedCreateCustomList)
-            }
+            initDialog(state, createCustomList = mockedCreateCustomList)
 
             // Act
             onNodeWithText(CREATE_BUTTON_TEXT).performClick()
@@ -113,9 +126,7 @@ class CreateCustomListDialogTest {
             val mockedCreateCustomList: (String) -> Unit = mockk(relaxed = true)
             val inputText = "NEW LIST"
             val state = CreateCustomListUiState()
-            setContentWithTheme {
-                CreateCustomListDialog(state = state, createCustomList = mockedCreateCustomList)
-            }
+            initDialog(state, createCustomList = mockedCreateCustomList)
 
             // Act
             onNodeWithTag(CREATE_CUSTOM_LIST_DIALOG_INPUT_TEST_TAG).performTextInput(inputText)
@@ -132,9 +143,7 @@ class CreateCustomListDialogTest {
             val mockedOnInputChanged: () -> Unit = mockk(relaxed = true)
             val inputText = "NEW LIST"
             val state = CreateCustomListUiState()
-            setContentWithTheme {
-                CreateCustomListDialog(state = state, onInputChanged = mockedOnInputChanged)
-            }
+            initDialog(state, onInputChanged = mockedOnInputChanged)
 
             // Act
             onNodeWithTag(CREATE_CUSTOM_LIST_DIALOG_INPUT_TEST_TAG).performTextInput(inputText)
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 7ca4d7f3d4..075c6cc360 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
@@ -1,7 +1,5 @@
 package net.mullvad.mullvadvpn.compose.dialog
 
-import android.annotation.SuppressLint
-import androidx.compose.runtime.Composable
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
@@ -13,6 +11,7 @@ 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.ComposeContext
 import io.mockk.MockKAnnotations
 import io.mockk.mockk
 import io.mockk.verify
@@ -36,30 +35,30 @@ class CustomPortDialogTest {
         MockKAnnotations.init(this)
     }
 
-    @SuppressLint("ComposableNaming")
-    @Composable
-    private fun testWireguardCustomPortDialog(
+    private fun ComposeContext.initDialog(
         title: String = "",
         portInput: String = "",
         isValidInput: Boolean = false,
-        showResetToDefault: Boolean = false,
         allowedPortRanges: List<PortRange> = emptyList(),
+        showResetToDefault: Boolean = false,
         onInputChanged: (String) -> Unit = { _ -> },
         onSavePort: (String) -> Unit = { _ -> },
         onResetPort: () -> Unit = {},
         onDismiss: () -> Unit = {},
     ) {
-        CustomPortDialog(
-            title = title,
-            portInput = portInput,
-            isValidInput = isValidInput,
-            showResetToDefault = showResetToDefault,
-            allowedPortRanges = allowedPortRanges,
-            onInputChanged = onInputChanged,
-            onSavePort = onSavePort,
-            onDismiss = onDismiss,
-            onResetPort = onResetPort,
-        )
+        setContentWithTheme {
+            CustomPortDialog(
+                title = title,
+                portInput = portInput,
+                isValidInput = isValidInput,
+                allowedPortRanges = allowedPortRanges,
+                showResetToDefault = showResetToDefault,
+                onInputChanged = onInputChanged,
+                onSavePort = onSavePort,
+                onDismiss = onDismiss,
+                onResetPort = onResetPort,
+            )
+        }
     }
 
     @Test
@@ -71,7 +70,17 @@ class CustomPortDialogTest {
             // Arrange
             setContentWithTheme {
                 var input by remember { mutableStateOf("") }
-                testWireguardCustomPortDialog(portInput = input, onInputChanged = { input = it })
+                CustomPortDialog(
+                    title = "",
+                    portInput = input,
+                    isValidInput = false,
+                    allowedPortRanges = emptyList(),
+                    showResetToDefault = false,
+                    onInputChanged = { input = it },
+                    onSavePort = {},
+                    onDismiss = {},
+                    onResetPort = {},
+                )
             }
 
             // Act
@@ -86,7 +95,7 @@ class CustomPortDialogTest {
     fun testEmptyInputResultsInSetPortButtonBeingDisabled() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme { testWireguardCustomPortDialog(isValidInput = false) }
+            initDialog(isValidInput = false)
 
             // Assert
             onNodeWithText("Set port").assertIsNotEnabled()
@@ -96,9 +105,7 @@ class CustomPortDialogTest {
     fun testValidInputResultsInSetPortButtonBeingEnabled() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                testWireguardCustomPortDialog(portInput = VALID_CUSTOM_PORT, isValidInput = true)
-            }
+            initDialog(portInput = VALID_CUSTOM_PORT, isValidInput = true)
 
             // Assert
             onNodeWithText("Set port").assertIsEnabled()
@@ -109,9 +116,7 @@ class CustomPortDialogTest {
     fun testInvalidInputResultsInSetPortButtonBeingDisabled() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                testWireguardCustomPortDialog(portInput = INVALID_CUSTOM_PORT, isValidInput = false)
-            }
+            initDialog(portInput = INVALID_CUSTOM_PORT, isValidInput = false)
 
             // Assert
             onNodeWithText("Set port").assertIsNotEnabled()
@@ -122,13 +127,11 @@ class CustomPortDialogTest {
         composeExtension.use {
             // Arrange
             val mockedSubmitHandler: (String) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                testWireguardCustomPortDialog(
-                    portInput = VALID_CUSTOM_PORT,
-                    isValidInput = true,
-                    onSavePort = mockedSubmitHandler,
-                )
-            }
+            initDialog(
+                portInput = VALID_CUSTOM_PORT,
+                isValidInput = true,
+                onSavePort = mockedSubmitHandler,
+            )
 
             // Act
             onNodeWithText("Set port").assertIsEnabled().performClick()
@@ -142,14 +145,12 @@ class CustomPortDialogTest {
         composeExtension.use {
             // Arrange
             val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                testWireguardCustomPortDialog(
-                    portInput = VALID_CUSTOM_PORT,
-                    isValidInput = true,
-                    showResetToDefault = true,
-                    onResetPort = mockedClickHandler,
-                )
-            }
+            initDialog(
+                portInput = VALID_CUSTOM_PORT,
+                isValidInput = true,
+                showResetToDefault = true,
+                onResetPort = mockedClickHandler,
+            )
 
             // Act
             onNodeWithText("Remove custom port").performClick()
@@ -163,7 +164,7 @@ class CustomPortDialogTest {
         composeExtension.use {
             // Arrange
             val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme { testWireguardCustomPortDialog(onDismiss = mockedClickHandler) }
+            initDialog(onDismiss = mockedClickHandler)
 
             // Assert
             onNodeWithText("Cancel").performClick()
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/DeleteCustomListConfirmationDialogTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/DeleteCustomListConfirmationDialogTest.kt
index abe7abfbcf..6749e620a5 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/DeleteCustomListConfirmationDialogTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/DeleteCustomListConfirmationDialogTest.kt
@@ -3,6 +3,7 @@ package net.mullvad.mullvadvpn.compose.dialog
 import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
+import de.mannodermaus.junit5.compose.ComposeContext
 import io.mockk.MockKAnnotations
 import io.mockk.mockk
 import io.mockk.verify
@@ -25,16 +26,23 @@ class DeleteCustomListConfirmationDialogTest {
         MockKAnnotations.init(this)
     }
 
+    private fun ComposeContext.initDialog(
+        state: DeleteCustomListUiState =
+            DeleteCustomListUiState(CustomListName.fromString("My Custom List"), null),
+        onDelete: () -> Unit = {},
+        onBack: () -> Unit = {},
+    ) {
+        setContentWithTheme {
+            DeleteCustomListConfirmationDialog(state = state, onDelete = onDelete, onBack = onBack)
+        }
+    }
+
     @Test
     fun givenNameShouldShowDeleteNameTitle() =
         composeExtension.use {
             // Arrange
             val name = CustomListName.fromString("List should be deleted")
-            setContentWithTheme {
-                DeleteCustomListConfirmationDialog(
-                    state = DeleteCustomListUiState(name = name, deleteError = null)
-                )
-            }
+            initDialog(state = DeleteCustomListUiState(name = name, deleteError = null))
 
             // Assert
             onNodeWithText(DELETE_TITLE.format(name)).assertExists()
@@ -46,12 +54,10 @@ class DeleteCustomListConfirmationDialogTest {
             // Arrange
             val name = CustomListName.fromString("List should be deleted")
             val mockedOnDelete: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                DeleteCustomListConfirmationDialog(
-                    state = DeleteCustomListUiState(name = name, deleteError = null),
-                    onDelete = mockedOnDelete,
-                )
-            }
+            initDialog(
+                state = DeleteCustomListUiState(name = name, deleteError = null),
+                onDelete = mockedOnDelete,
+            )
 
             // Act
             onNodeWithText(DELETE_BUTTON_TEXT).performClick()
@@ -66,12 +72,10 @@ class DeleteCustomListConfirmationDialogTest {
             // Arrange
             val name = CustomListName.fromString("List should be deleted")
             val mockedOnBack: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                DeleteCustomListConfirmationDialog(
-                    state = DeleteCustomListUiState(name = name, deleteError = null),
-                    onBack = mockedOnBack,
-                )
-            }
+            initDialog(
+                state = DeleteCustomListUiState(name = name, deleteError = null),
+                onBack = mockedOnBack,
+            )
 
             // Act
             onNodeWithText(CANCEL_BUTTON_TEXT).performClick()
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 faafce1137..56bdc562fd 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
@@ -1,10 +1,9 @@
 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.onNodeWithText
+import de.mannodermaus.junit5.compose.ComposeContext
 import net.mullvad.mullvadvpn.compose.createEdgeToEdgeComposeExtension
 import net.mullvad.mullvadvpn.compose.setContentWithTheme
 import net.mullvad.mullvadvpn.viewmodel.DnsDialogViewState
@@ -27,25 +26,29 @@ class DnsDialogTest {
             index = null,
         )
 
-    @SuppressLint("ComposableNaming")
-    @Composable
-    private fun testDnsDialog(
+    private fun ComposeContext.initDialog(
         state: DnsDialogViewState = defaultState,
         onDnsInputChange: (String) -> Unit = { _ -> },
         onSaveDnsClick: () -> Unit = {},
         onRemoveDnsClick: (Int) -> Unit = {},
         onDismiss: () -> Unit = {},
     ) {
-        DnsDialog(state, onDnsInputChange, onSaveDnsClick, onRemoveDnsClick, onDismiss)
+        setContentWithTheme {
+            DnsDialog(
+                state = state,
+                onDnsInputChange = onDnsInputChange,
+                onSaveDnsClick = onSaveDnsClick,
+                onRemoveDnsClick = onRemoveDnsClick,
+                onDismiss = onDismiss,
+            )
+        }
     }
 
     @Test
     fun testDnsDialogLanWarningShownWhenLanTrafficDisabledAndLocalAddressUsed() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                testDnsDialog(defaultState.copy(isAllowLanEnabled = false, isLocal = true))
-            }
+            initDialog(defaultState.copy(isAllowLanEnabled = false, isLocal = true))
 
             // Assert
             onNodeWithText(LOCAL_DNS_SERVER_WARNING).assertExists()
@@ -55,9 +58,7 @@ class DnsDialogTest {
     fun testDnsDialogLanWarningNotShownWhenLanTrafficEnabledAndLocalAddressUsed() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                testDnsDialog(defaultState.copy(isAllowLanEnabled = true, isLocal = true))
-            }
+            initDialog(defaultState.copy(isAllowLanEnabled = true, isLocal = true))
 
             // Assert
             onNodeWithText(LOCAL_DNS_SERVER_WARNING).assertDoesNotExist()
@@ -67,9 +68,7 @@ class DnsDialogTest {
     fun testDnsDialogLanWarningNotShownWhenLanTrafficEnabledAndNonLocalAddressUsed() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                testDnsDialog(defaultState.copy(isAllowLanEnabled = true, isLocal = false))
-            }
+            initDialog(defaultState.copy(isAllowLanEnabled = true, isLocal = false))
 
             // Assert
             onNodeWithText(LOCAL_DNS_SERVER_WARNING).assertDoesNotExist()
@@ -79,9 +78,7 @@ class DnsDialogTest {
     fun testDnsDialogLanWarningNotShownWhenLanTrafficDisabledAndNonLocalAddressUsed() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                testDnsDialog(defaultState.copy(isAllowLanEnabled = false, isLocal = false))
-            }
+            initDialog(defaultState.copy(isAllowLanEnabled = false, isLocal = false))
 
             // Assert
             onNodeWithText(LOCAL_DNS_SERVER_WARNING).assertDoesNotExist()
@@ -91,14 +88,12 @@ class DnsDialogTest {
     fun testDnsDialogSubmitButtonDisabledOnInvalidDnsAddress() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                testDnsDialog(
-                    defaultState.copy(
-                        input = invalidIpAddress,
-                        validationError = ValidationError.InvalidAddress,
-                    )
+            initDialog(
+                defaultState.copy(
+                    input = invalidIpAddress,
+                    validationError = ValidationError.InvalidAddress,
                 )
-            }
+            )
 
             // Assert
             onNodeWithText("Submit").assertIsNotEnabled()
@@ -108,14 +103,12 @@ class DnsDialogTest {
     fun testDnsDialogSubmitButtonDisabledOnDuplicateDnsAddress() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                testDnsDialog(
-                    defaultState.copy(
-                        input = "192.168.0.1",
-                        validationError = ValidationError.DuplicateAddress,
-                    )
+            initDialog(
+                defaultState.copy(
+                    input = "192.168.0.1",
+                    validationError = ValidationError.DuplicateAddress,
                 )
-            }
+            )
 
             // Assert
             onNodeWithText("Submit").assertIsNotEnabled()
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/EditCustomListNameDialogTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/EditCustomListNameDialogTest.kt
index d43a862037..a939b72ccf 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/EditCustomListNameDialogTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/EditCustomListNameDialogTest.kt
@@ -5,6 +5,7 @@ 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.ComposeContext
 import io.mockk.MockKAnnotations
 import io.mockk.mockk
 import io.mockk.verify
@@ -31,12 +32,28 @@ class EditCustomListNameDialogTest {
         MockKAnnotations.init(this)
     }
 
+    private fun ComposeContext.initDialog(
+        state: EditCustomListNameUiState = EditCustomListNameUiState(),
+        updateName: (String) -> Unit = {},
+        onInputChanged: (String) -> Unit = {},
+        onDismiss: () -> Unit = {},
+    ) {
+        setContentWithTheme {
+            EditCustomListNameDialog(
+                state = state,
+                updateName = updateName,
+                onInputChanged = onInputChanged,
+                onDismiss = onDismiss,
+            )
+        }
+    }
+
     @Test
     fun givenNoErrorShouldShowNoErrorMessage() =
         composeExtension.use {
             // Arrange
             val state = EditCustomListNameUiState(error = null)
-            setContentWithTheme { EditCustomListNameDialog(state = state) }
+            initDialog(state = state)
 
             // Assert
             onNodeWithText(NAME_EXIST_ERROR_TEXT).assertDoesNotExist()
@@ -51,7 +68,7 @@ class EditCustomListNameDialogTest {
                 EditCustomListNameUiState(
                     error = RenameError(NameAlreadyExists(CustomListName.fromString("name")))
                 )
-            setContentWithTheme { EditCustomListNameDialog(state = state) }
+            initDialog(state = state)
 
             // Assert
             onNodeWithText(NAME_EXIST_ERROR_TEXT).assertExists()
@@ -66,7 +83,7 @@ class EditCustomListNameDialogTest {
                 EditCustomListNameUiState(
                     error = RenameError(UnknownCustomListError(RuntimeException("")))
                 )
-            setContentWithTheme { EditCustomListNameDialog(state = state) }
+            initDialog(state = state)
 
             // Assert
             onNodeWithText(NAME_EXIST_ERROR_TEXT).assertDoesNotExist()
@@ -79,9 +96,7 @@ class EditCustomListNameDialogTest {
             // Arrange
             val mockedOnDismiss: () -> Unit = mockk(relaxed = true)
             val state = EditCustomListNameUiState()
-            setContentWithTheme {
-                EditCustomListNameDialog(state = state, onDismiss = mockedOnDismiss)
-            }
+            initDialog(state = state, onDismiss = mockedOnDismiss)
 
             // Act
             onNodeWithText(CANCEL_BUTTON_TEXT).performClick()
@@ -96,9 +111,7 @@ class EditCustomListNameDialogTest {
             // Arrange
             val mockedUpdateName: (String) -> Unit = mockk(relaxed = true)
             val state = EditCustomListNameUiState()
-            setContentWithTheme {
-                EditCustomListNameDialog(state = state, updateName = mockedUpdateName)
-            }
+            initDialog(state = state, updateName = mockedUpdateName)
 
             // Act
             onNodeWithText(SAVE_BUTTON_TEXT).performClick()
@@ -114,9 +127,7 @@ class EditCustomListNameDialogTest {
             val mockedUpdateName: (String) -> Unit = mockk(relaxed = true)
             val inputText = "NEW NAME"
             val state = EditCustomListNameUiState(name = inputText)
-            setContentWithTheme {
-                EditCustomListNameDialog(state = state, updateName = mockedUpdateName)
-            }
+            initDialog(state, updateName = mockedUpdateName)
 
             // Act
             onNodeWithText(SAVE_BUTTON_TEXT).performClick()
@@ -132,9 +143,7 @@ class EditCustomListNameDialogTest {
             val mockedOnInputChanged: (String) -> Unit = mockk(relaxed = true)
             val inputText = "NEW NAME"
             val state = EditCustomListNameUiState()
-            setContentWithTheme {
-                EditCustomListNameDialog(state = state, onInputChanged = mockedOnInputChanged)
-            }
+            initDialog(state, onInputChanged = mockedOnInputChanged)
 
             // Act
             onNodeWithTag(EDIT_CUSTOM_LIST_DIALOG_INPUT_TEST_TAG).performTextInput(inputText)
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 9d6b09fe2d..052c7846d9 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
@@ -1,12 +1,11 @@
 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.onNodeWithText
 import androidx.compose.ui.test.performClick
+import de.mannodermaus.junit5.compose.ComposeContext
 import io.mockk.MockKAnnotations
 import io.mockk.mockk
 import io.mockk.verify
@@ -28,31 +27,32 @@ class MtuDialogTest {
         MockKAnnotations.init(this)
     }
 
-    @SuppressLint("ComposableNaming")
-    @Composable
-    private fun testMtuDialog(
-        mtuInput: String = "",
-        isValidInput: Boolean = true,
-        showResetButton: Boolean = true,
-        onInputChanged: (String) -> Unit = { _ -> },
-        onSaveMtu: (String) -> Unit = { _ -> },
+    private val defaultState =
+        MtuDialogUiState(mtuInput = "", isValidInput = true, showResetToDefault = true)
+
+    private fun ComposeContext.initDialog(
+        state: MtuDialogUiState = defaultState,
+        onInputChanged: (String) -> Unit = {},
+        onSaveMtu: (String) -> Unit = {},
         onResetMtu: () -> Unit = {},
         onDismiss: () -> Unit = {},
     ) {
-        MtuDialog(
-            MtuDialogUiState(mtuInput, isValidInput, showResetButton),
-            onInputChanged = onInputChanged,
-            onSaveMtu = onSaveMtu,
-            onResetMtu = onResetMtu,
-            onDismiss = onDismiss,
-        )
+        setContentWithTheme {
+            MtuDialog(
+                state = state,
+                onInputChanged = onInputChanged,
+                onSaveMtu = onSaveMtu,
+                onResetMtu = onResetMtu,
+                onDismiss = onDismiss,
+            )
+        }
     }
 
     @Test
     fun testMtuDialogWithDefaultValue() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme { testMtuDialog() }
+            initDialog()
 
             // Assert
             onNodeWithText(EMPTY_STRING).assertExists()
@@ -62,7 +62,7 @@ class MtuDialogTest {
     fun testMtuDialogWithEditValue() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme { testMtuDialog(mtuInput = VALID_DUMMY_MTU_VALUE) }
+            initDialog(defaultState.copy(mtuInput = VALID_DUMMY_MTU_VALUE))
 
             // Assert
             onNodeWithText(VALID_DUMMY_MTU_VALUE).assertExists()
@@ -73,9 +73,10 @@ class MtuDialogTest {
         composeExtension.use {
             // Arrange
             val mockedSubmitHandler: (String) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                testMtuDialog(VALID_DUMMY_MTU_VALUE, onSaveMtu = mockedSubmitHandler)
-            }
+            initDialog(
+                defaultState.copy(mtuInput = VALID_DUMMY_MTU_VALUE),
+                onSaveMtu = mockedSubmitHandler,
+            )
 
             // Act
             onNodeWithText("Submit").assertIsEnabled().performClick()
@@ -88,7 +89,7 @@ class MtuDialogTest {
     fun testMtuDialogSubmitButtonDisabledWhenInvalidInput() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme { testMtuDialog(INVALID_DUMMY_MTU_VALUE, false) }
+            initDialog(defaultState.copy(mtuInput = INVALID_DUMMY_MTU_VALUE, isValidInput = false))
 
             // Assert
             onNodeWithText("Submit").assertIsNotEnabled()
@@ -99,9 +100,10 @@ class MtuDialogTest {
         composeExtension.use {
             // Arrange
             val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                testMtuDialog(mtuInput = VALID_DUMMY_MTU_VALUE, onResetMtu = mockedClickHandler)
-            }
+            initDialog(
+                defaultState.copy(mtuInput = VALID_DUMMY_MTU_VALUE),
+                onResetMtu = mockedClickHandler,
+            )
 
             // Act
             onNodeWithText("Reset to default").performClick()
@@ -115,7 +117,7 @@ class MtuDialogTest {
         composeExtension.use {
             // Arrange
             val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme { testMtuDialog(onDismiss = mockedClickHandler) }
+            initDialog(onDismiss = mockedClickHandler)
 
             // Assert
             onNodeWithText("Cancel").performClick()
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 3a77a5620a..09a5e9dd72 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
@@ -2,8 +2,10 @@ package net.mullvad.mullvadvpn.compose.dialog
 
 import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.onNodeWithText
+import de.mannodermaus.junit5.compose.ComposeContext
 import net.mullvad.mullvadvpn.compose.createEdgeToEdgeComposeExtension
 import net.mullvad.mullvadvpn.compose.dialog.payment.PaymentDialog
+import net.mullvad.mullvadvpn.compose.dialog.payment.PaymentDialogData
 import net.mullvad.mullvadvpn.compose.setContentWithTheme
 import net.mullvad.mullvadvpn.lib.payment.model.ProductId
 import net.mullvad.mullvadvpn.lib.payment.model.PurchaseResult
@@ -17,15 +19,25 @@ class PaymentDialogTest {
     @RegisterExtension
     val composeExtension = createEdgeToEdgeComposeExtension()
 
+    private fun ComposeContext.initDialog(
+        paymentDialogData: PaymentDialogData,
+        retryPurchase: (ProductId) -> Unit = {},
+        onCloseDialog: (isPaymentSuccessful: Boolean) -> Unit = {},
+    ) {
+        setContentWithTheme {
+            PaymentDialog(
+                paymentDialogData = paymentDialogData,
+                retryPurchase = retryPurchase,
+                onCloseDialog = onCloseDialog,
+            )
+        }
+    }
+
     @Test
     fun testShowPurchaseCompleteDialog() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                PaymentDialog(
-                    paymentDialogData = PurchaseResult.Completed.Success.toPaymentDialogData()!!
-                )
-            }
+            initDialog(paymentDialogData = PurchaseResult.Completed.Success.toPaymentDialogData()!!)
 
             // Assert
             onNodeWithText("Time was successfully added").assertExists()
@@ -35,12 +47,10 @@ class PaymentDialogTest {
     fun testShowVerificationErrorDialog() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                PaymentDialog(
-                    paymentDialogData =
-                        PurchaseResult.Error.VerificationError(null).toPaymentDialogData()!!
-                )
-            }
+            initDialog(
+                paymentDialogData =
+                    PurchaseResult.Error.VerificationError(null).toPaymentDialogData()!!
+            )
 
             // Assert
             onNodeWithText("Verifying purchase").assertExists()
@@ -50,13 +60,11 @@ class PaymentDialogTest {
     fun testShowFetchProductsErrorDialog() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                PaymentDialog(
-                    paymentDialogData =
-                        PurchaseResult.Error.FetchProductsError(ProductId(""), null)
-                            .toPaymentDialogData()!!
-                )
-            }
+            initDialog(
+                paymentDialogData =
+                    PurchaseResult.Error.FetchProductsError(ProductId(""), null)
+                        .toPaymentDialogData()!!
+            )
 
             // Assert
             onNodeWithText("Google Play unavailable").assertExists()
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/dialog/RedeemVoucherDialogTest.kt
index 10e7af7923..b87b43e3c6 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/RedeemVoucherDialogTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/RedeemVoucherDialogTest.kt
@@ -1,15 +1,15 @@
-package net.mullvad.mullvadvpn.compose.screen
+package net.mullvad.mullvadvpn.compose.dialog
 
 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.ComposeContext
 import io.mockk.mockk
 import io.mockk.mockkObject
 import io.mockk.verify
 import net.mullvad.mullvadvpn.compose.createEdgeToEdgeComposeExtension
-import net.mullvad.mullvadvpn.compose.dialog.RedeemVoucherDialog
 import net.mullvad.mullvadvpn.compose.setContentWithTheme
 import net.mullvad.mullvadvpn.compose.state.VoucherDialogState
 import net.mullvad.mullvadvpn.compose.state.VoucherDialogUiState
@@ -29,19 +29,28 @@ class RedeemVoucherDialogTest {
         mockkObject(VoucherRegexHelper)
     }
 
+    private fun ComposeContext.initDialog(
+        state: VoucherDialogUiState = VoucherDialogUiState.INITIAL,
+        onVoucherInputChange: (String) -> Unit = {},
+        onRedeem: (voucherCode: String) -> Unit = {},
+        onDismiss: (isTimeAdded: Boolean) -> Unit = {},
+    ) {
+        setContentWithTheme {
+            RedeemVoucherDialog(
+                state = state,
+                onVoucherInputChange = onVoucherInputChange,
+                onRedeem = onRedeem,
+                onDismiss = onDismiss,
+            )
+        }
+    }
+
     @Test
     fun testDismissDialog() =
         composeExtension.use {
             // Arrange
             val mockedClickHandler: (Boolean) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                RedeemVoucherDialog(
-                    state = VoucherDialogUiState.INITIAL,
-                    onVoucherInputChange = {},
-                    onRedeem = {},
-                    onDismiss = mockedClickHandler,
-                )
-            }
+            initDialog(onDismiss = mockedClickHandler)
 
             // Act
             onNodeWithText(CANCEL_BUTTON_TEXT).performClick()
@@ -55,14 +64,10 @@ class RedeemVoucherDialogTest {
         composeExtension.use {
             // Arrange
             val mockedClickHandler: (Boolean) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                RedeemVoucherDialog(
-                    state = VoucherDialogUiState(voucherState = VoucherDialogState.Success(0)),
-                    onVoucherInputChange = {},
-                    onRedeem = {},
-                    onDismiss = mockedClickHandler,
-                )
-            }
+            initDialog(
+                state = VoucherDialogUiState(voucherState = VoucherDialogState.Success(0)),
+                onDismiss = mockedClickHandler,
+            )
 
             // Act
             onNodeWithText(GOT_IT_BUTTON_TEXT).performClick()
@@ -76,14 +81,7 @@ class RedeemVoucherDialogTest {
         composeExtension.use {
             // Arrange
             val mockedClickHandler: (String) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                RedeemVoucherDialog(
-                    state = VoucherDialogUiState(),
-                    onVoucherInputChange = mockedClickHandler,
-                    onRedeem = {},
-                    onDismiss = {},
-                )
-            }
+            initDialog(state = VoucherDialogUiState(), onVoucherInputChange = mockedClickHandler)
 
             // Act
             onNodeWithTag(VOUCHER_INPUT_TEST_TAG).performTextInput(DUMMY_VOUCHER)
@@ -96,14 +94,7 @@ class RedeemVoucherDialogTest {
     fun testVerifyingState() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                RedeemVoucherDialog(
-                    state = VoucherDialogUiState(voucherState = VoucherDialogState.Verifying),
-                    onVoucherInputChange = {},
-                    onRedeem = {},
-                    onDismiss = {},
-                )
-            }
+            initDialog(state = VoucherDialogUiState(voucherState = VoucherDialogState.Verifying))
 
             // Assert
             onNodeWithText("Verifying voucher…").assertExists()
@@ -113,14 +104,7 @@ class RedeemVoucherDialogTest {
     fun testSuccessState() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                RedeemVoucherDialog(
-                    state = VoucherDialogUiState(voucherState = VoucherDialogState.Success(0)),
-                    onVoucherInputChange = {},
-                    onRedeem = {},
-                    onDismiss = {},
-                )
-            }
+            initDialog(state = VoucherDialogUiState(voucherState = VoucherDialogState.Success(0)))
 
             // Assert
             onNodeWithText("Voucher was successfully redeemed.").assertExists()
@@ -130,18 +114,12 @@ class RedeemVoucherDialogTest {
     fun testErrorState() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                RedeemVoucherDialog(
-                    state =
-                        VoucherDialogUiState(
-                            voucherState =
-                                VoucherDialogState.Error(RedeemVoucherError.InvalidVoucher)
-                        ),
-                    onVoucherInputChange = {},
-                    onRedeem = {},
-                    onDismiss = {},
-                )
-            }
+            initDialog(
+                state =
+                    VoucherDialogUiState(
+                        voucherState = VoucherDialogState.Error(RedeemVoucherError.InvalidVoucher)
+                    )
+            )
 
             // Assert
             onNodeWithText(VOUCHER_CODE_INVALID_ERROR_MESSAGE).assertExists()
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/ResetServerIPOverridesConfirmationDialogTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/ResetServerIPOverridesConfirmationDialogTest.kt
index 107c03f1d6..818cec27f3 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/ResetServerIPOverridesConfirmationDialogTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/ResetServerIPOverridesConfirmationDialogTest.kt
@@ -3,6 +3,7 @@ package net.mullvad.mullvadvpn.compose.dialog
 import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.onNodeWithTag
 import androidx.compose.ui.test.performClick
+import de.mannodermaus.junit5.compose.ComposeContext
 import io.mockk.MockKAnnotations
 import io.mockk.mockk
 import io.mockk.verify
@@ -25,18 +26,25 @@ class ResetServerIPOverridesConfirmationDialogTest {
         MockKAnnotations.init(this)
     }
 
+    private fun ComposeContext.initDialog(
+        onClearAllOverrides: () -> Unit = {},
+        onNavigateBack: () -> Unit = {},
+    ) {
+        setContentWithTheme {
+            ResetServerIpOverridesConfirmationDialog(
+                onClearAllOverrides = onClearAllOverrides,
+                onNavigateBack = onNavigateBack,
+            )
+        }
+    }
+
     @Test
-    fun ensure_cancel_click_works() =
+    fun ensureCancelClickWorks() =
         composeExtension.use {
             val clickHandler: () -> Unit = mockk(relaxed = true)
 
             // Arrange
-            setContentWithTheme {
-                ResetServerIpOverridesConfirmationDialog(
-                    onNavigateBack = clickHandler,
-                    onClearAllOverrides = {},
-                )
-            }
+            initDialog(onNavigateBack = clickHandler)
 
             // Act
             onNodeWithTag(RESET_SERVER_IP_OVERRIDE_CANCEL_TEST_TAG).performClick()
@@ -46,17 +54,12 @@ class ResetServerIPOverridesConfirmationDialogTest {
         }
 
     @Test
-    fun ensure_reset_click_works() =
+    fun ensureResetClickWorks() =
         composeExtension.use {
             val clickHandler: () -> Unit = mockk(relaxed = true)
 
             // Arrange
-            setContentWithTheme {
-                ResetServerIpOverridesConfirmationDialog(
-                    onNavigateBack = {},
-                    onClearAllOverrides = clickHandler,
-                )
-            }
+            initDialog(onClearAllOverrides = clickHandler)
 
             // Act
             onNodeWithTag(RESET_SERVER_IP_OVERRIDE_RESET_TEST_TAG).performClick()
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/SaveApiAccessMethodDialogTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/SaveApiAccessMethodDialogTest.kt
index afd6d0e15d..a9ad169b01 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/SaveApiAccessMethodDialogTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/SaveApiAccessMethodDialogTest.kt
@@ -4,6 +4,7 @@ import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.assertIsNotEnabled
 import androidx.compose.ui.test.onNodeWithTag
 import androidx.compose.ui.test.performClick
+import de.mannodermaus.junit5.compose.ComposeContext
 import io.mockk.mockk
 import io.mockk.verify
 import net.mullvad.mullvadvpn.compose.createEdgeToEdgeComposeExtension
@@ -20,19 +21,27 @@ import org.junit.jupiter.api.extension.RegisterExtension
 class SaveApiAccessMethodDialogTest {
     @JvmField @RegisterExtension val composeExtension = createEdgeToEdgeComposeExtension()
 
+    private fun ComposeContext.initDialog(
+        state: SaveApiAccessMethodUiState = SaveApiAccessMethodUiState(),
+        onCancel: () -> Unit = {},
+        onSave: () -> Unit = {},
+    ) {
+        setContentWithTheme {
+            SaveApiAccessMethodDialog(state = state, onCancel = onCancel, onSave = onSave)
+        }
+    }
+
     @Test
     fun whenTestingInProgressShouldShowSpinnerWithCancelButton() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                SaveApiAccessMethodDialog(
-                    state =
-                        SaveApiAccessMethodUiState(
-                            testingState = TestApiAccessMethodState.Testing,
-                            isSaving = false,
-                        )
-                )
-            }
+            initDialog(
+                state =
+                    SaveApiAccessMethodUiState(
+                        testingState = TestApiAccessMethodState.Testing,
+                        isSaving = false,
+                    )
+            )
 
             // Assert
             onNodeWithTag(SAVE_API_ACCESS_METHOD_LOADING_SPINNER_TEST_TAG).assertExists()
@@ -43,15 +52,13 @@ class SaveApiAccessMethodDialogTest {
     fun whenTestingFailedShouldShowSaveAndCancelButton() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                SaveApiAccessMethodDialog(
-                    state =
-                        SaveApiAccessMethodUiState(
-                            testingState = TestApiAccessMethodState.Result.Failure,
-                            isSaving = false,
-                        )
-                )
-            }
+            initDialog(
+                state =
+                    SaveApiAccessMethodUiState(
+                        testingState = TestApiAccessMethodState.Result.Failure,
+                        isSaving = false,
+                    )
+            )
 
             // Assert
             onNodeWithTag(SAVE_API_ACCESS_METHOD_SAVE_BUTTON_TEST_TAG).assertExists()
@@ -62,15 +69,13 @@ class SaveApiAccessMethodDialogTest {
     fun whenTestingSuccessfulAndSavingShouldShowDisabledCancelButton() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                SaveApiAccessMethodDialog(
-                    state =
-                        SaveApiAccessMethodUiState(
-                            testingState = TestApiAccessMethodState.Result.Successful,
-                            isSaving = true,
-                        )
-                )
-            }
+            initDialog(
+                state =
+                    SaveApiAccessMethodUiState(
+                        testingState = TestApiAccessMethodState.Result.Successful,
+                        isSaving = true,
+                    )
+            )
 
             // Assert
             onNodeWithTag(SAVE_API_ACCESS_METHOD_CANCEL_BUTTON_TEST_TAG).assertExists()
@@ -82,16 +87,14 @@ class SaveApiAccessMethodDialogTest {
         composeExtension.use {
             // Arrange
             val onCancelClick: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                SaveApiAccessMethodDialog(
-                    state =
-                        SaveApiAccessMethodUiState(
-                            testingState = TestApiAccessMethodState.Testing,
-                            isSaving = false,
-                        ),
-                    onCancel = onCancelClick,
-                )
-            }
+            initDialog(
+                state =
+                    SaveApiAccessMethodUiState(
+                        testingState = TestApiAccessMethodState.Testing,
+                        isSaving = false,
+                    ),
+                onCancel = onCancelClick,
+            )
 
             // Act
             onNodeWithTag(SAVE_API_ACCESS_METHOD_CANCEL_BUTTON_TEST_TAG).performClick()
@@ -105,16 +108,14 @@ class SaveApiAccessMethodDialogTest {
         composeExtension.use {
             // Arrange
             val onSaveClick: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                SaveApiAccessMethodDialog(
-                    state =
-                        SaveApiAccessMethodUiState(
-                            testingState = TestApiAccessMethodState.Result.Failure,
-                            isSaving = false,
-                        ),
-                    onSave = onSaveClick,
-                )
-            }
+            initDialog(
+                state =
+                    SaveApiAccessMethodUiState(
+                        testingState = TestApiAccessMethodState.Result.Failure,
+                        isSaving = false,
+                    ),
+                onSave = onSaveClick,
+            )
 
             // Act
             onNodeWithTag(SAVE_API_ACCESS_METHOD_SAVE_BUTTON_TEST_TAG).performClick()
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 dba40b40dd..88b783df83 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
@@ -5,6 +5,7 @@ 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.ComposeContext
 import io.mockk.MockKAnnotations
 import io.mockk.every
 import io.mockk.mockk
@@ -33,23 +34,47 @@ class AccountScreenTest {
         MockKAnnotations.init(this)
     }
 
+    private fun ComposeContext.initScreen(
+        state: AccountUiState = AccountUiState.default(),
+        onCopyAccountNumber: (String) -> Unit = {},
+        onRedeemVoucherClick: () -> Unit = {},
+        onManageAccountClick: () -> Unit = {},
+        onLogoutClick: () -> Unit = {},
+        onPurchaseBillingProductClick: (productId: ProductId) -> Unit = {},
+        navigateToDeviceInfo: () -> Unit = {},
+        navigateToVerificationPendingDialog: () -> Unit = {},
+        onBackClick: () -> Unit = {},
+    ) {
+        setContentWithTheme {
+            AccountScreen(
+                state = state,
+                onCopyAccountNumber = onCopyAccountNumber,
+                onRedeemVoucherClick = onRedeemVoucherClick,
+                onManageAccountClick = onManageAccountClick,
+                onLogoutClick = onLogoutClick,
+                onPurchaseBillingProductClick = onPurchaseBillingProductClick,
+                navigateToDeviceInfo = navigateToDeviceInfo,
+                navigateToVerificationPendingDialog = navigateToVerificationPendingDialog,
+                onBackClick = onBackClick,
+            )
+        }
+    }
+
     @Test
     fun testDefaultState() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                AccountScreen(
-                    state =
-                        AccountUiState(
-                            deviceName = DUMMY_DEVICE_NAME,
-                            accountNumber = DUMMY_ACCOUNT_NUMBER,
-                            accountExpiry = null,
-                            showSitePayment = false,
-                            showLogoutLoading = false,
-                            showManageAccountLoading = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    AccountUiState(
+                        deviceName = DUMMY_DEVICE_NAME,
+                        accountNumber = DUMMY_ACCOUNT_NUMBER,
+                        accountExpiry = null,
+                        showSitePayment = false,
+                        showLogoutLoading = false,
+                        showManageAccountLoading = false,
+                    )
+            )
 
             // Assert
             onNodeWithText("Redeem voucher").assertExists()
@@ -61,20 +86,18 @@ class AccountScreenTest {
         composeExtension.use {
             // Arrange
             val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                AccountScreen(
-                    state =
-                        AccountUiState(
-                            deviceName = DUMMY_DEVICE_NAME,
-                            accountNumber = DUMMY_ACCOUNT_NUMBER,
-                            accountExpiry = null,
-                            showSitePayment = true,
-                            showLogoutLoading = false,
-                            showManageAccountLoading = false,
-                        ),
-                    onManageAccountClick = mockedClickHandler,
-                )
-            }
+            initScreen(
+                state =
+                    AccountUiState(
+                        deviceName = DUMMY_DEVICE_NAME,
+                        accountNumber = DUMMY_ACCOUNT_NUMBER,
+                        accountExpiry = null,
+                        showSitePayment = true,
+                        showLogoutLoading = false,
+                        showManageAccountLoading = false,
+                    ),
+                onManageAccountClick = mockedClickHandler,
+            )
 
             // Act
             onNodeWithText("Manage account").performClick()
@@ -88,20 +111,18 @@ class AccountScreenTest {
         composeExtension.use {
             // Arrange
             val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                AccountScreen(
-                    state =
-                        AccountUiState(
-                            deviceName = DUMMY_DEVICE_NAME,
-                            accountNumber = DUMMY_ACCOUNT_NUMBER,
-                            accountExpiry = null,
-                            showSitePayment = false,
-                            showLogoutLoading = false,
-                            showManageAccountLoading = false,
-                        ),
-                    onRedeemVoucherClick = mockedClickHandler,
-                )
-            }
+            initScreen(
+                state =
+                    AccountUiState(
+                        deviceName = DUMMY_DEVICE_NAME,
+                        accountNumber = DUMMY_ACCOUNT_NUMBER,
+                        accountExpiry = null,
+                        showSitePayment = false,
+                        showLogoutLoading = false,
+                        showManageAccountLoading = false,
+                    ),
+                onRedeemVoucherClick = mockedClickHandler,
+            )
 
             // Act
             onNodeWithText("Redeem voucher").performClick()
@@ -115,20 +136,18 @@ class AccountScreenTest {
         composeExtension.use {
             // Arrange
             val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                AccountScreen(
-                    state =
-                        AccountUiState(
-                            deviceName = DUMMY_DEVICE_NAME,
-                            accountNumber = DUMMY_ACCOUNT_NUMBER,
-                            accountExpiry = null,
-                            showSitePayment = false,
-                            showLogoutLoading = false,
-                            showManageAccountLoading = false,
-                        ),
-                    onLogoutClick = mockedClickHandler,
-                )
-            }
+            initScreen(
+                state =
+                    AccountUiState(
+                        deviceName = DUMMY_DEVICE_NAME,
+                        accountNumber = DUMMY_ACCOUNT_NUMBER,
+                        accountExpiry = null,
+                        showSitePayment = false,
+                        showLogoutLoading = false,
+                        showManageAccountLoading = false,
+                    ),
+                onLogoutClick = mockedClickHandler,
+            )
 
             // Act
             onNodeWithText("Log out").performClick()
@@ -141,13 +160,10 @@ class AccountScreenTest {
     fun testShowBillingErrorPaymentButton() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                AccountScreen(
-                    state =
-                        AccountUiState.default()
-                            .copy(billingPaymentState = PaymentState.Error.Billing)
-                )
-            }
+            initScreen(
+                state =
+                    AccountUiState.default().copy(billingPaymentState = PaymentState.Error.Billing)
+            )
 
             // Assert
             onNodeWithText("Add 30 days time").assertExists()
@@ -160,16 +176,14 @@ class AccountScreenTest {
             val mockPaymentProduct: PaymentProduct = mockk()
             every { mockPaymentProduct.price } returns ProductPrice("$10")
             every { mockPaymentProduct.status } returns null
-            setContentWithTheme {
-                AccountScreen(
-                    state =
-                        AccountUiState.default()
-                            .copy(
-                                billingPaymentState =
-                                    PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                            )
-                )
-            }
+            initScreen(
+                state =
+                    AccountUiState.default()
+                        .copy(
+                            billingPaymentState =
+                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+                        )
+            )
 
             // Assert
             onNodeWithText("Add 30 days time ($10)").assertExists()
@@ -182,16 +196,14 @@ class AccountScreenTest {
             val mockPaymentProduct: PaymentProduct = mockk()
             every { mockPaymentProduct.price } returns ProductPrice("$10")
             every { mockPaymentProduct.status } returns PaymentStatus.PENDING
-            setContentWithTheme {
-                AccountScreen(
-                    state =
-                        AccountUiState.default()
-                            .copy(
-                                billingPaymentState =
-                                    PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                            )
-                )
-            }
+            initScreen(
+                state =
+                    AccountUiState.default()
+                        .copy(
+                            billingPaymentState =
+                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+                        )
+            )
 
             // Assert
             onNodeWithText("Google Play payment pending").assertExists()
@@ -205,17 +217,15 @@ class AccountScreenTest {
             every { mockPaymentProduct.price } returns ProductPrice("$10")
             every { mockPaymentProduct.status } returns PaymentStatus.PENDING
             val mockNavigateToVerificationPending: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                AccountScreen(
-                    state =
-                        AccountUiState.default()
-                            .copy(
-                                billingPaymentState =
-                                    PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                            ),
-                    navigateToVerificationPendingDialog = mockNavigateToVerificationPending,
-                )
-            }
+            initScreen(
+                state =
+                    AccountUiState.default()
+                        .copy(
+                            billingPaymentState =
+                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+                        ),
+                navigateToVerificationPendingDialog = mockNavigateToVerificationPending,
+            )
 
             // Act
             onNodeWithTag(PLAY_PAYMENT_INFO_ICON_TEST_TAG).performClick()
@@ -231,16 +241,14 @@ class AccountScreenTest {
             val mockPaymentProduct: PaymentProduct = mockk()
             every { mockPaymentProduct.price } returns ProductPrice("$10")
             every { mockPaymentProduct.status } returns PaymentStatus.VERIFICATION_IN_PROGRESS
-            setContentWithTheme {
-                AccountScreen(
-                    state =
-                        AccountUiState.default()
-                            .copy(
-                                billingPaymentState =
-                                    PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                            )
-                )
-            }
+            initScreen(
+                state =
+                    AccountUiState.default()
+                        .copy(
+                            billingPaymentState =
+                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+                        )
+            )
 
             // Assert
             onNodeWithText("Verifying purchase").assertExists()
@@ -255,17 +263,15 @@ class AccountScreenTest {
             every { mockPaymentProduct.price } returns ProductPrice("$10")
             every { mockPaymentProduct.productId } returns ProductId("PRODUCT_ID")
             every { mockPaymentProduct.status } returns null
-            setContentWithTheme {
-                AccountScreen(
-                    state =
-                        AccountUiState.default()
-                            .copy(
-                                billingPaymentState =
-                                    PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                            ),
-                    onPurchaseBillingProductClick = clickHandler,
-                )
-            }
+            initScreen(
+                state =
+                    AccountUiState.default()
+                        .copy(
+                            billingPaymentState =
+                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+                        ),
+                onPurchaseBillingProductClick = clickHandler,
+            )
 
             // Act
             onNodeWithText("Add 30 days time ($10)").performClick()
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ApiAccessListScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ApiAccessListScreenTest.kt
index e08d46d9d3..e0f24bcc5e 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ApiAccessListScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ApiAccessListScreenTest.kt
@@ -4,6 +4,7 @@ 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.ComposeContext
 import io.mockk.mockk
 import io.mockk.verify
 import net.mullvad.mullvadvpn.compose.createEdgeToEdgeComposeExtension
@@ -19,17 +20,32 @@ import org.junit.jupiter.api.extension.RegisterExtension
 class ApiAccessListScreenTest {
     @JvmField @RegisterExtension val composeExtension = createEdgeToEdgeComposeExtension()
 
+    private fun ComposeContext.initScreen(
+        state: ApiAccessListUiState = ApiAccessListUiState(),
+        onAddMethodClick: () -> Unit = {},
+        onApiAccessMethodClick: (apiAccessMethodSetting: ApiAccessMethodSetting) -> Unit = {},
+        onApiAccessInfoClick: () -> Unit = {},
+        onBackClick: () -> Unit = {},
+    ) {
+        setContentWithTheme {
+            ApiAccessListScreen(
+                state = state,
+                onAddMethodClick = onAddMethodClick,
+                onApiAccessMethodClick = onApiAccessMethodClick,
+                onApiAccessInfoClick = onApiAccessInfoClick,
+                onBackClick = onBackClick,
+            )
+        }
+    }
+
     @Test
     fun shouldShowCurrentApiAccessName() =
         composeExtension.use {
             // Arrange
             val currentApiAccessMethod = DIRECT_ACCESS_METHOD
-            setContentWithTheme {
-                ApiAccessListScreen(
-                    state =
-                        ApiAccessListUiState(currentApiAccessMethodSetting = currentApiAccessMethod)
-                )
-            }
+            initScreen(
+                state = ApiAccessListUiState(currentApiAccessMethodSetting = currentApiAccessMethod)
+            )
 
             // Assert
             onNodeWithText("Current: ${currentApiAccessMethod.name}")
@@ -40,11 +56,9 @@ class ApiAccessListScreenTest {
         composeExtension.use {
             // Arrange
             val apiAccessMethod = DIRECT_ACCESS_METHOD
-            setContentWithTheme {
-                ApiAccessListScreen(
-                    state = ApiAccessListUiState(apiAccessMethodSettings = listOf(apiAccessMethod))
-                )
-            }
+            initScreen(
+                state = ApiAccessListUiState(apiAccessMethodSettings = listOf(apiAccessMethod))
+            )
 
             // Assert
             onNodeWithText(apiAccessMethod.name.value)
@@ -56,12 +70,7 @@ class ApiAccessListScreenTest {
         composeExtension.use {
             // Arrange
             val onAddMethodClick: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                ApiAccessListScreen(
-                    state = ApiAccessListUiState(),
-                    onAddMethodClick = onAddMethodClick,
-                )
-            }
+            initScreen(state = ApiAccessListUiState(), onAddMethodClick = onAddMethodClick)
 
             // Act
             onNodeWithText("Add").performClick()
@@ -75,12 +84,7 @@ class ApiAccessListScreenTest {
         composeExtension.use {
             // Arrange
             val onApiAccessInfoClick: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                ApiAccessListScreen(
-                    state = ApiAccessListUiState(),
-                    onApiAccessInfoClick = onApiAccessInfoClick,
-                )
-            }
+            initScreen(state = ApiAccessListUiState(), onApiAccessInfoClick = onApiAccessInfoClick)
 
             // Act
             onNodeWithTag(API_ACCESS_LIST_INFO_TEST_TAG).performClick()
@@ -95,12 +99,10 @@ class ApiAccessListScreenTest {
             // Arrange
             val apiAccessMethod = DIRECT_ACCESS_METHOD
             val onApiAccessMethodClick: (ApiAccessMethodSetting) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                ApiAccessListScreen(
-                    state = ApiAccessListUiState(apiAccessMethodSettings = listOf(apiAccessMethod)),
-                    onApiAccessMethodClick = onApiAccessMethodClick,
-                )
-            }
+            initScreen(
+                state = ApiAccessListUiState(apiAccessMethodSettings = listOf(apiAccessMethod)),
+                onApiAccessMethodClick = onApiAccessMethodClick,
+            )
 
             // Act
             onNodeWithText(apiAccessMethod.name.value).performClick()
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ApiAccessMethodDetailsScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ApiAccessMethodDetailsScreenTest.kt
index 93f3fb0343..14dfdfbfd7 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ApiAccessMethodDetailsScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ApiAccessMethodDetailsScreenTest.kt
@@ -4,6 +4,7 @@ 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.ComposeContext
 import io.mockk.mockk
 import io.mockk.verify
 import net.mullvad.mullvadvpn.compose.createEdgeToEdgeComposeExtension
@@ -23,22 +24,44 @@ import org.junit.jupiter.api.extension.RegisterExtension
 class ApiAccessMethodDetailsScreenTest {
     @JvmField @RegisterExtension val composeExtension = createEdgeToEdgeComposeExtension()
 
+    private fun ComposeContext.initScreen(
+        state: ApiAccessMethodDetailsUiState,
+        onEditMethodClicked: () -> Unit = {},
+        onEnableClicked: (Boolean) -> Unit = {},
+        onTestMethodClicked: () -> Unit = {},
+        onUseMethodClicked: () -> Unit = {},
+        onDeleteApiAccessMethodClicked: (ApiAccessMethodId) -> Unit = {},
+        onNavigateToEncryptedDnsInfoDialog: () -> Unit = {},
+        onBackClicked: () -> Unit = {},
+    ) {
+        setContentWithTheme {
+            ApiAccessMethodDetailsScreen(
+                state = state,
+                onEditMethodClicked = onEditMethodClicked,
+                onEnableClicked = onEnableClicked,
+                onTestMethodClicked = onTestMethodClicked,
+                onUseMethodClicked = onUseMethodClicked,
+                onDeleteApiAccessMethodClicked = onDeleteApiAccessMethodClicked,
+                onNavigateToEncryptedDnsInfoDialog = onNavigateToEncryptedDnsInfoDialog,
+                onBackClicked = onBackClicked,
+            )
+        }
+    }
+
     @Test
     fun whenApiAccessMethodIsNotEditableShouldNotShowDeleteAndEdit() =
         composeExtension.use {
             // Arrange
             val apiAccessMethod = DIRECT_ACCESS_METHOD
-            setContentWithTheme {
-                ApiAccessMethodDetailsScreen(
-                    state =
-                        ApiAccessMethodDetailsUiState.Content(
-                            apiAccessMethodSetting = apiAccessMethod,
-                            isDisableable = true,
-                            isCurrentMethod = true,
-                            isTestingAccessMethod = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    ApiAccessMethodDetailsUiState.Content(
+                        apiAccessMethodSetting = apiAccessMethod,
+                        isDisableable = true,
+                        isCurrentMethod = true,
+                        isTestingAccessMethod = false,
+                    )
+            )
 
             // Assert
             onNodeWithTag(API_ACCESS_DETAILS_TOP_BAR_DROPDOWN_BUTTON_TEST_TAG).assertDoesNotExist()
@@ -51,18 +74,16 @@ class ApiAccessMethodDetailsScreenTest {
             // Arrange
             val onEnableClicked: (Boolean) -> Unit = mockk(relaxed = true)
             val apiAccessMethod = DIRECT_ACCESS_METHOD
-            setContentWithTheme {
-                ApiAccessMethodDetailsScreen(
-                    state =
-                        ApiAccessMethodDetailsUiState.Content(
-                            apiAccessMethodSetting = apiAccessMethod,
-                            isDisableable = false,
-                            isCurrentMethod = true,
-                            isTestingAccessMethod = false,
-                        ),
-                    onEnableClicked = onEnableClicked,
-                )
-            }
+            initScreen(
+                state =
+                    ApiAccessMethodDetailsUiState.Content(
+                        apiAccessMethodSetting = apiAccessMethod,
+                        isDisableable = false,
+                        isCurrentMethod = true,
+                        isTestingAccessMethod = false,
+                    ),
+                onEnableClicked = onEnableClicked,
+            )
 
             // Act
             onNodeWithText("Enable method").performClick()
@@ -78,18 +99,16 @@ class ApiAccessMethodDetailsScreenTest {
             // Arrange
             val onDeleteApiAccessMethodClicked: (ApiAccessMethodId) -> Unit = mockk(relaxed = true)
             val apiAccessMethod = CUSTOM_ACCESS_METHOD
-            setContentWithTheme {
-                ApiAccessMethodDetailsScreen(
-                    state =
-                        ApiAccessMethodDetailsUiState.Content(
-                            apiAccessMethodSetting = apiAccessMethod,
-                            isDisableable = false,
-                            isCurrentMethod = true,
-                            isTestingAccessMethod = false,
-                        ),
-                    onDeleteApiAccessMethodClicked = onDeleteApiAccessMethodClicked,
-                )
-            }
+            initScreen(
+                state =
+                    ApiAccessMethodDetailsUiState.Content(
+                        apiAccessMethodSetting = apiAccessMethod,
+                        isDisableable = false,
+                        isCurrentMethod = true,
+                        isTestingAccessMethod = false,
+                    ),
+                onDeleteApiAccessMethodClicked = onDeleteApiAccessMethodClicked,
+            )
 
             // Act
             onNodeWithTag(API_ACCESS_DETAILS_TOP_BAR_DROPDOWN_BUTTON_TEST_TAG).performClick()
@@ -105,18 +124,16 @@ class ApiAccessMethodDetailsScreenTest {
             // Arrange
             val onEditMethodClicked: () -> Unit = mockk(relaxed = true)
             val apiAccessMethod = CUSTOM_ACCESS_METHOD
-            setContentWithTheme {
-                ApiAccessMethodDetailsScreen(
-                    state =
-                        ApiAccessMethodDetailsUiState.Content(
-                            apiAccessMethodSetting = apiAccessMethod,
-                            isDisableable = false,
-                            isCurrentMethod = true,
-                            isTestingAccessMethod = false,
-                        ),
-                    onEditMethodClicked = onEditMethodClicked,
-                )
-            }
+            initScreen(
+                state =
+                    ApiAccessMethodDetailsUiState.Content(
+                        apiAccessMethodSetting = apiAccessMethod,
+                        isDisableable = false,
+                        isCurrentMethod = true,
+                        isTestingAccessMethod = false,
+                    ),
+                onEditMethodClicked = onEditMethodClicked,
+            )
 
             // Act
             onNodeWithTag(API_ACCESS_DETAILS_EDIT_BUTTON).performClick()
@@ -131,18 +148,16 @@ class ApiAccessMethodDetailsScreenTest {
             // Arrange
             val onEnableClicked: (Boolean) -> Unit = mockk(relaxed = true)
             val apiAccessMethod = DIRECT_ACCESS_METHOD
-            setContentWithTheme {
-                ApiAccessMethodDetailsScreen(
-                    state =
-                        ApiAccessMethodDetailsUiState.Content(
-                            apiAccessMethodSetting = apiAccessMethod,
-                            isDisableable = true,
-                            isCurrentMethod = true,
-                            isTestingAccessMethod = false,
-                        ),
-                    onEnableClicked = onEnableClicked,
-                )
-            }
+            initScreen(
+                state =
+                    ApiAccessMethodDetailsUiState.Content(
+                        apiAccessMethodSetting = apiAccessMethod,
+                        isDisableable = true,
+                        isCurrentMethod = true,
+                        isTestingAccessMethod = false,
+                    ),
+                onEnableClicked = onEnableClicked,
+            )
 
             // Act
             onNodeWithText("Enable method").performClick()
@@ -157,18 +172,16 @@ class ApiAccessMethodDetailsScreenTest {
             // Arrange
             val onTestMethodClicked: () -> Unit = mockk(relaxed = true)
             val apiAccessMethod = DIRECT_ACCESS_METHOD
-            setContentWithTheme {
-                ApiAccessMethodDetailsScreen(
-                    state =
-                        ApiAccessMethodDetailsUiState.Content(
-                            apiAccessMethodSetting = apiAccessMethod,
-                            isDisableable = true,
-                            isCurrentMethod = true,
-                            isTestingAccessMethod = false,
-                        ),
-                    onTestMethodClicked = onTestMethodClicked,
-                )
-            }
+            initScreen(
+                state =
+                    ApiAccessMethodDetailsUiState.Content(
+                        apiAccessMethodSetting = apiAccessMethod,
+                        isDisableable = true,
+                        isCurrentMethod = true,
+                        isTestingAccessMethod = false,
+                    ),
+                onTestMethodClicked = onTestMethodClicked,
+            )
 
             // Act
             onNodeWithTag(API_ACCESS_TEST_METHOD_BUTTON).performClick()
@@ -183,18 +196,16 @@ class ApiAccessMethodDetailsScreenTest {
             // Arrange
             val onUseMethodClicked: () -> Unit = mockk(relaxed = true)
             val apiAccessMethod = DIRECT_ACCESS_METHOD
-            setContentWithTheme {
-                ApiAccessMethodDetailsScreen(
-                    state =
-                        ApiAccessMethodDetailsUiState.Content(
-                            apiAccessMethodSetting = apiAccessMethod,
-                            isDisableable = true,
-                            isCurrentMethod = false,
-                            isTestingAccessMethod = false,
-                        ),
-                    onUseMethodClicked = onUseMethodClicked,
-                )
-            }
+            initScreen(
+                state =
+                    ApiAccessMethodDetailsUiState.Content(
+                        apiAccessMethodSetting = apiAccessMethod,
+                        isDisableable = true,
+                        isCurrentMethod = false,
+                        isTestingAccessMethod = false,
+                    ),
+                onUseMethodClicked = onUseMethodClicked,
+            )
 
             // Act
             onNodeWithTag(API_ACCESS_USE_METHOD_BUTTON).performClick()
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 9087fe2f6c..be998a3728 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
@@ -4,6 +4,7 @@ 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.ComposeContext
 import io.mockk.MockKAnnotations
 import io.mockk.every
 import io.mockk.mockk
@@ -52,11 +53,41 @@ class ConnectScreenTest {
         unmockkAll()
     }
 
+    private fun ComposeContext.initScreen(
+        state: ConnectUiState = ConnectUiState.INITIAL,
+        onDisconnectClick: () -> Unit = {},
+        onReconnectClick: () -> Unit = {},
+        onConnectClick: () -> Unit = {},
+        onCancelClick: () -> Unit = {},
+        onSwitchLocationClick: () -> Unit = {},
+        onOpenAppListing: () -> Unit = {},
+        onManageAccountClick: () -> Unit = {},
+        onSettingsClick: () -> Unit = {},
+        onAccountClick: () -> Unit = {},
+        onDismissNewDeviceClick: () -> Unit = {},
+    ) {
+        setContentWithTheme {
+            ConnectScreen(
+                state = state,
+                onDisconnectClick = onDisconnectClick,
+                onReconnectClick = onReconnectClick,
+                onConnectClick = onConnectClick,
+                onCancelClick = onCancelClick,
+                onSwitchLocationClick = onSwitchLocationClick,
+                onOpenAppListing = onOpenAppListing,
+                onManageAccountClick = onManageAccountClick,
+                onSettingsClick = onSettingsClick,
+                onAccountClick = onAccountClick,
+                onDismissNewDeviceClick = onDismissNewDeviceClick,
+            )
+        }
+    }
+
     @Test
     fun testDefaultState() {
         composeExtension.use {
             // Arrange
-            setContentWithTheme { ConnectScreen(state = ConnectUiState.INITIAL) }
+            initScreen()
 
             // Assert
             onNodeWithText("DISCONNECTED").assertExists()
@@ -68,21 +99,19 @@ class ConnectScreenTest {
     fun testConnectingState() {
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                ConnectScreen(
-                    state =
-                        ConnectUiState(
-                            location = null,
-                            selectedRelayItemTitle = null,
-                            tunnelState = TunnelState.Connecting(null, null, emptyList()),
-                            showLocation = false,
-                            deviceName = "",
-                            daysLeftUntilExpiry = null,
-                            inAppNotification = InAppNotification.TunnelStateBlocked,
-                            isPlayBuild = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    ConnectUiState(
+                        location = null,
+                        selectedRelayItemTitle = null,
+                        tunnelState = TunnelState.Connecting(null, null, emptyList()),
+                        showLocation = false,
+                        deviceName = "",
+                        daysLeftUntilExpiry = null,
+                        inAppNotification = InAppNotification.TunnelStateBlocked,
+                        isPlayBuild = false,
+                    )
+            )
 
             // Assert
             onNodeWithTag(CIRCULAR_PROGRESS_INDICATOR).assertExists()
@@ -98,22 +127,19 @@ class ConnectScreenTest {
         composeExtension.use {
             // Arrange
             val mockTunnelEndpoint: TunnelEndpoint = mockk(relaxed = true)
-            setContentWithTheme {
-                ConnectScreen(
-                    state =
-                        ConnectUiState(
-                            location = null,
-                            selectedRelayItemTitle = null,
-                            tunnelState =
-                                TunnelState.Connected(mockTunnelEndpoint, null, emptyList()),
-                            showLocation = false,
-                            deviceName = "",
-                            daysLeftUntilExpiry = null,
-                            inAppNotification = null,
-                            isPlayBuild = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    ConnectUiState(
+                        location = null,
+                        selectedRelayItemTitle = null,
+                        tunnelState = TunnelState.Connected(mockTunnelEndpoint, null, emptyList()),
+                        showLocation = false,
+                        deviceName = "",
+                        daysLeftUntilExpiry = null,
+                        inAppNotification = null,
+                        isPlayBuild = false,
+                    )
+            )
 
             // Assert
             onNodeWithText("CONNECTED").assertExists()
@@ -127,21 +153,19 @@ class ConnectScreenTest {
         composeExtension.use {
             // Arrange
             val mockLocationName = "Home"
-            setContentWithTheme {
-                ConnectScreen(
-                    state =
-                        ConnectUiState(
-                            location = null,
-                            selectedRelayItemTitle = mockLocationName,
-                            tunnelState = TunnelState.Disconnecting(ActionAfterDisconnect.Nothing),
-                            showLocation = true,
-                            deviceName = "",
-                            daysLeftUntilExpiry = null,
-                            inAppNotification = null,
-                            isPlayBuild = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    ConnectUiState(
+                        location = null,
+                        selectedRelayItemTitle = mockLocationName,
+                        tunnelState = TunnelState.Disconnecting(ActionAfterDisconnect.Nothing),
+                        showLocation = true,
+                        deviceName = "",
+                        daysLeftUntilExpiry = null,
+                        inAppNotification = null,
+                        isPlayBuild = false,
+                    )
+            )
 
             // Assert
             onNodeWithText("DISCONNECTED").assertExists()
@@ -155,21 +179,19 @@ class ConnectScreenTest {
         composeExtension.use {
             // Arrange
             val mockLocationName = "Home"
-            setContentWithTheme {
-                ConnectScreen(
-                    state =
-                        ConnectUiState(
-                            location = null,
-                            selectedRelayItemTitle = mockLocationName,
-                            tunnelState = TunnelState.Disconnected(),
-                            showLocation = true,
-                            deviceName = "",
-                            daysLeftUntilExpiry = null,
-                            inAppNotification = null,
-                            isPlayBuild = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    ConnectUiState(
+                        location = null,
+                        selectedRelayItemTitle = mockLocationName,
+                        tunnelState = TunnelState.Disconnected(),
+                        showLocation = true,
+                        deviceName = "",
+                        daysLeftUntilExpiry = null,
+                        inAppNotification = null,
+                        isPlayBuild = false,
+                    )
+            )
 
             // Assert
             onNodeWithText("DISCONNECTED").assertExists()
@@ -183,27 +205,23 @@ class ConnectScreenTest {
         composeExtension.use {
             // Arrange
             val mockLocationName = "Home"
-            setContentWithTheme {
-                ConnectScreen(
-                    state =
-                        ConnectUiState(
-                            location = null,
-                            selectedRelayItemTitle = mockLocationName,
-                            tunnelState =
-                                TunnelState.Error(
-                                    ErrorState(ErrorStateCause.StartTunnelError, true)
-                                ),
-                            showLocation = true,
-                            deviceName = "",
-                            daysLeftUntilExpiry = null,
-                            inAppNotification =
-                                InAppNotification.TunnelStateError(
-                                    ErrorState(ErrorStateCause.StartTunnelError, true)
-                                ),
-                            isPlayBuild = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    ConnectUiState(
+                        location = null,
+                        selectedRelayItemTitle = mockLocationName,
+                        tunnelState =
+                            TunnelState.Error(ErrorState(ErrorStateCause.StartTunnelError, true)),
+                        showLocation = true,
+                        deviceName = "",
+                        daysLeftUntilExpiry = null,
+                        inAppNotification =
+                            InAppNotification.TunnelStateError(
+                                ErrorState(ErrorStateCause.StartTunnelError, true)
+                            ),
+                        isPlayBuild = false,
+                    )
+            )
 
             // Assert
             onNodeWithText("BLOCKED CONNECTION").assertExists()
@@ -218,27 +236,23 @@ class ConnectScreenTest {
         composeExtension.use {
             // Arrange
             val mockLocationName = "Home"
-            setContentWithTheme {
-                ConnectScreen(
-                    state =
-                        ConnectUiState(
-                            location = null,
-                            selectedRelayItemTitle = mockLocationName,
-                            tunnelState =
-                                TunnelState.Error(
-                                    ErrorState(ErrorStateCause.StartTunnelError, false)
-                                ),
-                            showLocation = true,
-                            deviceName = "",
-                            daysLeftUntilExpiry = null,
-                            inAppNotification =
-                                InAppNotification.TunnelStateError(
-                                    ErrorState(ErrorStateCause.StartTunnelError, false)
-                                ),
-                            isPlayBuild = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    ConnectUiState(
+                        location = null,
+                        selectedRelayItemTitle = mockLocationName,
+                        tunnelState =
+                            TunnelState.Error(ErrorState(ErrorStateCause.StartTunnelError, false)),
+                        showLocation = true,
+                        deviceName = "",
+                        daysLeftUntilExpiry = null,
+                        inAppNotification =
+                            InAppNotification.TunnelStateError(
+                                ErrorState(ErrorStateCause.StartTunnelError, false)
+                            ),
+                        isPlayBuild = false,
+                    )
+            )
 
             // Assert
             onNodeWithText("FAILED TO CONNECT").assertExists()
@@ -253,22 +267,19 @@ class ConnectScreenTest {
     fun testReconnectingState() {
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                ConnectScreen(
-                    state =
-                        ConnectUiState(
-                            location = null,
-                            selectedRelayItemTitle = null,
-                            tunnelState =
-                                TunnelState.Disconnecting(ActionAfterDisconnect.Reconnect),
-                            showLocation = false,
-                            deviceName = "",
-                            daysLeftUntilExpiry = null,
-                            inAppNotification = InAppNotification.TunnelStateBlocked,
-                            isPlayBuild = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    ConnectUiState(
+                        location = null,
+                        selectedRelayItemTitle = null,
+                        tunnelState = TunnelState.Disconnecting(ActionAfterDisconnect.Reconnect),
+                        showLocation = false,
+                        deviceName = "",
+                        daysLeftUntilExpiry = null,
+                        inAppNotification = InAppNotification.TunnelStateBlocked,
+                        isPlayBuild = false,
+                    )
+            )
 
             // Assert
             onNodeWithTag(CIRCULAR_PROGRESS_INDICATOR).assertExists()
@@ -284,21 +295,19 @@ class ConnectScreenTest {
         composeExtension.use {
             // Arrange
             val mockLocationName = "Home"
-            setContentWithTheme {
-                ConnectScreen(
-                    state =
-                        ConnectUiState(
-                            location = null,
-                            selectedRelayItemTitle = mockLocationName,
-                            tunnelState = TunnelState.Disconnecting(ActionAfterDisconnect.Block),
-                            showLocation = true,
-                            deviceName = "",
-                            daysLeftUntilExpiry = null,
-                            inAppNotification = InAppNotification.TunnelStateBlocked,
-                            isPlayBuild = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    ConnectUiState(
+                        location = null,
+                        selectedRelayItemTitle = mockLocationName,
+                        tunnelState = TunnelState.Disconnecting(ActionAfterDisconnect.Block),
+                        showLocation = true,
+                        deviceName = "",
+                        daysLeftUntilExpiry = null,
+                        inAppNotification = InAppNotification.TunnelStateBlocked,
+                        isPlayBuild = false,
+                    )
+            )
 
             // Assert
             onNodeWithText("CONNECTED").assertExists()
@@ -314,22 +323,20 @@ class ConnectScreenTest {
             // Arrange
             val mockLocationName = "Home"
             val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                ConnectScreen(
-                    state =
-                        ConnectUiState(
-                            location = null,
-                            selectedRelayItemTitle = mockLocationName,
-                            tunnelState = TunnelState.Disconnected(),
-                            showLocation = false,
-                            deviceName = "",
-                            daysLeftUntilExpiry = null,
-                            inAppNotification = null,
-                            isPlayBuild = false,
-                        ),
-                    onSwitchLocationClick = mockedClickHandler,
-                )
-            }
+            initScreen(
+                state =
+                    ConnectUiState(
+                        location = null,
+                        selectedRelayItemTitle = mockLocationName,
+                        tunnelState = TunnelState.Disconnected(),
+                        showLocation = false,
+                        deviceName = "",
+                        daysLeftUntilExpiry = null,
+                        inAppNotification = null,
+                        isPlayBuild = false,
+                    ),
+                onSwitchLocationClick = mockedClickHandler,
+            )
 
             // Act
             onNodeWithTag(SELECT_LOCATION_BUTTON_TEST_TAG).performClick()
@@ -345,23 +352,20 @@ class ConnectScreenTest {
             // Arrange
             val mockTunnelEndpoint: TunnelEndpoint = mockk(relaxed = true)
             val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                ConnectScreen(
-                    state =
-                        ConnectUiState(
-                            location = null,
-                            selectedRelayItemTitle = null,
-                            tunnelState =
-                                TunnelState.Connected(mockTunnelEndpoint, null, emptyList()),
-                            showLocation = false,
-                            deviceName = "",
-                            daysLeftUntilExpiry = null,
-                            inAppNotification = null,
-                            isPlayBuild = false,
-                        ),
-                    onDisconnectClick = mockedClickHandler,
-                )
-            }
+            initScreen(
+                state =
+                    ConnectUiState(
+                        location = null,
+                        selectedRelayItemTitle = null,
+                        tunnelState = TunnelState.Connected(mockTunnelEndpoint, null, emptyList()),
+                        showLocation = false,
+                        deviceName = "",
+                        daysLeftUntilExpiry = null,
+                        inAppNotification = null,
+                        isPlayBuild = false,
+                    ),
+                onDisconnectClick = mockedClickHandler,
+            )
 
             // Act
             onNodeWithTag(CONNECT_BUTTON_TEST_TAG).performClick()
@@ -377,23 +381,20 @@ class ConnectScreenTest {
             // Arrange
             val mockTunnelEndpoint: TunnelEndpoint = mockk(relaxed = true)
             val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                ConnectScreen(
-                    state =
-                        ConnectUiState(
-                            location = null,
-                            selectedRelayItemTitle = null,
-                            tunnelState =
-                                TunnelState.Connected(mockTunnelEndpoint, null, emptyList()),
-                            showLocation = false,
-                            deviceName = "",
-                            daysLeftUntilExpiry = null,
-                            inAppNotification = null,
-                            isPlayBuild = false,
-                        ),
-                    onReconnectClick = mockedClickHandler,
-                )
-            }
+            initScreen(
+                state =
+                    ConnectUiState(
+                        location = null,
+                        selectedRelayItemTitle = null,
+                        tunnelState = TunnelState.Connected(mockTunnelEndpoint, null, emptyList()),
+                        showLocation = false,
+                        deviceName = "",
+                        daysLeftUntilExpiry = null,
+                        inAppNotification = null,
+                        isPlayBuild = false,
+                    ),
+                onReconnectClick = mockedClickHandler,
+            )
 
             // Act
             onNodeWithTag(RECONNECT_BUTTON_TEST_TAG).performClick()
@@ -408,22 +409,20 @@ class ConnectScreenTest {
         composeExtension.use {
             // Arrange
             val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                ConnectScreen(
-                    state =
-                        ConnectUiState(
-                            location = null,
-                            selectedRelayItemTitle = null,
-                            tunnelState = TunnelState.Disconnected(),
-                            showLocation = false,
-                            deviceName = "",
-                            daysLeftUntilExpiry = null,
-                            inAppNotification = null,
-                            isPlayBuild = false,
-                        ),
-                    onConnectClick = mockedClickHandler,
-                )
-            }
+            initScreen(
+                state =
+                    ConnectUiState(
+                        location = null,
+                        selectedRelayItemTitle = null,
+                        tunnelState = TunnelState.Disconnected(),
+                        showLocation = false,
+                        deviceName = "",
+                        daysLeftUntilExpiry = null,
+                        inAppNotification = null,
+                        isPlayBuild = false,
+                    ),
+                onConnectClick = mockedClickHandler,
+            )
 
             // Act
             onNodeWithTag(CONNECT_BUTTON_TEST_TAG).performClick()
@@ -438,22 +437,20 @@ class ConnectScreenTest {
         composeExtension.use {
             // Arrange
             val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                ConnectScreen(
-                    state =
-                        ConnectUiState(
-                            location = null,
-                            selectedRelayItemTitle = null,
-                            tunnelState = TunnelState.Connecting(null, null, emptyList()),
-                            showLocation = false,
-                            deviceName = "",
-                            daysLeftUntilExpiry = null,
-                            inAppNotification = null,
-                            isPlayBuild = false,
-                        ),
-                    onCancelClick = mockedClickHandler,
-                )
-            }
+            initScreen(
+                state =
+                    ConnectUiState(
+                        location = null,
+                        selectedRelayItemTitle = null,
+                        tunnelState = TunnelState.Connecting(null, null, emptyList()),
+                        showLocation = false,
+                        deviceName = "",
+                        daysLeftUntilExpiry = null,
+                        inAppNotification = null,
+                        isPlayBuild = false,
+                    ),
+                onCancelClick = mockedClickHandler,
+            )
 
             // Act
             onNodeWithTag(CONNECT_BUTTON_TEST_TAG).performClick()
@@ -490,26 +487,20 @@ class ConnectScreenTest {
             val outIpv6 = "ipv6address"
             every { mockLocation.ipv6?.hostAddress } returns outIpv6
 
-            setContentWithTheme {
-                ConnectScreen(
-                    state =
-                        ConnectUiState(
-                            location = mockLocation,
-                            selectedRelayItemTitle = null,
-                            tunnelState =
-                                TunnelState.Connected(
-                                    mockTunnelEndpoint,
-                                    mockLocation,
-                                    emptyList(),
-                                ),
-                            showLocation = false,
-                            deviceName = "",
-                            daysLeftUntilExpiry = null,
-                            inAppNotification = null,
-                            isPlayBuild = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    ConnectUiState(
+                        location = mockLocation,
+                        selectedRelayItemTitle = null,
+                        tunnelState =
+                            TunnelState.Connected(mockTunnelEndpoint, mockLocation, emptyList()),
+                        showLocation = false,
+                        deviceName = "",
+                        daysLeftUntilExpiry = null,
+                        inAppNotification = null,
+                        isPlayBuild = false,
+                    )
+            )
 
             // Act
             onNodeWithTag(CONNECT_CARD_HEADER_TEST_TAG).performClick()
@@ -532,21 +523,19 @@ class ConnectScreenTest {
         composeExtension.use {
             // Arrange
             val versionInfo = VersionInfo(currentVersion = "1.0", isSupported = false)
-            setContentWithTheme {
-                ConnectScreen(
-                    state =
-                        ConnectUiState(
-                            location = null,
-                            selectedRelayItemTitle = null,
-                            tunnelState = TunnelState.Connecting(null, null, emptyList()),
-                            showLocation = false,
-                            deviceName = "",
-                            daysLeftUntilExpiry = null,
-                            inAppNotification = InAppNotification.UnsupportedVersion(versionInfo),
-                            isPlayBuild = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    ConnectUiState(
+                        location = null,
+                        selectedRelayItemTitle = null,
+                        tunnelState = TunnelState.Connecting(null, null, emptyList()),
+                        showLocation = false,
+                        deviceName = "",
+                        daysLeftUntilExpiry = null,
+                        inAppNotification = InAppNotification.UnsupportedVersion(versionInfo),
+                        isPlayBuild = false,
+                    )
+            )
 
             // Assert
             onNodeWithText("UNSUPPORTED VERSION").assertExists()
@@ -562,24 +551,20 @@ class ConnectScreenTest {
         composeExtension.use {
             // Arrange
             val expiryDate = DateTime(2020, 11, 11, 10, 10)
-            setContentWithTheme {
-                ConnectScreen(
-                    state =
-                        ConnectUiState(
-                            location = null,
-                            selectedRelayItemTitle = null,
-                            tunnelState = TunnelState.Connecting(null, null, emptyList()),
-                            showLocation = false,
-                            deviceName = "",
-                            daysLeftUntilExpiry = null,
-                            inAppNotification =
-                                InAppNotification.AccountExpiry(
-                                    Duration(DateTime.now(), expiryDate)
-                                ),
-                            isPlayBuild = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    ConnectUiState(
+                        location = null,
+                        selectedRelayItemTitle = null,
+                        tunnelState = TunnelState.Connecting(null, null, emptyList()),
+                        showLocation = false,
+                        deviceName = "",
+                        daysLeftUntilExpiry = null,
+                        inAppNotification =
+                            InAppNotification.AccountExpiry(Duration(DateTime.now(), expiryDate)),
+                        isPlayBuild = false,
+                    )
+            )
 
             // Assert
             onNodeWithText("ACCOUNT CREDIT EXPIRES SOON").assertExists()
@@ -593,22 +578,20 @@ class ConnectScreenTest {
             // Arrange
             val mockedClickHandler: () -> Unit = mockk(relaxed = true)
             val versionInfo = VersionInfo(isSupported = false, currentVersion = "")
-            setContentWithTheme {
-                ConnectScreen(
-                    onOpenAppListing = mockedClickHandler,
-                    state =
-                        ConnectUiState(
-                            location = null,
-                            selectedRelayItemTitle = null,
-                            tunnelState = TunnelState.Connecting(null, null, emptyList()),
-                            showLocation = false,
-                            deviceName = "",
-                            daysLeftUntilExpiry = null,
-                            inAppNotification = InAppNotification.UnsupportedVersion(versionInfo),
-                            isPlayBuild = false,
-                        ),
-                )
-            }
+            initScreen(
+                onOpenAppListing = mockedClickHandler,
+                state =
+                    ConnectUiState(
+                        location = null,
+                        selectedRelayItemTitle = null,
+                        tunnelState = TunnelState.Connecting(null, null, emptyList()),
+                        showLocation = false,
+                        deviceName = "",
+                        daysLeftUntilExpiry = null,
+                        inAppNotification = InAppNotification.UnsupportedVersion(versionInfo),
+                        isPlayBuild = false,
+                    ),
+            )
 
             // Act
             onNodeWithTag(NOTIFICATION_BANNER_ACTION).performClick()
@@ -624,25 +607,21 @@ class ConnectScreenTest {
             // Arrange
             val mockedClickHandler: () -> Unit = mockk(relaxed = true)
             val expiryDate = DateTime(2020, 11, 11, 10, 10)
-            setContentWithTheme {
-                ConnectScreen(
-                    onManageAccountClick = mockedClickHandler,
-                    state =
-                        ConnectUiState(
-                            location = null,
-                            selectedRelayItemTitle = null,
-                            tunnelState = TunnelState.Connecting(null, null, emptyList()),
-                            showLocation = false,
-                            deviceName = "",
-                            daysLeftUntilExpiry = null,
-                            inAppNotification =
-                                InAppNotification.AccountExpiry(
-                                    Duration(DateTime.now(), expiryDate)
-                                ),
-                            isPlayBuild = false,
-                        ),
-                )
-            }
+            initScreen(
+                onManageAccountClick = mockedClickHandler,
+                state =
+                    ConnectUiState(
+                        location = null,
+                        selectedRelayItemTitle = null,
+                        tunnelState = TunnelState.Connecting(null, null, emptyList()),
+                        showLocation = false,
+                        deviceName = "",
+                        daysLeftUntilExpiry = null,
+                        inAppNotification =
+                            InAppNotification.AccountExpiry(Duration(DateTime.now(), expiryDate)),
+                        isPlayBuild = false,
+                    ),
+            )
 
             // Act
             onNodeWithTag(NOTIFICATION_BANNER_ACTION).performClick()
@@ -657,9 +636,7 @@ class ConnectScreenTest {
         composeExtension.use {
             // Arrange
             val onAccountClickMockk: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                ConnectScreen(state = ConnectUiState.INITIAL, onAccountClick = onAccountClickMockk)
-            }
+            initScreen(state = ConnectUiState.INITIAL, onAccountClick = onAccountClickMockk)
 
             // Assert
             onNodeWithTag(TOP_BAR_ACCOUNT_BUTTON).performClick()
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/CustomListLocationsScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/CustomListLocationsScreenTest.kt
index 484bb132d6..dee081ba03 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/CustomListLocationsScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/CustomListLocationsScreenTest.kt
@@ -5,6 +5,7 @@ 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.ComposeContext
 import io.mockk.MockKAnnotations
 import io.mockk.mockk
 import io.mockk.verify
@@ -29,15 +30,32 @@ class CustomListLocationsScreenTest {
         MockKAnnotations.init(this)
     }
 
+    private fun ComposeContext.initScreen(
+        state: CustomListLocationsUiState,
+        onSearchTermInput: (String) -> Unit = {},
+        onSaveClick: () -> Unit = {},
+        onRelaySelectionClick: (RelayItem.Location, selected: Boolean) -> Unit = { _, _ -> },
+        onExpand: (RelayItem.Location, selected: Boolean) -> Unit = { _, _ -> },
+        onBackClick: () -> Unit = {},
+    ) {
+
+        setContentWithTheme {
+            CustomListLocationsScreen(
+                state = state,
+                onSearchTermInput = onSearchTermInput,
+                onSaveClick = onSaveClick,
+                onRelaySelectionClick = onRelaySelectionClick,
+                onExpand = onExpand,
+                onBackClick = onBackClick,
+            )
+        }
+    }
+
     @Test
     fun givenLoadingStateShouldShowLoadingSpinner() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                CustomListLocationsScreen(
-                    state = CustomListLocationsUiState.Loading(newList = false)
-                )
-            }
+            initScreen(state = CustomListLocationsUiState.Loading(newList = false))
 
             // Assert
             onNodeWithTag(CIRCULAR_PROGRESS_INDICATOR).assertExists()
@@ -48,11 +66,7 @@ class CustomListLocationsScreenTest {
         composeExtension.use {
             // Arrange
             val newList = true
-            setContentWithTheme {
-                CustomListLocationsScreen(
-                    state = CustomListLocationsUiState.Loading(newList = newList)
-                )
-            }
+            initScreen(state = CustomListLocationsUiState.Loading(newList = newList))
 
             // Assert
             onNodeWithText(ADD_LOCATIONS_TEXT).assertExists()
@@ -63,11 +77,7 @@ class CustomListLocationsScreenTest {
         composeExtension.use {
             // Arrange
             val newList = false
-            setContentWithTheme {
-                CustomListLocationsScreen(
-                    state = CustomListLocationsUiState.Loading(newList = newList)
-                )
-            }
+            initScreen(state = CustomListLocationsUiState.Loading(newList = newList))
 
             // Assert
             onNodeWithText(EDIT_LOCATIONS_TEXT).assertExists()
@@ -77,19 +87,17 @@ class CustomListLocationsScreenTest {
     fun givenListOfAvailableLocationsShouldShowThem() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                CustomListLocationsScreen(
-                    state =
-                        CustomListLocationsUiState.Content.Data(
-                            locations =
-                                listOf(
-                                    RelayLocationListItem(DUMMY_RELAY_COUNTRIES[0], checked = true),
-                                    RelayLocationListItem(DUMMY_RELAY_COUNTRIES[1], checked = false),
-                                ),
-                            searchTerm = "",
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    CustomListLocationsUiState.Content.Data(
+                        locations =
+                            listOf(
+                                RelayLocationListItem(DUMMY_RELAY_COUNTRIES[0], checked = true),
+                                RelayLocationListItem(DUMMY_RELAY_COUNTRIES[1], checked = false),
+                            ),
+                        searchTerm = "",
+                    )
+            )
 
             // Assert
             onNodeWithText("Relay Country 1").assertExists()
@@ -102,17 +110,14 @@ class CustomListLocationsScreenTest {
             // Arrange
             val selectedCountry = DUMMY_RELAY_COUNTRIES[0]
             val mockedOnRelaySelectionClicked: (RelayItem, Boolean) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                CustomListLocationsScreen(
-                    state =
-                        CustomListLocationsUiState.Content.Data(
-                            newList = false,
-                            locations =
-                                listOf(RelayLocationListItem(selectedCountry, checked = true)),
-                        ),
-                    onRelaySelectionClick = mockedOnRelaySelectionClicked,
-                )
-            }
+            initScreen(
+                state =
+                    CustomListLocationsUiState.Content.Data(
+                        newList = false,
+                        locations = listOf(RelayLocationListItem(selectedCountry, checked = true)),
+                    ),
+                onRelaySelectionClick = mockedOnRelaySelectionClicked,
+            )
 
             // Act
             onNodeWithText(selectedCountry.name).performClick()
@@ -126,16 +131,14 @@ class CustomListLocationsScreenTest {
         composeExtension.use {
             // Arrange
             val mockedSearchTermInput: (String) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                CustomListLocationsScreen(
-                    state =
-                        CustomListLocationsUiState.Content.Data(
-                            newList = false,
-                            locations = emptyList(),
-                        ),
-                    onSearchTermInput = mockedSearchTermInput,
-                )
-            }
+            initScreen(
+                state =
+                    CustomListLocationsUiState.Content.Data(
+                        newList = false,
+                        locations = emptyList(),
+                    ),
+                onSearchTermInput = mockedSearchTermInput,
+            )
             val mockSearchString = "SEARCH"
 
             // Act
@@ -151,16 +154,14 @@ class CustomListLocationsScreenTest {
             // Arrange
             val mockedSearchTermInput: (String) -> Unit = mockk(relaxed = true)
             val mockSearchString = "SEARCH"
-            setContentWithTheme {
-                CustomListLocationsScreen(
-                    state =
-                        CustomListLocationsUiState.Content.Empty(
-                            newList = false,
-                            searchTerm = mockSearchString,
-                        ),
-                    onSearchTermInput = mockedSearchTermInput,
-                )
-            }
+            initScreen(
+                state =
+                    CustomListLocationsUiState.Content.Empty(
+                        newList = false,
+                        searchTerm = mockSearchString,
+                    ),
+                onSearchTermInput = mockedSearchTermInput,
+            )
 
             // Assert
             onNodeWithText(EMPTY_SEARCH.format(mockSearchString)).assertExists()
@@ -171,15 +172,13 @@ class CustomListLocationsScreenTest {
         composeExtension.use {
             // Arrange
             val emptySearchString = ""
-            setContentWithTheme {
-                CustomListLocationsScreen(
-                    state =
-                        CustomListLocationsUiState.Content.Empty(
-                            newList = false,
-                            searchTerm = emptySearchString,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    CustomListLocationsUiState.Content.Empty(
+                        newList = false,
+                        searchTerm = emptySearchString,
+                    )
+            )
 
             // Assert
             onNodeWithText(NO_LOCATIONS_FOUND_TEXT).assertExists()
@@ -190,17 +189,15 @@ class CustomListLocationsScreenTest {
         composeExtension.use {
             // Arrange
             val mockOnSaveClick: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                CustomListLocationsScreen(
-                    state =
-                        CustomListLocationsUiState.Content.Data(
-                            newList = false,
-                            locations = emptyList(),
-                            saveEnabled = true,
-                        ),
-                    onSaveClick = mockOnSaveClick,
-                )
-            }
+            initScreen(
+                state =
+                    CustomListLocationsUiState.Content.Data(
+                        newList = false,
+                        locations = emptyList(),
+                        saveEnabled = true,
+                    ),
+                onSaveClick = mockOnSaveClick,
+            )
 
             // Act
             onNodeWithTag(SAVE_BUTTON_TEST_TAG).performClick()
@@ -214,17 +211,15 @@ class CustomListLocationsScreenTest {
         composeExtension.use {
             // Arrange
             val mockOnSaveClick: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                CustomListLocationsScreen(
-                    state =
-                        CustomListLocationsUiState.Content.Data(
-                            newList = false,
-                            locations = emptyList(),
-                            saveEnabled = false,
-                        ),
-                    onSaveClick = mockOnSaveClick,
-                )
-            }
+            initScreen(
+                state =
+                    CustomListLocationsUiState.Content.Data(
+                        newList = false,
+                        locations = emptyList(),
+                        saveEnabled = false,
+                    ),
+                onSaveClick = mockOnSaveClick,
+            )
 
             // Act
             onNodeWithTag(SAVE_BUTTON_TEST_TAG).performClick()
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/CustomListsScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/CustomListsScreenTest.kt
index 47e88bb52f..f319ce92b3 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/CustomListsScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/CustomListsScreenTest.kt
@@ -1,10 +1,10 @@
 package net.mullvad.mullvadvpn.compose.screen
 
-import androidx.compose.material3.SnackbarHostState
 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.ComposeContext
 import io.mockk.MockKAnnotations
 import io.mockk.mockk
 import io.mockk.verify
@@ -28,16 +28,28 @@ class CustomListsScreenTest {
         MockKAnnotations.init(this)
     }
 
+    private fun ComposeContext.initScreen(
+        state: CustomListsUiState = CustomListsUiState.Loading,
+        addCustomList: () -> Unit = {},
+        openCustomList: (CustomList) -> Unit = {},
+        onBackClick: () -> Unit = {},
+    ) {
+
+        setContentWithTheme {
+            CustomListsScreen(
+                state = state,
+                addCustomList = addCustomList,
+                openCustomList = openCustomList,
+                onBackClick = onBackClick,
+            )
+        }
+    }
+
     @Test
     fun givenLoadingStateShouldShowLoadingSpinner() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                CustomListsScreen(
-                    state = CustomListsUiState.Loading,
-                    snackbarHostState = SnackbarHostState(),
-                )
-            }
+            initScreen(state = CustomListsUiState.Loading)
 
             // Assert
             onNodeWithTag(CIRCULAR_PROGRESS_INDICATOR).assertExists()
@@ -48,12 +60,7 @@ class CustomListsScreenTest {
         composeExtension.use {
             // Arrange
             val customLists = DUMMY_CUSTOM_LISTS
-            setContentWithTheme {
-                CustomListsScreen(
-                    state = CustomListsUiState.Content(customLists = customLists),
-                    snackbarHostState = SnackbarHostState(),
-                )
-            }
+            initScreen(state = CustomListsUiState.Content(customLists = customLists))
 
             // Assert
             onNodeWithText(customLists[0].name.value).assertExists()
@@ -66,13 +73,10 @@ class CustomListsScreenTest {
             // Arrange
             val customLists = DUMMY_CUSTOM_LISTS
             val mockedAddCustomList: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                CustomListsScreen(
-                    state = CustomListsUiState.Content(customLists = customLists),
-                    snackbarHostState = SnackbarHostState(),
-                    addCustomList = mockedAddCustomList,
-                )
-            }
+            initScreen(
+                state = CustomListsUiState.Content(customLists = customLists),
+                addCustomList = mockedAddCustomList,
+            )
 
             // Act
             onNodeWithTag(NEW_LIST_BUTTON_TEST_TAG).performClick()
@@ -88,13 +92,10 @@ class CustomListsScreenTest {
             val customLists = DUMMY_CUSTOM_LISTS
             val clickedList = DUMMY_CUSTOM_LISTS[0]
             val mockedOpenCustomList: (CustomList) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                CustomListsScreen(
-                    state = CustomListsUiState.Content(customLists = customLists),
-                    snackbarHostState = SnackbarHostState(),
-                    openCustomList = mockedOpenCustomList,
-                )
-            }
+            initScreen(
+                state = CustomListsUiState.Content(customLists = customLists),
+                openCustomList = mockedOpenCustomList,
+            )
 
             // Act
             onNodeWithText(clickedList.name.value).performClick()
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 cd1c03a59e..879d76ddd4 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
@@ -3,6 +3,7 @@ package net.mullvad.mullvadvpn.compose.screen
 import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
+import de.mannodermaus.junit5.compose.ComposeContext
 import io.mockk.MockKAnnotations
 import io.mockk.mockk
 import io.mockk.verify
@@ -22,6 +23,20 @@ class DeviceRevokedScreenTest {
         MockKAnnotations.init(this)
     }
 
+    private fun ComposeContext.initScreen(
+        state: DeviceRevokedUiState,
+        onSettingsClicked: () -> Unit = {},
+        onGoToLoginClicked: () -> Unit = {},
+    ) {
+        setContentWithTheme {
+            DeviceRevokedScreen(
+                state = state,
+                onSettingsClicked = onSettingsClicked,
+                onGoToLoginClicked = onGoToLoginClicked,
+            )
+        }
+    }
+
     @Test
     fun testUnblockWarningShowingWhenSecured() =
         composeExtension.use {
@@ -29,7 +44,7 @@ class DeviceRevokedScreenTest {
             val state = DeviceRevokedUiState.SECURED
 
             // Act
-            setContentWithTheme { DeviceRevokedScreen(state) }
+            initScreen(state)
 
             // Assert
             onNodeWithText(UNBLOCK_WARNING).assertExists()
@@ -42,7 +57,7 @@ class DeviceRevokedScreenTest {
             val state = DeviceRevokedUiState.UNSECURED
 
             // Act
-            setContentWithTheme { DeviceRevokedScreen(state) }
+            initScreen(state)
 
             // Assert
             onNodeWithText(UNBLOCK_WARNING).assertDoesNotExist()
@@ -54,9 +69,7 @@ class DeviceRevokedScreenTest {
             // Arrange
             val state = DeviceRevokedUiState.UNSECURED
             val mockOnGoToLoginClicked: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                DeviceRevokedScreen(state = state, onGoToLoginClicked = mockOnGoToLoginClicked)
-            }
+            initScreen(state = state, onGoToLoginClicked = mockOnGoToLoginClicked)
 
             // Act
             onNodeWithText(GO_TO_LOGIN_BUTTON_TEXT).performClick()
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/EditApiAccessMethodScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/EditApiAccessMethodScreenTest.kt
index d737f5459f..1c763219ea 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/EditApiAccessMethodScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/EditApiAccessMethodScreenTest.kt
@@ -5,6 +5,7 @@ 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.ComposeContext
 import io.mockk.mockk
 import io.mockk.verify
 import net.mullvad.mullvadvpn.compose.createEdgeToEdgeComposeExtension
@@ -13,6 +14,7 @@ import net.mullvad.mullvadvpn.compose.state.ApiAccessMethodTypes
 import net.mullvad.mullvadvpn.compose.state.EditApiAccessFormData
 import net.mullvad.mullvadvpn.compose.state.EditApiAccessMethodUiState
 import net.mullvad.mullvadvpn.compose.test.EDIT_API_ACCESS_NAME_INPUT
+import net.mullvad.mullvadvpn.lib.model.Cipher
 import net.mullvad.mullvadvpn.lib.model.InvalidDataError
 import net.mullvad.mullvadvpn.lib.model.ParsePortError
 import org.junit.jupiter.api.Test
@@ -22,21 +24,51 @@ import org.junit.jupiter.api.extension.RegisterExtension
 class EditApiAccessMethodScreenTest {
     @JvmField @RegisterExtension val composeExtension = createEdgeToEdgeComposeExtension()
 
+    private fun ComposeContext.initScreen(
+        state: EditApiAccessMethodUiState,
+        onNameChanged: (String) -> Unit = {},
+        onTypeSelected: (ApiAccessMethodTypes) -> Unit = {},
+        onIpChanged: (String) -> Unit = {},
+        onPortChanged: (String) -> Unit = {},
+        onPasswordChanged: (String) -> Unit = {},
+        onCipherChange: (Cipher) -> Unit = {},
+        onToggleAuthenticationEnabled: (Boolean) -> Unit = {},
+        onUsernameChanged: (String) -> Unit = {},
+        onTestMethod: () -> Unit = {},
+        onAddMethod: () -> Unit = {},
+        onNavigateBack: () -> Unit = {},
+    ) {
+        setContentWithTheme {
+            EditApiAccessMethodScreen(
+                state = state,
+                onNameChanged = onNameChanged,
+                onTypeSelected = onTypeSelected,
+                onIpChanged = onIpChanged,
+                onPortChanged = onPortChanged,
+                onPasswordChanged = onPasswordChanged,
+                onCipherChange = onCipherChange,
+                onToggleAuthenticationEnabled = onToggleAuthenticationEnabled,
+                onUsernameChanged = onUsernameChanged,
+                onTestMethod = onTestMethod,
+                onAddMethod = onAddMethod,
+                onNavigateBack = onNavigateBack,
+            )
+        }
+    }
+
     @Test
     fun whenInEditModeAddButtonShouldSaySave() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                EditApiAccessMethodScreen(
-                    state =
-                        EditApiAccessMethodUiState.Content(
-                            editMode = true,
-                            formData = EditApiAccessFormData.empty(),
-                            hasChanges = false,
-                            isTestingApiAccessMethod = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    EditApiAccessMethodUiState.Content(
+                        editMode = true,
+                        formData = EditApiAccessFormData.empty(),
+                        hasChanges = false,
+                        isTestingApiAccessMethod = false,
+                    )
+            )
 
             // Assert
             onNodeWithText("Save").assertExists()
@@ -46,17 +78,15 @@ class EditApiAccessMethodScreenTest {
     fun whenNotInEditModeAddButtonShouldSayAdd() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                EditApiAccessMethodScreen(
-                    state =
-                        EditApiAccessMethodUiState.Content(
-                            editMode = false,
-                            formData = EditApiAccessFormData.empty(),
-                            hasChanges = false,
-                            isTestingApiAccessMethod = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    EditApiAccessMethodUiState.Content(
+                        editMode = false,
+                        formData = EditApiAccessFormData.empty(),
+                        hasChanges = false,
+                        isTestingApiAccessMethod = false,
+                    )
+            )
 
             // Assert
             onNodeWithText("Add").assertExists()
@@ -66,25 +96,23 @@ class EditApiAccessMethodScreenTest {
     fun whenNameInputHasErrorShouldShowError() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                EditApiAccessMethodScreen(
-                    state =
-                        EditApiAccessMethodUiState.Content(
-                            editMode = false,
-                            formData =
-                                EditApiAccessFormData(
-                                    name = "",
-                                    nameError = InvalidDataError.NameError.Required,
-                                    serverIp = "",
-                                    username = "",
-                                    password = "",
-                                    port = "",
-                                ),
-                            hasChanges = false,
-                            isTestingApiAccessMethod = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    EditApiAccessMethodUiState.Content(
+                        editMode = false,
+                        formData =
+                            EditApiAccessFormData(
+                                name = "",
+                                nameError = InvalidDataError.NameError.Required,
+                                serverIp = "",
+                                username = "",
+                                password = "",
+                                port = "",
+                            ),
+                        hasChanges = false,
+                        isTestingApiAccessMethod = false,
+                    )
+            )
 
             // Assert
             onNodeWithText("This field is required").assertExists()
@@ -94,25 +122,23 @@ class EditApiAccessMethodScreenTest {
     fun whenServerInputIsNotIpAddressShouldShowError() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                EditApiAccessMethodScreen(
-                    state =
-                        EditApiAccessMethodUiState.Content(
-                            editMode = false,
-                            formData =
-                                EditApiAccessFormData(
-                                    name = "",
-                                    serverIp = "123",
-                                    serverIpError = InvalidDataError.ServerIpError.Invalid,
-                                    username = "",
-                                    password = "",
-                                    port = "",
-                                ),
-                            hasChanges = false,
-                            isTestingApiAccessMethod = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    EditApiAccessMethodUiState.Content(
+                        editMode = false,
+                        formData =
+                            EditApiAccessFormData(
+                                name = "",
+                                serverIp = "123",
+                                serverIpError = InvalidDataError.ServerIpError.Invalid,
+                                username = "",
+                                password = "",
+                                port = "",
+                            ),
+                        hasChanges = false,
+                        isTestingApiAccessMethod = false,
+                    )
+            )
 
             // Assert
             onNodeWithText("Please enter a valid IPv4 or IPv6 address").assertExists()
@@ -122,28 +148,26 @@ class EditApiAccessMethodScreenTest {
     fun whenPortInputIsNotWithinRangeShouldShowError() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                EditApiAccessMethodScreen(
-                    state =
-                        EditApiAccessMethodUiState.Content(
-                            editMode = false,
-                            formData =
-                                EditApiAccessFormData(
-                                    name = "",
-                                    serverIp = "",
-                                    username = "",
-                                    password = "",
-                                    port = "1111111111",
-                                    portError =
-                                        InvalidDataError.PortError.Invalid(
-                                            ParsePortError.OutOfRange(1111111111)
-                                        ),
-                                ),
-                            hasChanges = false,
-                            isTestingApiAccessMethod = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    EditApiAccessMethodUiState.Content(
+                        editMode = false,
+                        formData =
+                            EditApiAccessFormData(
+                                name = "",
+                                serverIp = "",
+                                username = "",
+                                password = "",
+                                port = "1111111111",
+                                portError =
+                                    InvalidDataError.PortError.Invalid(
+                                        ParsePortError.OutOfRange(1111111111)
+                                    ),
+                            ),
+                        hasChanges = false,
+                        isTestingApiAccessMethod = false,
+                    )
+            )
 
             // Assert
             onNodeWithText("Please enter a valid remote server port").assertExists()
@@ -155,18 +179,16 @@ class EditApiAccessMethodScreenTest {
             // Arrange
             val onNameChanged: (String) -> Unit = mockk(relaxed = true)
             val mockInput = "Name"
-            setContentWithTheme {
-                EditApiAccessMethodScreen(
-                    state =
-                        EditApiAccessMethodUiState.Content(
-                            editMode = false,
-                            formData = EditApiAccessFormData.empty(),
-                            hasChanges = false,
-                            isTestingApiAccessMethod = false,
-                        ),
-                    onNameChanged = onNameChanged,
-                )
-            }
+            initScreen(
+                state =
+                    EditApiAccessMethodUiState.Content(
+                        editMode = false,
+                        formData = EditApiAccessFormData.empty(),
+                        hasChanges = false,
+                        isTestingApiAccessMethod = false,
+                    ),
+                onNameChanged = onNameChanged,
+            )
 
             // Act
             onNodeWithTag(EDIT_API_ACCESS_NAME_INPUT).performTextInput(mockInput)
@@ -179,26 +201,24 @@ class EditApiAccessMethodScreenTest {
     fun whenSocks5IsSelectedAndAuthenticationIsEnabledShouldShowUsernameAndPassword() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                EditApiAccessMethodScreen(
-                    state =
-                        EditApiAccessMethodUiState.Content(
-                            editMode = false,
-                            formData =
-                                EditApiAccessFormData(
-                                    name = "",
-                                    serverIp = "",
-                                    username = "",
-                                    password = "",
-                                    port = "",
-                                    enableAuthentication = true,
-                                    apiAccessMethodTypes = ApiAccessMethodTypes.SOCKS5_REMOTE,
-                                ),
-                            hasChanges = false,
-                            isTestingApiAccessMethod = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    EditApiAccessMethodUiState.Content(
+                        editMode = false,
+                        formData =
+                            EditApiAccessFormData(
+                                name = "",
+                                serverIp = "",
+                                username = "",
+                                password = "",
+                                port = "",
+                                enableAuthentication = true,
+                                apiAccessMethodTypes = ApiAccessMethodTypes.SOCKS5_REMOTE,
+                            ),
+                        hasChanges = false,
+                        isTestingApiAccessMethod = false,
+                    )
+            )
 
             // Assert
             onNodeWithText("Username").assertExists()
@@ -210,18 +230,16 @@ class EditApiAccessMethodScreenTest {
         composeExtension.use {
             // Arrange
             val onTestMethod: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                EditApiAccessMethodScreen(
-                    state =
-                        EditApiAccessMethodUiState.Content(
-                            editMode = false,
-                            formData = EditApiAccessFormData.empty(),
-                            hasChanges = false,
-                            isTestingApiAccessMethod = false,
-                        ),
-                    onTestMethod = onTestMethod,
-                )
-            }
+            initScreen(
+                state =
+                    EditApiAccessMethodUiState.Content(
+                        editMode = false,
+                        formData = EditApiAccessFormData.empty(),
+                        hasChanges = false,
+                        isTestingApiAccessMethod = false,
+                    ),
+                onTestMethod = onTestMethod,
+            )
 
             // Act
             onNodeWithText("Test method").performClick()
@@ -235,18 +253,16 @@ class EditApiAccessMethodScreenTest {
         composeExtension.use {
             // Arrange
             val onAddMethod: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                EditApiAccessMethodScreen(
-                    state =
-                        EditApiAccessMethodUiState.Content(
-                            editMode = false,
-                            formData = EditApiAccessFormData.empty(),
-                            hasChanges = false,
-                            isTestingApiAccessMethod = false,
-                        ),
-                    onAddMethod = onAddMethod,
-                )
-            }
+            initScreen(
+                state =
+                    EditApiAccessMethodUiState.Content(
+                        editMode = false,
+                        formData = EditApiAccessFormData.empty(),
+                        hasChanges = false,
+                        isTestingApiAccessMethod = false,
+                    ),
+                onAddMethod = onAddMethod,
+            )
 
             // Act
             onNodeWithText("Add").performClick()
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/EditCustomListScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/EditCustomListScreenTest.kt
index 1207f2417b..f3a4e3b2af 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/EditCustomListScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/EditCustomListScreenTest.kt
@@ -4,6 +4,7 @@ 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.ComposeContext
 import io.mockk.MockKAnnotations
 import io.mockk.mockk
 import io.mockk.verify
@@ -29,11 +30,29 @@ class EditCustomListScreenTest {
         MockKAnnotations.init(this)
     }
 
+    private fun ComposeContext.initScreen(
+        state: EditCustomListUiState = EditCustomListUiState.Loading,
+        onDeleteList: (id: CustomListId, name: CustomListName) -> Unit = { _, _ -> },
+        onNameClicked: (id: CustomListId, name: CustomListName) -> Unit = { _, _ -> },
+        onLocationsClicked: (CustomListId) -> Unit = {},
+        onBackClick: () -> Unit = {},
+    ) {
+        setContentWithTheme {
+            EditCustomListScreen(
+                state = state,
+                onDeleteList = onDeleteList,
+                onNameClicked = onNameClicked,
+                onLocationsClicked = onLocationsClicked,
+                onBackClick = onBackClick,
+            )
+        }
+    }
+
     @Test
     fun givenLoadingStateShouldShowLoadingSpinner() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme { EditCustomListScreen(state = EditCustomListUiState.Loading) }
+            initScreen()
 
             // Assert
             onNodeWithTag(CIRCULAR_PROGRESS_INDICATOR).assertExists()
@@ -43,7 +62,7 @@ class EditCustomListScreenTest {
     fun givenNotFoundStateShouldShowNotFound() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme { EditCustomListScreen(state = EditCustomListUiState.NotFound) }
+            initScreen(state = EditCustomListUiState.NotFound)
 
             // Assert
             onNodeWithText(NOT_FOUND_TEXT).assertExists()
@@ -54,16 +73,15 @@ class EditCustomListScreenTest {
         composeExtension.use {
             // Arrange
             val customList = DUMMY_CUSTOM_LISTS[0]
-            setContentWithTheme {
-                EditCustomListScreen(
-                    state =
-                        EditCustomListUiState.Content(
-                            id = customList.id,
-                            name = customList.name,
-                            locations = customList.locations,
-                        )
-                )
-            }
+
+            initScreen(
+                state =
+                    EditCustomListUiState.Content(
+                        id = customList.id,
+                        name = customList.name,
+                        locations = customList.locations,
+                    )
+            )
 
             // Assert
             onNodeWithText(customList.name.value)
@@ -74,16 +92,14 @@ class EditCustomListScreenTest {
         composeExtension.use {
             // Arrange
             val customList = DUMMY_CUSTOM_LISTS[0]
-            setContentWithTheme {
-                EditCustomListScreen(
-                    state =
-                        EditCustomListUiState.Content(
-                            id = customList.id,
-                            name = customList.name,
-                            locations = customList.locations,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    EditCustomListUiState.Content(
+                        id = customList.id,
+                        name = customList.name,
+                        locations = customList.locations,
+                    )
+            )
 
             // Assert
             onNodeWithText(LOCATIONS_TEXT.format(customList.locations.size))
@@ -95,17 +111,15 @@ class EditCustomListScreenTest {
             // Arrange
             val mockedOnDelete: (CustomListId, CustomListName) -> Unit = mockk(relaxed = true)
             val customList = DUMMY_CUSTOM_LISTS[0]
-            setContentWithTheme {
-                EditCustomListScreen(
-                    state =
-                        EditCustomListUiState.Content(
-                            id = customList.id,
-                            name = customList.name,
-                            locations = customList.locations,
-                        ),
-                    onDeleteList = mockedOnDelete,
-                )
-            }
+            initScreen(
+                state =
+                    EditCustomListUiState.Content(
+                        id = customList.id,
+                        name = customList.name,
+                        locations = customList.locations,
+                    ),
+                onDeleteList = mockedOnDelete,
+            )
 
             // Act
             onNodeWithTag(TOP_BAR_DROPDOWN_BUTTON_TEST_TAG).performClick()
@@ -121,17 +135,15 @@ class EditCustomListScreenTest {
             // Arrange
             val mockedOnNameClicked: (CustomListId, CustomListName) -> Unit = mockk(relaxed = true)
             val customList = DUMMY_CUSTOM_LISTS[0]
-            setContentWithTheme {
-                EditCustomListScreen(
-                    state =
-                        EditCustomListUiState.Content(
-                            id = customList.id,
-                            name = customList.name,
-                            locations = customList.locations,
-                        ),
-                    onNameClicked = mockedOnNameClicked,
-                )
-            }
+            initScreen(
+                state =
+                    EditCustomListUiState.Content(
+                        id = customList.id,
+                        name = customList.name,
+                        locations = customList.locations,
+                    ),
+                onNameClicked = mockedOnNameClicked,
+            )
 
             // Act
             onNodeWithText(customList.name.value).performClick()
@@ -146,17 +158,15 @@ class EditCustomListScreenTest {
             // Arrange
             val mockedOnLocationsClicked: (CustomListId) -> Unit = mockk(relaxed = true)
             val customList = DUMMY_CUSTOM_LISTS[0]
-            setContentWithTheme {
-                EditCustomListScreen(
-                    state =
-                        EditCustomListUiState.Content(
-                            id = customList.id,
-                            name = customList.name,
-                            locations = customList.locations,
-                        ),
-                    onLocationsClicked = mockedOnLocationsClicked,
-                )
-            }
+            initScreen(
+                state =
+                    EditCustomListUiState.Content(
+                        id = customList.id,
+                        name = customList.name,
+                        locations = customList.locations,
+                    ),
+                onLocationsClicked = mockedOnLocationsClicked,
+            )
 
             // Act
             onNodeWithText(LOCATIONS_TEXT.format(customList.locations.size)).performClick()
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 7ca3d6941e..2f16b27c23 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
@@ -3,6 +3,7 @@ package net.mullvad.mullvadvpn.compose.screen
 import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
+import de.mannodermaus.junit5.compose.ComposeContext
 import io.mockk.MockKAnnotations
 import io.mockk.mockk
 import io.mockk.verify
@@ -23,20 +24,37 @@ class FilterScreenTest {
         MockKAnnotations.init(this)
     }
 
+    private fun ComposeContext.initScreen(
+        state: RelayFilterUiState = RelayFilterUiState(),
+        onBackClick: () -> Unit = {},
+        onApplyClick: () -> Unit = {},
+        onSelectedOwnership: (ownership: Ownership?) -> Unit = {},
+        onAllProviderCheckChange: (isChecked: Boolean) -> Unit = {},
+        onSelectedProvider: (checked: Boolean, provider: Provider) -> Unit = { _, _ -> },
+    ) {
+        setContentWithTheme {
+            FilterScreen(
+                state = state,
+                onBackClick = onBackClick,
+                onApplyClick = onApplyClick,
+                onSelectedOwnership = onSelectedOwnership,
+                onAllProviderCheckChange = onAllProviderCheckChange,
+                onSelectedProvider = onSelectedProvider,
+            )
+        }
+    }
+
     @Test
     fun testDefaultState() =
         composeExtension.use {
-            setContentWithTheme {
-                FilterScreen(
-                    state =
-                        RelayFilterUiState(
-                            allProviders = DUMMY_RELAY_ALL_PROVIDERS,
-                            selectedOwnership = null,
-                            selectedProviders = DUMMY_SELECTED_PROVIDERS,
-                        ),
-                    onSelectedProvider = { _, _ -> },
-                )
-            }
+            initScreen(
+                state =
+                    RelayFilterUiState(
+                        allProviders = DUMMY_RELAY_ALL_PROVIDERS,
+                        selectedOwnership = null,
+                        selectedProviders = DUMMY_SELECTED_PROVIDERS,
+                    )
+            )
             onNodeWithText("Ownership").assertExists()
             onNodeWithText("Providers").assertExists()
         }
@@ -44,17 +62,14 @@ class FilterScreenTest {
     @Test
     fun testIsAnyCellShowing() =
         composeExtension.use {
-            setContentWithTheme {
-                FilterScreen(
-                    state =
-                        RelayFilterUiState(
-                            allProviders = DUMMY_RELAY_ALL_PROVIDERS,
-                            selectedOwnership = null,
-                            selectedProviders = DUMMY_SELECTED_PROVIDERS,
-                        ),
-                    onSelectedProvider = { _, _ -> },
-                )
-            }
+            initScreen(
+                state =
+                    RelayFilterUiState(
+                        allProviders = DUMMY_RELAY_ALL_PROVIDERS,
+                        selectedOwnership = null,
+                        selectedProviders = DUMMY_SELECTED_PROVIDERS,
+                    )
+            )
             onNodeWithText("Ownership").performClick()
             onNodeWithText("Any").assertExists()
         }
@@ -62,17 +77,14 @@ class FilterScreenTest {
     @Test
     fun testIsMullvadCellShowing() =
         composeExtension.use {
-            setContentWithTheme {
-                FilterScreen(
-                    state =
-                        RelayFilterUiState(
-                            allProviders = DUMMY_RELAY_ALL_PROVIDERS,
-                            selectedOwnership = Ownership.MullvadOwned,
-                            selectedProviders = DUMMY_SELECTED_PROVIDERS,
-                        ),
-                    onSelectedProvider = { _, _ -> },
-                )
-            }
+            initScreen(
+                state =
+                    RelayFilterUiState(
+                        allProviders = DUMMY_RELAY_ALL_PROVIDERS,
+                        selectedOwnership = Ownership.MullvadOwned,
+                        selectedProviders = DUMMY_SELECTED_PROVIDERS,
+                    )
+            )
             onNodeWithText("Ownership").performClick()
             onNodeWithText("Mullvad owned only").assertExists()
         }
@@ -80,17 +92,14 @@ class FilterScreenTest {
     @Test
     fun testIsRentedCellShowing() =
         composeExtension.use {
-            setContentWithTheme {
-                FilterScreen(
-                    state =
-                        RelayFilterUiState(
-                            allProviders = DUMMY_RELAY_ALL_PROVIDERS,
-                            selectedOwnership = Ownership.Rented,
-                            selectedProviders = DUMMY_SELECTED_PROVIDERS,
-                        ),
-                    onSelectedProvider = { _, _ -> },
-                )
-            }
+            initScreen(
+                state =
+                    RelayFilterUiState(
+                        allProviders = DUMMY_RELAY_ALL_PROVIDERS,
+                        selectedOwnership = Ownership.Rented,
+                        selectedProviders = DUMMY_SELECTED_PROVIDERS,
+                    )
+            )
             onNodeWithText("Ownership").performClick()
             onNodeWithText("Rented only").assertExists()
         }
@@ -98,17 +107,14 @@ class FilterScreenTest {
     @Test
     fun testShowProviders() =
         composeExtension.use {
-            setContentWithTheme {
-                FilterScreen(
-                    state =
-                        RelayFilterUiState(
-                            allProviders = DUMMY_RELAY_ALL_PROVIDERS,
-                            selectedOwnership = null,
-                            selectedProviders = DUMMY_SELECTED_PROVIDERS,
-                        ),
-                    onSelectedProvider = { _, _ -> },
-                )
-            }
+            initScreen(
+                state =
+                    RelayFilterUiState(
+                        allProviders = DUMMY_RELAY_ALL_PROVIDERS,
+                        selectedOwnership = null,
+                        selectedProviders = DUMMY_SELECTED_PROVIDERS,
+                    )
+            )
 
             onNodeWithText("Providers").performClick()
             onNodeWithText("Creanova").assertExists()
@@ -119,19 +125,16 @@ class FilterScreenTest {
     fun testApplyButtonClick() =
         composeExtension.use {
             val mockClickListener: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                FilterScreen(
-                    state =
-                        RelayFilterUiState(
-                            allProviders = listOf(),
-                            selectedOwnership = null,
-                            selectedProviders =
-                                listOf(Provider(ProviderId("31173"), Ownership.MullvadOwned)),
-                        ),
-                    onSelectedProvider = { _, _ -> },
-                    onApplyClick = mockClickListener,
-                )
-            }
+            initScreen(
+                state =
+                    RelayFilterUiState(
+                        allProviders = listOf(),
+                        selectedOwnership = null,
+                        selectedProviders =
+                            listOf(Provider(ProviderId("31173"), Ownership.MullvadOwned)),
+                    ),
+                onApplyClick = mockClickListener,
+            )
             onNodeWithText("Apply").performClick()
             verify { mockClickListener() }
         }
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 4294577836..ce683cc7fd 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,11 @@
 package net.mullvad.mullvadvpn.compose.screen
 
 import androidx.compose.ui.test.ExperimentalTestApi
+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 de.mannodermaus.junit5.compose.ComposeContext
 import io.mockk.MockKAnnotations
 import io.mockk.every
 import io.mockk.mockk
@@ -31,20 +33,36 @@ class OutOfTimeScreenTest {
         MockKAnnotations.init(this)
     }
 
+    private fun ComposeContext.initScreen(
+        state: OutOfTimeUiState = OutOfTimeUiState(),
+        onDisconnectClick: () -> Unit = {},
+        onSitePaymentClick: () -> Unit = {},
+        onRedeemVoucherClick: () -> Unit = {},
+        onSettingsClick: () -> Unit = {},
+        onAccountClick: () -> Unit = {},
+        onPurchaseBillingProductClick: (ProductId) -> Unit = {},
+        navigateToVerificationPendingDialog: () -> Unit = {},
+    ) {
+
+        setContentWithTheme {
+            OutOfTimeScreen(
+                state = state,
+                onDisconnectClick = onDisconnectClick,
+                onSitePaymentClick = onSitePaymentClick,
+                onRedeemVoucherClick = onRedeemVoucherClick,
+                onSettingsClick = onSettingsClick,
+                onAccountClick = onAccountClick,
+                onPurchaseBillingProductClick = onPurchaseBillingProductClick,
+                navigateToVerificationPendingDialog = navigateToVerificationPendingDialog,
+            )
+        }
+    }
+
     @Test
     fun testDisableSitePayment() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                OutOfTimeScreen(
-                    state = OutOfTimeUiState(deviceName = ""),
-                    onSitePaymentClick = {},
-                    onRedeemVoucherClick = {},
-                    onSettingsClick = {},
-                    onAccountClick = {},
-                    onDisconnectClick = {},
-                )
-            }
+            initScreen(state = OutOfTimeUiState(deviceName = ""))
 
             // Assert
             onNodeWithText(
@@ -58,20 +76,18 @@ class OutOfTimeScreenTest {
     @Test
     fun testOpenAccountView() =
         composeExtension.use {
+            val mockClickListener: () -> Unit = mockk(relaxed = true)
+
             // Arrange
-            setContentWithTheme {
-                OutOfTimeScreen(
-                    state = OutOfTimeUiState(deviceName = "", showSitePayment = true),
-                    onSitePaymentClick = {},
-                    onRedeemVoucherClick = {},
-                    onSettingsClick = {},
-                    onAccountClick = {},
-                    onDisconnectClick = {},
-                )
-            }
+            initScreen(
+                state = OutOfTimeUiState(deviceName = "", showSitePayment = true),
+                onAccountClick = mockClickListener,
+            )
+
+            onNodeWithContentDescription(label = "Account").performClick()
 
             // Assert
-            onNodeWithText("Congrats!").assertDoesNotExist()
+            verify(exactly = 1) { mockClickListener.invoke() }
         }
 
     @Test
@@ -79,16 +95,10 @@ class OutOfTimeScreenTest {
         composeExtension.use {
             // Arrange
             val mockClickListener: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                OutOfTimeScreen(
-                    state = OutOfTimeUiState(deviceName = "", showSitePayment = true),
-                    onSitePaymentClick = mockClickListener,
-                    onRedeemVoucherClick = {},
-                    onSettingsClick = {},
-                    onAccountClick = {},
-                    onDisconnectClick = {},
-                )
-            }
+            initScreen(
+                state = OutOfTimeUiState(deviceName = "", showSitePayment = true),
+                onSitePaymentClick = mockClickListener,
+            )
 
             // Act
             onNodeWithText("Buy credit").performClick()
@@ -102,16 +112,10 @@ class OutOfTimeScreenTest {
         composeExtension.use {
             // Arrange
             val mockClickListener: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                OutOfTimeScreen(
-                    state = OutOfTimeUiState(deviceName = "", showSitePayment = true),
-                    onSitePaymentClick = {},
-                    onRedeemVoucherClick = mockClickListener,
-                    onSettingsClick = {},
-                    onAccountClick = {},
-                    onDisconnectClick = {},
-                )
-            }
+            initScreen(
+                state = OutOfTimeUiState(deviceName = "", showSitePayment = true),
+                onRedeemVoucherClick = mockClickListener,
+            )
 
             // Act
             onNodeWithText("Redeem voucher").performClick()
@@ -125,21 +129,15 @@ class OutOfTimeScreenTest {
         composeExtension.use {
             // Arrange
             val mockClickListener: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                OutOfTimeScreen(
-                    state =
-                        OutOfTimeUiState(
-                            tunnelState = TunnelState.Connecting(null, null, emptyList()),
-                            deviceName = "",
-                            showSitePayment = true,
-                        ),
-                    onSitePaymentClick = {},
-                    onRedeemVoucherClick = {},
-                    onSettingsClick = {},
-                    onAccountClick = {},
-                    onDisconnectClick = mockClickListener,
-                )
-            }
+            initScreen(
+                state =
+                    OutOfTimeUiState(
+                        tunnelState = TunnelState.Connecting(null, null, emptyList()),
+                        deviceName = "",
+                        showSitePayment = true,
+                    ),
+                onDisconnectClick = mockClickListener,
+            )
 
             // Act
             onNodeWithText("Disconnect").performClick()
@@ -152,20 +150,13 @@ class OutOfTimeScreenTest {
     fun testShowBillingErrorPaymentButton() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                OutOfTimeScreen(
-                    state =
-                        OutOfTimeUiState(
-                            showSitePayment = true,
-                            billingPaymentState = PaymentState.Error.Billing,
-                        ),
-                    onSitePaymentClick = {},
-                    onRedeemVoucherClick = {},
-                    onSettingsClick = {},
-                    onAccountClick = {},
-                    onPurchaseBillingProductClick = { _ -> },
-                )
-            }
+            initScreen(
+                state =
+                    OutOfTimeUiState(
+                        showSitePayment = true,
+                        billingPaymentState = PaymentState.Error.Billing,
+                    )
+            )
 
             // Assert
             onNodeWithText("Add 30 days time").assertExists()
@@ -178,21 +169,14 @@ class OutOfTimeScreenTest {
             val mockPaymentProduct: PaymentProduct = mockk()
             every { mockPaymentProduct.price } returns ProductPrice("$10")
             every { mockPaymentProduct.status } returns null
-            setContentWithTheme {
-                OutOfTimeScreen(
-                    state =
-                        OutOfTimeUiState(
-                            showSitePayment = true,
-                            billingPaymentState =
-                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
-                        ),
-                    onSitePaymentClick = {},
-                    onRedeemVoucherClick = {},
-                    onSettingsClick = {},
-                    onAccountClick = {},
-                    onPurchaseBillingProductClick = { _ -> },
-                )
-            }
+            initScreen(
+                state =
+                    OutOfTimeUiState(
+                        showSitePayment = true,
+                        billingPaymentState =
+                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
+                    )
+            )
 
             // Assert
             onNodeWithText("Add 30 days time ($10)").assertExists()
@@ -205,16 +189,14 @@ class OutOfTimeScreenTest {
             val mockPaymentProduct: PaymentProduct = mockk()
             every { mockPaymentProduct.price } returns ProductPrice("$10")
             every { mockPaymentProduct.status } returns PaymentStatus.PENDING
-            setContentWithTheme {
-                OutOfTimeScreen(
-                    state =
-                        OutOfTimeUiState(
-                            showSitePayment = true,
-                            billingPaymentState =
-                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    OutOfTimeUiState(
+                        showSitePayment = true,
+                        billingPaymentState =
+                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
+                    )
+            )
 
             // Assert
             onNodeWithText("Google Play payment pending").assertExists()
@@ -228,17 +210,15 @@ class OutOfTimeScreenTest {
             every { mockPaymentProduct.price } returns ProductPrice("$10")
             every { mockPaymentProduct.status } returns PaymentStatus.PENDING
             val mockNavigateToVerificationPending: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                OutOfTimeScreen(
-                    state =
-                        OutOfTimeUiState(
-                            showSitePayment = true,
-                            billingPaymentState =
-                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
-                        ),
-                    navigateToVerificationPendingDialog = mockNavigateToVerificationPending,
-                )
-            }
+            initScreen(
+                state =
+                    OutOfTimeUiState(
+                        showSitePayment = true,
+                        billingPaymentState =
+                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
+                    ),
+                navigateToVerificationPendingDialog = mockNavigateToVerificationPending,
+            )
 
             // Act
             onNodeWithTag(PLAY_PAYMENT_INFO_ICON_TEST_TAG).performClick()
@@ -255,16 +235,14 @@ class OutOfTimeScreenTest {
             val mockPaymentProduct: PaymentProduct = mockk()
             every { mockPaymentProduct.price } returns ProductPrice("$10")
             every { mockPaymentProduct.status } returns PaymentStatus.VERIFICATION_IN_PROGRESS
-            setContentWithTheme {
-                OutOfTimeScreen(
-                    state =
-                        OutOfTimeUiState(
-                            billingPaymentState =
-                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
-                            showSitePayment = true,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    OutOfTimeUiState(
+                        billingPaymentState =
+                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
+                        showSitePayment = true,
+                    )
+            )
 
             // Assert
             onNodeWithText("Verifying purchase").assertExists()
@@ -279,21 +257,15 @@ class OutOfTimeScreenTest {
             every { mockPaymentProduct.price } returns ProductPrice("$10")
             every { mockPaymentProduct.productId } returns ProductId("PRODUCT_ID")
             every { mockPaymentProduct.status } returns null
-            setContentWithTheme {
-                OutOfTimeScreen(
-                    state =
-                        OutOfTimeUiState(
-                            billingPaymentState =
-                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
-                            showSitePayment = true,
-                        ),
-                    onSitePaymentClick = {},
-                    onRedeemVoucherClick = {},
-                    onSettingsClick = {},
-                    onAccountClick = {},
-                    onPurchaseBillingProductClick = clickHandler,
-                )
-            }
+            initScreen(
+                state =
+                    OutOfTimeUiState(
+                        billingPaymentState =
+                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
+                        showSitePayment = true,
+                    ),
+                onPurchaseBillingProductClick = clickHandler,
+            )
 
             // Act
             onNodeWithText("Add 30 days time ($10)").performClick()
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ServerIpOverridesScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ServerIpOverridesScreenTest.kt
index 8d33d848f1..0b11616184 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ServerIpOverridesScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ServerIpOverridesScreenTest.kt
@@ -1,10 +1,10 @@
 package net.mullvad.mullvadvpn.compose.screen
 
-import androidx.compose.runtime.Composable
 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.ComposeContext
 import io.mockk.MockKAnnotations
 import io.mockk.mockk
 import io.mockk.verify
@@ -30,9 +30,7 @@ class ServerIpOverridesScreenTest {
         MockKAnnotations.init(this)
     }
 
-    @Suppress("TestFunctionName")
-    @Composable
-    private fun ScreenWithDefault(
+    private fun ComposeContext.initScreen(
         state: ServerIpOverridesUiState,
         onBackClick: () -> Unit = {},
         onInfoClick: () -> Unit = {},
@@ -40,43 +38,43 @@ class ServerIpOverridesScreenTest {
         onImportByFile: () -> Unit = {},
         onImportByText: () -> Unit = {},
     ) {
-        ServerIpOverridesScreen(
-            state = state,
-            onBackClick = onBackClick,
-            onInfoClick = onInfoClick,
-            onResetOverridesClick = onResetOverridesClick,
-            onImportByFile = onImportByFile,
-            onImportByText = onImportByText,
-        )
+        setContentWithTheme {
+            ServerIpOverridesScreen(
+                state = state,
+                onBackClick = onBackClick,
+                onInfoClick = onInfoClick,
+                onResetOverridesClick = onResetOverridesClick,
+                onImportByFile = onImportByFile,
+                onImportByText = onImportByText,
+            )
+        }
     }
 
     @Test
-    fun ensure_overrides_inactive_is_displayed() =
+    fun ensureOverridesInactiveIsDisplayed() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                ScreenWithDefault(state = ServerIpOverridesUiState.Loaded(false))
-            }
+            initScreen(state = ServerIpOverridesUiState.Loaded(false))
 
             // Assert
             onNodeWithText("Overrides inactive").assertExists()
         }
 
     @Test
-    fun ensure_overrides_active_is_displayed() =
+    fun ensureOverridesActiveIsDisplayed() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme { ScreenWithDefault(state = ServerIpOverridesUiState.Loaded(true)) }
+            initScreen(state = ServerIpOverridesUiState.Loaded(true))
 
             // Assert
             onNodeWithText("Overrides active").assertExists()
         }
 
     @Test
-    fun ensure_overrides_active_shows_warning_on_import() =
+    fun ensureOverridesActiveShowsWarningOnImport() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme { ScreenWithDefault(state = ServerIpOverridesUiState.Loaded(true)) }
+            initScreen(state = ServerIpOverridesUiState.Loaded(true))
 
             // Act
             onNodeWithTag(testTag = SERVER_IP_OVERRIDE_IMPORT_TEST_TAG).performClick()
@@ -89,16 +87,11 @@ class ServerIpOverridesScreenTest {
         }
 
     @Test
-    fun ensure_info_click_works() =
+    fun ensureInfoClickWorks() =
         composeExtension.use {
             // Arrange
             val clickHandler: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                ScreenWithDefault(
-                    state = ServerIpOverridesUiState.Loaded(false),
-                    onInfoClick = clickHandler,
-                )
-            }
+            initScreen(state = ServerIpOverridesUiState.Loaded(false), onInfoClick = clickHandler)
 
             // Act
             onNodeWithTag(SERVER_IP_OVERRIDE_INFO_TEST_TAG).performClick()
@@ -108,16 +101,14 @@ class ServerIpOverridesScreenTest {
         }
 
     @Test
-    fun ensure_reset_click_works() =
+    fun ensureResetClickWorks() =
         composeExtension.use {
             // Arrange
             val clickHandler: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                ScreenWithDefault(
-                    state = ServerIpOverridesUiState.Loaded(true),
-                    onResetOverridesClick = clickHandler,
-                )
-            }
+            initScreen(
+                state = ServerIpOverridesUiState.Loaded(true),
+                onResetOverridesClick = clickHandler,
+            )
 
             // Act
             onNodeWithTag(SERVER_IP_OVERRIDE_MORE_VERT_TEST_TAG).performClick()
@@ -128,16 +119,14 @@ class ServerIpOverridesScreenTest {
         }
 
     @Test
-    fun ensure_import_by_file_works() =
+    fun ensureImportByFileWorks() =
         composeExtension.use {
             // Arrange
             val clickHandler: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                ScreenWithDefault(
-                    state = ServerIpOverridesUiState.Loaded(false),
-                    onImportByFile = clickHandler,
-                )
-            }
+            initScreen(
+                state = ServerIpOverridesUiState.Loaded(false),
+                onImportByFile = clickHandler,
+            )
 
             // Act
             onNodeWithTag(SERVER_IP_OVERRIDE_IMPORT_TEST_TAG).performClick()
@@ -148,16 +137,14 @@ class ServerIpOverridesScreenTest {
         }
 
     @Test
-    fun ensure_import_by_text() =
+    fun ensureImportByText() =
         composeExtension.use {
             // Arrange
             val clickHandler: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                ScreenWithDefault(
-                    state = ServerIpOverridesUiState.Loaded(false),
-                    onImportByText = clickHandler,
-                )
-            }
+            initScreen(
+                state = ServerIpOverridesUiState.Loaded(false),
+                onImportByText = clickHandler,
+            )
 
             // Act
             onNodeWithTag(SERVER_IP_OVERRIDE_IMPORT_TEST_TAG).performClick()
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 464fa1181a..b343c44c95 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
@@ -3,6 +3,7 @@ package net.mullvad.mullvadvpn.compose.screen
 import androidx.compose.material3.ExperimentalMaterial3Api
 import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.onNodeWithText
+import de.mannodermaus.junit5.compose.ComposeContext
 import io.mockk.MockKAnnotations
 import net.mullvad.mullvadvpn.compose.createEdgeToEdgeComposeExtension
 import net.mullvad.mullvadvpn.compose.setContentWithTheme
@@ -20,24 +21,48 @@ class SettingsScreenTest {
         MockKAnnotations.init(this)
     }
 
+    private fun ComposeContext.initScreen(
+        state: SettingsUiState,
+        onVpnSettingCellClick: () -> Unit = {},
+        onSplitTunnelingCellClick: () -> Unit = {},
+        onAppInfoClick: () -> Unit = {},
+        onReportProblemCellClick: () -> Unit = {},
+        onApiAccessClick: () -> Unit = {},
+        onMultihopClick: () -> Unit = {},
+        onDaitaClick: () -> Unit = {},
+        onBackClick: () -> Unit = {},
+    ) {
+        setContentWithTheme {
+            SettingsScreen(
+                state = state,
+                onVpnSettingCellClick = onVpnSettingCellClick,
+                onSplitTunnelingCellClick = onSplitTunnelingCellClick,
+                onAppInfoClick = onAppInfoClick,
+                onReportProblemCellClick = onReportProblemCellClick,
+                onApiAccessClick = onApiAccessClick,
+                onMultihopClick = onMultihopClick,
+                onDaitaClick = onDaitaClick,
+                onBackClick = onBackClick,
+            )
+        }
+    }
+
     @Test
     @OptIn(ExperimentalMaterial3Api::class)
     fun testLoggedInState() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                SettingsScreen(
-                    state =
-                        SettingsUiState(
-                            appVersion = "",
-                            isLoggedIn = true,
-                            isSupportedVersion = true,
-                            isPlayBuild = false,
-                            multihopEnabled = false,
-                            isDaitaEnabled = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    SettingsUiState(
+                        appVersion = "",
+                        isLoggedIn = true,
+                        isSupportedVersion = true,
+                        isPlayBuild = false,
+                        multihopEnabled = false,
+                        isDaitaEnabled = false,
+                    )
+            )
             // Assert
             onNodeWithText("VPN settings").assertExists()
             onNodeWithText("Split tunneling").assertExists()
@@ -50,19 +75,17 @@ class SettingsScreenTest {
     fun testLoggedOutState() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                SettingsScreen(
-                    state =
-                        SettingsUiState(
-                            appVersion = "",
-                            isLoggedIn = false,
-                            isSupportedVersion = true,
-                            isPlayBuild = false,
-                            multihopEnabled = false,
-                            isDaitaEnabled = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    SettingsUiState(
+                        appVersion = "",
+                        isLoggedIn = false,
+                        isSupportedVersion = true,
+                        isPlayBuild = false,
+                        multihopEnabled = false,
+                        isDaitaEnabled = false,
+                    )
+            )
             // Assert
             onNodeWithText("VPN settings").assertDoesNotExist()
             onNodeWithText("Split tunneling").assertDoesNotExist()
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ShadowsocksSettingsScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ShadowsocksSettingsScreenTest.kt
index 23d94f58c6..1340e96c58 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ShadowsocksSettingsScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ShadowsocksSettingsScreenTest.kt
@@ -4,6 +4,7 @@ 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.ComposeContext
 import io.mockk.mockk
 import io.mockk.verify
 import net.mullvad.mullvadvpn.compose.createEdgeToEdgeComposeExtension
@@ -19,13 +20,27 @@ import org.junit.jupiter.api.extension.RegisterExtension
 class ShadowsocksSettingsScreenTest {
     @JvmField @RegisterExtension val composeExtension = createEdgeToEdgeComposeExtension()
 
+    private fun ComposeContext.initScreen(
+        state: ShadowsocksSettingsState = ShadowsocksSettingsState(),
+        navigateToCustomPortDialog: () -> Unit = {},
+        onObfuscationPortSelected: (Constraint<Port>) -> Unit = {},
+        onBackClick: () -> Unit = {},
+    ) {
+        setContentWithTheme {
+            ShadowsocksSettingsScreen(
+                state = state,
+                navigateToCustomPortDialog = navigateToCustomPortDialog,
+                onObfuscationPortSelected = onObfuscationPortSelected,
+                onBackClick = onBackClick,
+            )
+        }
+    }
+
     @Test
     fun testShowShadowsocksCustomPort() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                ShadowsocksSettingsScreen(state = ShadowsocksSettingsState(customPort = Port(4000)))
-            }
+            initScreen(state = ShadowsocksSettingsState(customPort = Port(4000)))
 
             // Assert
             onNodeWithText("4000").assertExists()
@@ -36,16 +51,14 @@ class ShadowsocksSettingsScreenTest {
         composeExtension.use {
             // Arrange
             val onObfuscationPortSelected: (Constraint<Port>) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                ShadowsocksSettingsScreen(
-                    state =
-                        ShadowsocksSettingsState(
-                            port = Constraint.Only(Port(4000)),
-                            customPort = Port(4000),
-                        ),
-                    onObfuscationPortSelected = onObfuscationPortSelected,
-                )
-            }
+            initScreen(
+                state =
+                    ShadowsocksSettingsState(
+                        port = Constraint.Only(Port(4000)),
+                        customPort = Port(4000),
+                    ),
+                onObfuscationPortSelected = onObfuscationPortSelected,
+            )
 
             // Act
             onNodeWithTag(testTag = SHADOWSOCKS_CUSTOM_PORT_TEXT_TEST_TAG).performClick()
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 500349c861..8215ccde69 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,8 +1,10 @@
 package net.mullvad.mullvadvpn.compose.screen
 
+import android.graphics.Bitmap
 import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
+import de.mannodermaus.junit5.compose.ComposeContext
 import io.mockk.MockKAnnotations
 import io.mockk.mockk
 import io.mockk.unmockkAll
@@ -30,13 +32,33 @@ class SplitTunnelingScreenTest {
         unmockkAll()
     }
 
+    private fun ComposeContext.initScreen(
+        state: SplitTunnelingUiState,
+        onEnableSplitTunneling: (Boolean) -> Unit = {},
+        onShowSystemAppsClick: (show: Boolean) -> Unit = {},
+        onExcludeAppClick: (packageName: String) -> Unit = {},
+        onIncludeAppClick: (packageName: String) -> Unit = {},
+        onBackClick: () -> Unit = {},
+        onResolveIcon: (String) -> Bitmap? = { null },
+    ) {
+        setContentWithTheme {
+            SplitTunnelingScreen(
+                state = state,
+                onEnableSplitTunneling = onEnableSplitTunneling,
+                onShowSystemAppsClick = onShowSystemAppsClick,
+                onExcludeAppClick = onExcludeAppClick,
+                onIncludeAppClick = onIncludeAppClick,
+                onBackClick = onBackClick,
+                onResolveIcon = onResolveIcon,
+            )
+        }
+    }
+
     @Test
     fun testLoadingState() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                SplitTunnelingScreen(state = SplitTunnelingUiState.Loading(enabled = true))
-            }
+            initScreen(state = SplitTunnelingUiState.Loading(enabled = true))
 
             // Assert
             onNodeWithText(TITLE).assertExists()
@@ -62,17 +84,15 @@ class SplitTunnelingScreenTest {
                     iconRes = 0,
                     name = INCLUDED_APP_NAME,
                 )
-            setContentWithTheme {
-                SplitTunnelingScreen(
-                    state =
-                        SplitTunnelingUiState.ShowAppList(
-                            enabled = true,
-                            excludedApps = listOf(excludedApp),
-                            includedApps = listOf(includedApp),
-                            showSystemApps = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    SplitTunnelingUiState.ShowAppList(
+                        enabled = true,
+                        excludedApps = listOf(excludedApp),
+                        includedApps = listOf(includedApp),
+                        showSystemApps = false,
+                    )
+            )
 
             // Assert
             onNodeWithText(TITLE).assertExists()
@@ -94,17 +114,15 @@ class SplitTunnelingScreenTest {
                     iconRes = 0,
                     name = INCLUDED_APP_NAME,
                 )
-            setContentWithTheme {
-                SplitTunnelingScreen(
-                    state =
-                        SplitTunnelingUiState.ShowAppList(
-                            enabled = true,
-                            excludedApps = emptyList(),
-                            includedApps = listOf(includedApp),
-                            showSystemApps = false,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    SplitTunnelingUiState.ShowAppList(
+                        enabled = true,
+                        excludedApps = emptyList(),
+                        includedApps = listOf(includedApp),
+                        showSystemApps = false,
+                    )
+            )
 
             // Assert
             onNodeWithText(TITLE).assertExists()
@@ -133,18 +151,16 @@ class SplitTunnelingScreenTest {
                     name = INCLUDED_APP_NAME,
                 )
             val mockedClickHandler: (String) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                SplitTunnelingScreen(
-                    state =
-                        SplitTunnelingUiState.ShowAppList(
-                            enabled = true,
-                            excludedApps = listOf(excludedApp),
-                            includedApps = listOf(includedApp),
-                            showSystemApps = false,
-                        ),
-                    onExcludeAppClick = mockedClickHandler,
-                )
-            }
+            initScreen(
+                state =
+                    SplitTunnelingUiState.ShowAppList(
+                        enabled = true,
+                        excludedApps = listOf(excludedApp),
+                        includedApps = listOf(includedApp),
+                        showSystemApps = false,
+                    ),
+                onExcludeAppClick = mockedClickHandler,
+            )
 
             // Act
             onNodeWithText(INCLUDED_APP_NAME).performClick()
@@ -170,18 +186,16 @@ class SplitTunnelingScreenTest {
                     name = INCLUDED_APP_NAME,
                 )
             val mockedClickHandler: (String) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                SplitTunnelingScreen(
-                    state =
-                        SplitTunnelingUiState.ShowAppList(
-                            enabled = true,
-                            excludedApps = listOf(excludedApp),
-                            includedApps = listOf(includedApp),
-                            showSystemApps = false,
-                        ),
-                    onIncludeAppClick = mockedClickHandler,
-                )
-            }
+            initScreen(
+                state =
+                    SplitTunnelingUiState.ShowAppList(
+                        enabled = true,
+                        excludedApps = listOf(excludedApp),
+                        includedApps = listOf(includedApp),
+                        showSystemApps = false,
+                    ),
+                onIncludeAppClick = mockedClickHandler,
+            )
 
             // Act
             onNodeWithText(EXCLUDED_APP_NAME).performClick()
@@ -207,18 +221,16 @@ class SplitTunnelingScreenTest {
                     name = INCLUDED_APP_NAME,
                 )
             val mockedClickHandler: (Boolean) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                SplitTunnelingScreen(
-                    state =
-                        SplitTunnelingUiState.ShowAppList(
-                            enabled = true,
-                            excludedApps = listOf(excludedApp),
-                            includedApps = listOf(includedApp),
-                            showSystemApps = false,
-                        ),
-                    onShowSystemAppsClick = mockedClickHandler,
-                )
-            }
+            initScreen(
+                state =
+                    SplitTunnelingUiState.ShowAppList(
+                        enabled = true,
+                        excludedApps = listOf(excludedApp),
+                        includedApps = listOf(includedApp),
+                        showSystemApps = false,
+                    ),
+                onShowSystemAppsClick = mockedClickHandler,
+            )
 
             // Act
             onNodeWithText(SHOW_SYSTEM_APPS).performClick()
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/Udp2TcpSettingsScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/Udp2TcpSettingsScreenTest.kt
index 77b9965b26..1c16ff0bd1 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/Udp2TcpSettingsScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/Udp2TcpSettingsScreenTest.kt
@@ -2,6 +2,7 @@ package net.mullvad.mullvadvpn.compose.screen
 
 import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.performClick
+import de.mannodermaus.junit5.compose.ComposeContext
 import io.mockk.coVerify
 import io.mockk.mockk
 import net.mullvad.mullvadvpn.compose.createEdgeToEdgeComposeExtension
@@ -18,17 +19,31 @@ import org.junit.jupiter.api.extension.RegisterExtension
 class Udp2TcpSettingsScreenTest {
     @JvmField @RegisterExtension val composeExtension = createEdgeToEdgeComposeExtension()
 
+    private fun ComposeContext.initScreen(
+        state: Udp2TcpSettingsState = Udp2TcpSettingsState(),
+        onObfuscationPortSelected: (Constraint<Port>) -> Unit = {},
+        navigateUdp2TcpInfo: () -> Unit = {},
+        onBackClick: () -> Unit = {},
+    ) {
+        setContentWithTheme {
+            Udp2TcpSettingsScreen(
+                state = state,
+                onObfuscationPortSelected = onObfuscationPortSelected,
+                navigateUdp2TcpInfo = navigateUdp2TcpInfo,
+                onBackClick = onBackClick,
+            )
+        }
+    }
+
     @Test
     fun testSelectTcpOverUdpPortOption() =
         composeExtension.use {
             // Arrange
             val onObfuscationPortSelected: (Constraint<Port>) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                Udp2TcpSettingsScreen(
-                    state = Udp2TcpSettingsState(port = Constraint.Any),
-                    onObfuscationPortSelected = onObfuscationPortSelected,
-                )
-            }
+            initScreen(
+                state = Udp2TcpSettingsState(port = Constraint.Any),
+                onObfuscationPortSelected = onObfuscationPortSelected,
+            )
 
             // Act
             onNodeWithTagAndText(
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 19802571f6..43c5559a38 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
@@ -7,6 +7,7 @@ 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.ComposeContext
 import io.mockk.MockKAnnotations
 import io.mockk.mockk
 import io.mockk.verify
@@ -23,6 +24,7 @@ import net.mullvad.mullvadvpn.compose.test.LAZY_LIST_WIREGUARD_OBFUSCATION_TITLE
 import net.mullvad.mullvadvpn.compose.test.LAZY_LIST_WIREGUARD_PORT_ITEM_X_TEST_TAG
 import net.mullvad.mullvadvpn.lib.model.Constraint
 import net.mullvad.mullvadvpn.lib.model.Mtu
+import net.mullvad.mullvadvpn.lib.model.ObfuscationMode
 import net.mullvad.mullvadvpn.lib.model.Port
 import net.mullvad.mullvadvpn.lib.model.PortRange
 import net.mullvad.mullvadvpn.lib.model.QuantumResistantState
@@ -41,11 +43,75 @@ class VpnSettingsScreenTest {
         MockKAnnotations.init(this)
     }
 
+    private fun ComposeContext.initScreen(
+        state: VpnSettingsUiState = VpnSettingsUiState.createDefault(),
+        navigateToContentBlockersInfo: () -> Unit = {},
+        navigateToAutoConnectScreen: () -> Unit = {},
+        navigateToCustomDnsInfo: () -> Unit = {},
+        navigateToMalwareInfo: () -> Unit = {},
+        navigateToObfuscationInfo: () -> Unit = {},
+        navigateToQuantumResistanceInfo: () -> Unit = {},
+        navigateToWireguardPortInfo: (availablePortRanges: List<PortRange>) -> Unit = {},
+        navigateToLocalNetworkSharingInfo: () -> Unit = {},
+        navigateToWireguardPortDialog: () -> Unit = {},
+        navigateToServerIpOverrides: () -> Unit = {},
+        onToggleBlockTrackers: (Boolean) -> Unit = {},
+        onToggleBlockAds: (Boolean) -> Unit = {},
+        onToggleBlockMalware: (Boolean) -> Unit = {},
+        onToggleLocalNetworkSharing: (Boolean) -> Unit = {},
+        onToggleBlockAdultContent: (Boolean) -> Unit = {},
+        onToggleBlockGambling: (Boolean) -> Unit = {},
+        onToggleBlockSocialMedia: (Boolean) -> Unit = {},
+        navigateToMtuDialog: (mtu: Mtu?) -> Unit = {},
+        navigateToDns: (index: Int?, address: String?) -> Unit = { _, _ -> },
+        onToggleDnsClick: (Boolean) -> Unit = {},
+        onBackClick: () -> Unit = {},
+        onSelectObfuscationMode: (obfuscationMode: ObfuscationMode) -> Unit = {},
+        onSelectQuantumResistanceSetting: (quantumResistant: QuantumResistantState) -> Unit = {},
+        onWireguardPortSelected: (port: Constraint<Port>) -> Unit = {},
+        navigateToShadowSocksSettings: () -> Unit = {},
+        navigateToUdp2TcpSettings: () -> Unit = {},
+        onToggleAutoStartAndConnectOnBoot: (Boolean) -> Unit = {},
+    ) {
+        setContentWithTheme {
+            VpnSettingsScreen(
+                state = state,
+                navigateToContentBlockersInfo = navigateToContentBlockersInfo,
+                navigateToAutoConnectScreen = navigateToAutoConnectScreen,
+                navigateToCustomDnsInfo = navigateToCustomDnsInfo,
+                navigateToMalwareInfo = navigateToMalwareInfo,
+                navigateToObfuscationInfo = navigateToObfuscationInfo,
+                navigateToQuantumResistanceInfo = navigateToQuantumResistanceInfo,
+                navigateToWireguardPortInfo = navigateToWireguardPortInfo,
+                navigateToLocalNetworkSharingInfo = navigateToLocalNetworkSharingInfo,
+                navigateToWireguardPortDialog = navigateToWireguardPortDialog,
+                navigateToServerIpOverrides = navigateToServerIpOverrides,
+                onToggleBlockTrackers = onToggleBlockTrackers,
+                onToggleBlockAds = onToggleBlockAds,
+                onToggleBlockMalware = onToggleBlockMalware,
+                onToggleLocalNetworkSharing = onToggleLocalNetworkSharing,
+                onToggleBlockAdultContent = onToggleBlockAdultContent,
+                onToggleBlockGambling = onToggleBlockGambling,
+                onToggleBlockSocialMedia = onToggleBlockSocialMedia,
+                navigateToMtuDialog = navigateToMtuDialog,
+                navigateToDns = navigateToDns,
+                onToggleDnsClick = onToggleDnsClick,
+                onBackClick = onBackClick,
+                onSelectObfuscationMode = onSelectObfuscationMode,
+                onSelectQuantumResistanceSetting = onSelectQuantumResistanceSetting,
+                onWireguardPortSelected = onWireguardPortSelected,
+                navigateToShadowSocksSettings = navigateToShadowSocksSettings,
+                navigateToUdp2TcpSettings = navigateToUdp2TcpSettings,
+                onToggleAutoStartAndConnectOnBoot = onToggleAutoStartAndConnectOnBoot,
+            )
+        }
+    }
+
     @Test
     fun testDefaultState() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme { VpnSettingsScreen(state = VpnSettingsUiState.createDefault()) }
+            initScreen(state = VpnSettingsUiState.createDefault())
 
             onNodeWithTag(LAZY_LIST_VPN_SETTINGS_TEST_TAG)
                 .performScrollToNode(hasTestTag(LAZY_LIST_LAST_ITEM_TEST_TAG))
@@ -61,14 +127,12 @@ class VpnSettingsScreenTest {
     fun testMtuCustomValue() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state =
-                        VpnSettingsUiState.createDefault(
-                            mtu = Mtu.fromString(VALID_DUMMY_MTU_VALUE).getOrNull()!!
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    VpnSettingsUiState.createDefault(
+                        mtu = Mtu.fromString(VALID_DUMMY_MTU_VALUE).getOrNull()!!
+                    )
+            )
 
             onNodeWithTag(LAZY_LIST_VPN_SETTINGS_TEST_TAG)
                 .performScrollToNode(hasTestTag(LAZY_LIST_LAST_ITEM_TEST_TAG))
@@ -81,20 +145,18 @@ class VpnSettingsScreenTest {
     fun testCustomDnsAddressesAndAddButtonVisibleWhenCustomDnsEnabled() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state =
-                        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),
-                                ),
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    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
             onNodeWithText(DUMMY_DNS_ADDRESS).assertExists()
@@ -107,16 +169,13 @@ class VpnSettingsScreenTest {
     fun testCustomDnsAddressesAndAddButtonNotVisibleWhenCustomDnsDisabled() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state =
-                        VpnSettingsUiState.createDefault(
-                            isCustomDnsEnabled = false,
-                            customDnsItems =
-                                listOf(CustomDnsItem(address = DUMMY_DNS_ADDRESS, false)),
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    VpnSettingsUiState.createDefault(
+                        isCustomDnsEnabled = false,
+                        customDnsItems = listOf(CustomDnsItem(address = DUMMY_DNS_ADDRESS, false)),
+                    )
+            )
             onNodeWithTag(LAZY_LIST_VPN_SETTINGS_TEST_TAG)
                 .performScrollToNode(hasTestTag(LAZY_LIST_LAST_ITEM_TEST_TAG))
             // Assert
@@ -128,17 +187,15 @@ class VpnSettingsScreenTest {
     fun testLanWarningNotShownWhenLanTrafficEnabledAndLocalAddressIsUsed() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state =
-                        VpnSettingsUiState.createDefault(
-                            isCustomDnsEnabled = true,
-                            isLocalNetworkSharingEnabled = true,
-                            customDnsItems =
-                                listOf(CustomDnsItem(address = DUMMY_DNS_ADDRESS, isLocal = true)),
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    VpnSettingsUiState.createDefault(
+                        isCustomDnsEnabled = true,
+                        isLocalNetworkSharingEnabled = true,
+                        customDnsItems =
+                            listOf(CustomDnsItem(address = DUMMY_DNS_ADDRESS, isLocal = true)),
+                    )
+            )
 
             // Assert
             onNodeWithContentDescription(LOCAL_DNS_SERVER_WARNING).assertDoesNotExist()
@@ -148,16 +205,14 @@ class VpnSettingsScreenTest {
     fun testLanWarningNotShowedWhenLanTrafficDisabledAndLocalAddressIsNotUsed() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state =
-                        VpnSettingsUiState.createDefault(
-                            isCustomDnsEnabled = true,
-                            customDnsItems =
-                                listOf(CustomDnsItem(address = DUMMY_DNS_ADDRESS, isLocal = false)),
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    VpnSettingsUiState.createDefault(
+                        isCustomDnsEnabled = true,
+                        customDnsItems =
+                            listOf(CustomDnsItem(address = DUMMY_DNS_ADDRESS, isLocal = false)),
+                    )
+            )
 
             // Assert
             onNodeWithContentDescription(LOCAL_DNS_SERVER_WARNING).assertDoesNotExist()
@@ -167,16 +222,14 @@ class VpnSettingsScreenTest {
     fun testLanWarningNotShowedWhenLanTrafficEnabledAndLocalAddressIsNotUsed() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state =
-                        VpnSettingsUiState.createDefault(
-                            isCustomDnsEnabled = true,
-                            customDnsItems =
-                                listOf(CustomDnsItem(address = DUMMY_DNS_ADDRESS, isLocal = false)),
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    VpnSettingsUiState.createDefault(
+                        isCustomDnsEnabled = true,
+                        customDnsItems =
+                            listOf(CustomDnsItem(address = DUMMY_DNS_ADDRESS, isLocal = false)),
+                    )
+            )
 
             // Assert
             onNodeWithContentDescription(LOCAL_DNS_SERVER_WARNING).assertDoesNotExist()
@@ -186,16 +239,14 @@ class VpnSettingsScreenTest {
     fun testLanWarningShowedWhenAllowLanEnabledAndLocalDnsAddressIsUsed() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state =
-                        VpnSettingsUiState.createDefault(
-                            isCustomDnsEnabled = true,
-                            customDnsItems =
-                                listOf(CustomDnsItem(address = DUMMY_DNS_ADDRESS, isLocal = true)),
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    VpnSettingsUiState.createDefault(
+                        isCustomDnsEnabled = true,
+                        customDnsItems =
+                            listOf(CustomDnsItem(address = DUMMY_DNS_ADDRESS, isLocal = true)),
+                    )
+            )
 
             // Assert
             onNodeWithContentDescription(LOCAL_DNS_SERVER_WARNING).assertExists()
@@ -205,14 +256,10 @@ class VpnSettingsScreenTest {
     fun testShowSelectedTunnelQuantumOption() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state =
-                        VpnSettingsUiState.createDefault(
-                            quantumResistant = QuantumResistantState.On
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    VpnSettingsUiState.createDefault(quantumResistant = QuantumResistantState.On)
+            )
             onNodeWithTag(LAZY_LIST_VPN_SETTINGS_TEST_TAG)
                 .performScrollToNode(hasTestTag(LAZY_LIST_QUANTUM_ITEM_OFF_TEST_TAG))
 
@@ -227,15 +274,11 @@ class VpnSettingsScreenTest {
             // Arrange
             val mockSelectQuantumResistantSettingListener: (QuantumResistantState) -> Unit =
                 mockk(relaxed = true)
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state =
-                        VpnSettingsUiState.createDefault(
-                            quantumResistant = QuantumResistantState.Auto
-                        ),
-                    onSelectQuantumResistanceSetting = mockSelectQuantumResistantSettingListener,
-                )
-            }
+            initScreen(
+                state =
+                    VpnSettingsUiState.createDefault(quantumResistant = QuantumResistantState.Auto),
+                onSelectQuantumResistanceSetting = mockSelectQuantumResistantSettingListener,
+            )
             onNodeWithTag(LAZY_LIST_VPN_SETTINGS_TEST_TAG)
                 .performScrollToNode(hasTestTag(LAZY_LIST_QUANTUM_ITEM_OFF_TEST_TAG))
 
@@ -251,14 +294,12 @@ class VpnSettingsScreenTest {
     fun testShowWireguardPortOptions() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state =
-                        VpnSettingsUiState.createDefault(
-                            selectedWireguardPort = Constraint.Only(Port(53))
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    VpnSettingsUiState.createDefault(
+                        selectedWireguardPort = Constraint.Only(Port(53))
+                    )
+            )
 
             // Act
             onNodeWithTag(LAZY_LIST_VPN_SETTINGS_TEST_TAG)
@@ -280,15 +321,13 @@ class VpnSettingsScreenTest {
             // Arrange
             val mockSelectWireguardPortSelectionListener: (Constraint<Port>) -> Unit =
                 mockk(relaxed = true)
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state =
-                        VpnSettingsUiState.createDefault(
-                            selectedWireguardPort = Constraint.Only(Port(53))
-                        ),
-                    onWireguardPortSelected = mockSelectWireguardPortSelectionListener,
-                )
-            }
+            initScreen(
+                state =
+                    VpnSettingsUiState.createDefault(
+                        selectedWireguardPort = Constraint.Only(Port(53))
+                    ),
+                onWireguardPortSelected = mockSelectWireguardPortSelectionListener,
+            )
 
             // Act
             onNodeWithTag(LAZY_LIST_VPN_SETTINGS_TEST_TAG)
@@ -311,11 +350,7 @@ class VpnSettingsScreenTest {
     fun testShowWireguardCustomPort() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state = VpnSettingsUiState.createDefault(customWireguardPort = Port(4000))
-                )
-            }
+            initScreen(state = VpnSettingsUiState.createDefault(customWireguardPort = Port(4000)))
 
             // Act
             onNodeWithTag(LAZY_LIST_VPN_SETTINGS_TEST_TAG)
@@ -330,16 +365,14 @@ class VpnSettingsScreenTest {
         composeExtension.use {
             // Arrange
             val onWireguardPortSelected: (Constraint<Port>) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state =
-                        VpnSettingsUiState.createDefault(
-                            selectedWireguardPort = Constraint.Only(Port(4000)),
-                            customWireguardPort = Port(4000),
-                        ),
-                    onWireguardPortSelected = onWireguardPortSelected,
-                )
-            }
+            initScreen(
+                state =
+                    VpnSettingsUiState.createDefault(
+                        selectedWireguardPort = Constraint.Only(Port(4000)),
+                        customWireguardPort = Port(4000),
+                    ),
+                onWireguardPortSelected = onWireguardPortSelected,
+            )
 
             // Act
             onNodeWithTag(LAZY_LIST_VPN_SETTINGS_TEST_TAG)
@@ -357,12 +390,10 @@ class VpnSettingsScreenTest {
         composeExtension.use {
             // Arrange
             val mockedClickHandler: (Mtu?) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state = VpnSettingsUiState.createDefault(),
-                    navigateToMtuDialog = mockedClickHandler,
-                )
-            }
+            initScreen(
+                state = VpnSettingsUiState.createDefault(),
+                navigateToMtuDialog = mockedClickHandler,
+            )
 
             onNodeWithTag(LAZY_LIST_VPN_SETTINGS_TEST_TAG)
                 .performScrollToNode(hasTestTag(LAZY_LIST_LAST_ITEM_TEST_TAG))
@@ -379,12 +410,10 @@ class VpnSettingsScreenTest {
         composeExtension.use {
             // Arrange
             val mockedClickHandler: (Int?, String?) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state = VpnSettingsUiState.createDefault(isCustomDnsEnabled = true),
-                    navigateToDns = mockedClickHandler,
-                )
-            }
+            initScreen(
+                state = VpnSettingsUiState.createDefault(isCustomDnsEnabled = true),
+                navigateToDns = mockedClickHandler,
+            )
 
             // Act
             onNodeWithText("Add a server").performClick()
@@ -399,12 +428,10 @@ class VpnSettingsScreenTest {
             val mockedNavigateToObfuscationInfo: () -> Unit = mockk(relaxed = true)
 
             // Arrange
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state = VpnSettingsUiState.createDefault(),
-                    navigateToObfuscationInfo = mockedNavigateToObfuscationInfo,
-                )
-            }
+            initScreen(
+                state = VpnSettingsUiState.createDefault(),
+                navigateToObfuscationInfo = mockedNavigateToObfuscationInfo,
+            )
 
             // Act
             onNodeWithTag(LAZY_LIST_VPN_SETTINGS_TEST_TAG)
@@ -421,15 +448,12 @@ class VpnSettingsScreenTest {
             val mockedShowTunnelQuantumInfoClick: () -> Unit = mockk(relaxed = true)
 
             // Arrange
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state = VpnSettingsUiState.createDefault(),
-                    navigateToQuantumResistanceInfo = mockedShowTunnelQuantumInfoClick,
-                )
-            }
+            initScreen(
+                state = VpnSettingsUiState.createDefault(),
+                navigateToQuantumResistanceInfo = mockedShowTunnelQuantumInfoClick,
+            )
 
             // Act
-
             onNodeWithTag(LAZY_LIST_VPN_SETTINGS_TEST_TAG)
                 .performScrollToNode(hasTestTag(LAZY_LIST_QUANTUM_ITEM_ON_TEST_TAG))
             onNodeWithText("Quantum-resistant tunnel").performClick()
@@ -444,12 +468,10 @@ class VpnSettingsScreenTest {
             val mockedClickHandler: (List<PortRange>) -> Unit = mockk(relaxed = true)
 
             // Arrange
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state = VpnSettingsUiState.createDefault(),
-                    navigateToWireguardPortInfo = mockedClickHandler,
-                )
-            }
+            initScreen(
+                state = VpnSettingsUiState.createDefault(),
+                navigateToWireguardPortInfo = mockedClickHandler,
+            )
 
             onNodeWithText("WireGuard port").performClick()
 
@@ -462,12 +484,10 @@ class VpnSettingsScreenTest {
             val mockedClickHandler: () -> Unit = mockk(relaxed = true)
 
             // Arrange
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state = VpnSettingsUiState.createDefault(),
-                    navigateToWireguardPortDialog = mockedClickHandler,
-                )
-            }
+            initScreen(
+                state = VpnSettingsUiState.createDefault(),
+                navigateToWireguardPortDialog = mockedClickHandler,
+            )
 
             onNodeWithTag(LAZY_LIST_VPN_SETTINGS_TEST_TAG)
                 .performScrollToNode(hasTestTag(LAZY_LIST_WIREGUARD_CUSTOM_PORT_NUMBER_TEST_TAG))
@@ -482,12 +502,10 @@ class VpnSettingsScreenTest {
         composeExtension.use {
             // Arrange
             val mockOnShowCustomPortDialog: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state = VpnSettingsUiState.createDefault(),
-                    navigateToWireguardPortDialog = mockOnShowCustomPortDialog,
-                )
-            }
+            initScreen(
+                state = VpnSettingsUiState.createDefault(),
+                navigateToWireguardPortDialog = mockOnShowCustomPortDialog,
+            )
 
             // Act
             onNodeWithTag(LAZY_LIST_VPN_SETTINGS_TEST_TAG)
@@ -503,15 +521,13 @@ class VpnSettingsScreenTest {
         composeExtension.use {
             // Arrange
             val mockOnShowCustomPortDialog: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state =
-                        VpnSettingsUiState.createDefault(
-                            selectedWireguardPort = Constraint.Only(Port(4000))
-                        ),
-                    navigateToWireguardPortDialog = mockOnShowCustomPortDialog,
-                )
-            }
+            initScreen(
+                state =
+                    VpnSettingsUiState.createDefault(
+                        selectedWireguardPort = Constraint.Only(Port(4000))
+                    ),
+                navigateToWireguardPortDialog = mockOnShowCustomPortDialog,
+            )
 
             // Act
             onNodeWithTag(LAZY_LIST_VPN_SETTINGS_TEST_TAG)
@@ -526,11 +542,7 @@ class VpnSettingsScreenTest {
     fun ensureConnectOnStartIsShownWhenSystemVpnSettingsAvailableIsFalse() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state = VpnSettingsUiState.createDefault(systemVpnSettingsAvailable = false)
-                )
-            }
+            initScreen(state = VpnSettingsUiState.createDefault(systemVpnSettingsAvailable = false))
 
             // Assert
             onNodeWithText("Connect on device start-up").assertExists()
@@ -541,16 +553,14 @@ class VpnSettingsScreenTest {
         composeExtension.use {
             // Arrange
             val mockOnToggleAutoStartAndConnectOnBoot: (Boolean) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                VpnSettingsScreen(
-                    state =
-                        VpnSettingsUiState.createDefault(
-                            systemVpnSettingsAvailable = false,
-                            autoStartAndConnectOnBoot = false,
-                        ),
-                    onToggleAutoStartAndConnectOnBoot = mockOnToggleAutoStartAndConnectOnBoot,
-                )
-            }
+            initScreen(
+                state =
+                    VpnSettingsUiState.createDefault(
+                        systemVpnSettingsAvailable = false,
+                        autoStartAndConnectOnBoot = false,
+                    ),
+                onToggleAutoStartAndConnectOnBoot = mockOnToggleAutoStartAndConnectOnBoot,
+            )
 
             // Act
             onNodeWithText("Connect on device start-up").performClick()
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 a9f1658308..8e9cc3fdf8 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
@@ -4,6 +4,7 @@ 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.ComposeContext
 import io.mockk.MockKAnnotations
 import io.mockk.every
 import io.mockk.mockk
@@ -32,23 +33,37 @@ class WelcomeScreenTest {
         MockKAnnotations.init(this)
     }
 
+    private fun ComposeContext.initScreen(
+        state: WelcomeUiState = WelcomeUiState(),
+        onSitePaymentClick: () -> Unit = {},
+        onRedeemVoucherClick: () -> Unit = {},
+        onSettingsClick: () -> Unit = {},
+        onAccountClick: () -> Unit = {},
+        onPurchaseBillingProductClick: (productId: ProductId) -> Unit = {},
+        onDisconnectClick: () -> Unit = {},
+        navigateToDeviceInfoDialog: () -> Unit = {},
+        navigateToVerificationPendingDialog: () -> Unit = {},
+    ) {
+        setContentWithTheme {
+            WelcomeScreen(
+                state = state,
+                onSitePaymentClick = onSitePaymentClick,
+                onRedeemVoucherClick = onRedeemVoucherClick,
+                onSettingsClick = onSettingsClick,
+                onAccountClick = onAccountClick,
+                onPurchaseBillingProductClick = onPurchaseBillingProductClick,
+                navigateToDeviceInfoDialog = navigateToDeviceInfoDialog,
+                navigateToVerificationPendingDialog = navigateToVerificationPendingDialog,
+                onDisconnectClick = onDisconnectClick,
+            )
+        }
+    }
+
     @Test
     fun testDefaultState() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                WelcomeScreen(
-                    state = WelcomeUiState(),
-                    onSitePaymentClick = {},
-                    onRedeemVoucherClick = {},
-                    onSettingsClick = {},
-                    onAccountClick = {},
-                    onPurchaseBillingProductClick = { _ -> },
-                    navigateToDeviceInfoDialog = {},
-                    navigateToVerificationPendingDialog = {},
-                    onDisconnectClick = {},
-                )
-            }
+            initScreen()
 
             // Assert
             onNodeWithText("Congrats!").assertExists()
@@ -59,19 +74,7 @@ class WelcomeScreenTest {
     fun testDisableSitePayment() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                WelcomeScreen(
-                    state = WelcomeUiState(),
-                    onSitePaymentClick = {},
-                    onRedeemVoucherClick = {},
-                    onSettingsClick = {},
-                    onAccountClick = {},
-                    onPurchaseBillingProductClick = { _ -> },
-                    navigateToDeviceInfoDialog = {},
-                    navigateToVerificationPendingDialog = {},
-                    onDisconnectClick = {},
-                )
-            }
+            initScreen()
 
             // Assert
             onNodeWithText(
@@ -88,19 +91,7 @@ class WelcomeScreenTest {
             // Arrange
             val rawAccountNumber = AccountNumber("1111222233334444")
             val expectedAccountNumber = "1111 2222 3333 4444"
-            setContentWithTheme {
-                WelcomeScreen(
-                    state = WelcomeUiState(accountNumber = rawAccountNumber),
-                    onSitePaymentClick = {},
-                    onRedeemVoucherClick = {},
-                    onSettingsClick = {},
-                    onAccountClick = {},
-                    onPurchaseBillingProductClick = { _ -> },
-                    navigateToDeviceInfoDialog = {},
-                    navigateToVerificationPendingDialog = {},
-                    onDisconnectClick = {},
-                )
-            }
+            initScreen(state = WelcomeUiState(accountNumber = rawAccountNumber))
 
             // Assert
             onNodeWithText(expectedAccountNumber).assertExists()
@@ -111,19 +102,10 @@ class WelcomeScreenTest {
         composeExtension.use {
             // Arrange
             val mockClickListener: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                WelcomeScreen(
-                    state = WelcomeUiState(showSitePayment = true),
-                    onSitePaymentClick = mockClickListener,
-                    onRedeemVoucherClick = {},
-                    onSettingsClick = {},
-                    onAccountClick = {},
-                    onPurchaseBillingProductClick = { _ -> },
-                    navigateToDeviceInfoDialog = {},
-                    navigateToVerificationPendingDialog = {},
-                    onDisconnectClick = {},
-                )
-            }
+            initScreen(
+                state = WelcomeUiState(showSitePayment = true),
+                onSitePaymentClick = mockClickListener,
+            )
 
             // Act
             onNodeWithText("Buy credit").performClick()
@@ -137,19 +119,7 @@ class WelcomeScreenTest {
         composeExtension.use {
             // Arrange
             val mockClickListener: () -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                WelcomeScreen(
-                    state = WelcomeUiState(),
-                    onSitePaymentClick = {},
-                    onRedeemVoucherClick = mockClickListener,
-                    onSettingsClick = {},
-                    onAccountClick = {},
-                    onPurchaseBillingProductClick = { _ -> },
-                    navigateToDeviceInfoDialog = {},
-                    navigateToVerificationPendingDialog = {},
-                    onDisconnectClick = {},
-                )
-            }
+            initScreen(state = WelcomeUiState(), onRedeemVoucherClick = mockClickListener)
 
             // Act
             onNodeWithText("Redeem voucher").performClick()
@@ -162,19 +132,9 @@ class WelcomeScreenTest {
     fun testShowBillingErrorPaymentButton() =
         composeExtension.use {
             // Arrange
-            setContentWithTheme {
-                WelcomeScreen(
-                    state = WelcomeUiState().copy(billingPaymentState = PaymentState.Error.Billing),
-                    onSitePaymentClick = {},
-                    onRedeemVoucherClick = {},
-                    onSettingsClick = {},
-                    onAccountClick = {},
-                    onPurchaseBillingProductClick = { _ -> },
-                    navigateToDeviceInfoDialog = {},
-                    navigateToVerificationPendingDialog = {},
-                    onDisconnectClick = {},
-                )
-            }
+            initScreen(
+                state = WelcomeUiState().copy(billingPaymentState = PaymentState.Error.Billing)
+            )
 
             // Assert
             onNodeWithText("Add 30 days time").assertExists()
@@ -187,23 +147,13 @@ class WelcomeScreenTest {
             val mockPaymentProduct: PaymentProduct = mockk()
             every { mockPaymentProduct.price } returns ProductPrice("$10")
             every { mockPaymentProduct.status } returns null
-            setContentWithTheme {
-                WelcomeScreen(
-                    state =
-                        WelcomeUiState(
-                            billingPaymentState =
-                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                        ),
-                    onSitePaymentClick = {},
-                    onRedeemVoucherClick = {},
-                    onSettingsClick = {},
-                    onAccountClick = {},
-                    onPurchaseBillingProductClick = { _ -> },
-                    navigateToDeviceInfoDialog = {},
-                    navigateToVerificationPendingDialog = {},
-                    onDisconnectClick = {},
-                )
-            }
+            initScreen(
+                state =
+                    WelcomeUiState(
+                        billingPaymentState =
+                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+                    )
+            )
 
             // Assert
             onNodeWithText("Add 30 days time ($10)").assertExists()
@@ -216,24 +166,13 @@ class WelcomeScreenTest {
             val mockPaymentProduct: PaymentProduct = mockk()
             every { mockPaymentProduct.price } returns ProductPrice("$10")
             every { mockPaymentProduct.status } returns PaymentStatus.PENDING
-            setContentWithTheme {
-                WelcomeScreen(
-                    state =
-                        WelcomeUiState()
-                            .copy(
-                                billingPaymentState =
-                                    PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                            ),
-                    onSitePaymentClick = {},
-                    onRedeemVoucherClick = {},
-                    onSettingsClick = {},
-                    onAccountClick = {},
-                    onPurchaseBillingProductClick = { _ -> },
-                    navigateToDeviceInfoDialog = {},
-                    navigateToVerificationPendingDialog = {},
-                    onDisconnectClick = {},
-                )
-            }
+            initScreen(
+                state =
+                    WelcomeUiState(
+                        billingPaymentState =
+                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+                    )
+            )
 
             // Assert
             onNodeWithText("Google Play payment pending").assertExists()
@@ -247,24 +186,14 @@ class WelcomeScreenTest {
             every { mockPaymentProduct.price } returns ProductPrice("$10")
             every { mockPaymentProduct.status } returns PaymentStatus.PENDING
             val mockShowPendingInfo = mockk<() -> Unit>(relaxed = true)
-            setContentWithTheme {
-                WelcomeScreen(
-                    state =
-                        WelcomeUiState()
-                            .copy(
-                                billingPaymentState =
-                                    PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                            ),
-                    onSitePaymentClick = {},
-                    onRedeemVoucherClick = {},
-                    onSettingsClick = {},
-                    onAccountClick = {},
-                    onPurchaseBillingProductClick = { _ -> },
-                    navigateToVerificationPendingDialog = mockShowPendingInfo,
-                    navigateToDeviceInfoDialog = {},
-                    onDisconnectClick = {},
-                )
-            }
+            initScreen(
+                state =
+                    WelcomeUiState(
+                        billingPaymentState =
+                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+                    ),
+                navigateToVerificationPendingDialog = mockShowPendingInfo,
+            )
 
             // Act
             onNodeWithTag(PLAY_PAYMENT_INFO_ICON_TEST_TAG).performClick()
@@ -280,24 +209,14 @@ class WelcomeScreenTest {
             val mockPaymentProduct: PaymentProduct = mockk()
             every { mockPaymentProduct.price } returns ProductPrice("$10")
             every { mockPaymentProduct.status } returns PaymentStatus.VERIFICATION_IN_PROGRESS
-            setContentWithTheme {
-                WelcomeScreen(
-                    state =
-                        WelcomeUiState()
-                            .copy(
-                                billingPaymentState =
-                                    PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                            ),
-                    onSitePaymentClick = {},
-                    onRedeemVoucherClick = {},
-                    onSettingsClick = {},
-                    onAccountClick = {},
-                    onPurchaseBillingProductClick = { _ -> },
-                    navigateToDeviceInfoDialog = {},
-                    navigateToVerificationPendingDialog = {},
-                    onDisconnectClick = {},
-                )
-            }
+
+            initScreen(
+                state =
+                    WelcomeUiState(
+                        billingPaymentState =
+                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+                    )
+            )
 
             // Assert
             onNodeWithText("Verifying purchase").assertExists()
@@ -312,23 +231,14 @@ class WelcomeScreenTest {
             every { mockPaymentProduct.price } returns ProductPrice("$10")
             every { mockPaymentProduct.productId } returns ProductId("PRODUCT_ID")
             every { mockPaymentProduct.status } returns null
-            setContentWithTheme {
-                WelcomeScreen(
-                    state =
-                        WelcomeUiState(
-                            billingPaymentState =
-                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                        ),
-                    onSitePaymentClick = {},
-                    onRedeemVoucherClick = {},
-                    onSettingsClick = {},
-                    onAccountClick = {},
-                    onPurchaseBillingProductClick = clickHandler,
-                    navigateToDeviceInfoDialog = {},
-                    navigateToVerificationPendingDialog = {},
-                    onDisconnectClick = {},
-                )
-            }
+            initScreen(
+                state =
+                    WelcomeUiState(
+                        billingPaymentState =
+                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+                    ),
+                onPurchaseBillingProductClick = clickHandler,
+            )
 
             // Act
             onNodeWithText("Add 30 days time ($10)").performClick()
@@ -344,19 +254,10 @@ class WelcomeScreenTest {
             val clickHandler: () -> Unit = mockk(relaxed = true)
             val tunnelState: TunnelState = mockk(relaxed = true)
             every { tunnelState.isSecured() } returns true
-            setContentWithTheme {
-                WelcomeScreen(
-                    state = WelcomeUiState(tunnelState = tunnelState),
-                    onSitePaymentClick = {},
-                    onRedeemVoucherClick = {},
-                    onSettingsClick = {},
-                    onAccountClick = {},
-                    onPurchaseBillingProductClick = {},
-                    navigateToDeviceInfoDialog = {},
-                    navigateToVerificationPendingDialog = {},
-                    onDisconnectClick = clickHandler,
-                )
-            }
+            initScreen(
+                state = WelcomeUiState(tunnelState = tunnelState),
+                onDisconnectClick = clickHandler,
+            )
 
             // Act
             onNodeWithText("Disconnect").performClick()
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/location/SearchLocationScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/location/SearchLocationScreenTest.kt
index 5901599df9..8e87c03e13 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/location/SearchLocationScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/location/SearchLocationScreenTest.kt
@@ -4,6 +4,7 @@ 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.ComposeContext
 import io.mockk.MockKAnnotations
 import io.mockk.mockk
 import io.mockk.unmockkAll
@@ -14,6 +15,9 @@ import net.mullvad.mullvadvpn.compose.setContentWithTheme
 import net.mullvad.mullvadvpn.compose.state.RelayListItem
 import net.mullvad.mullvadvpn.compose.state.SearchLocationUiState
 import net.mullvad.mullvadvpn.compose.test.SELECT_LOCATION_CUSTOM_LIST_HEADER_TEST_TAG
+import net.mullvad.mullvadvpn.lib.model.CustomListId
+import net.mullvad.mullvadvpn.lib.model.RelayItem
+import net.mullvad.mullvadvpn.lib.model.RelayItemId
 import org.junit.jupiter.api.AfterEach
 import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
@@ -33,18 +37,57 @@ class SearchLocationScreenTest {
         unmockkAll()
     }
 
+    private fun ComposeContext.initScreen(
+        state: SearchLocationUiState,
+        onSelectRelay: (RelayItem) -> Unit = {},
+        onToggleExpand: (RelayItemId, CustomListId?, Boolean) -> Unit = { _, _, _ -> },
+        onSearchInputChanged: (String) -> Unit = {},
+        onCreateCustomList: (location: RelayItem.Location?) -> Unit = {},
+        onEditCustomLists: () -> Unit = {},
+        onAddLocationToList:
+            (location: RelayItem.Location, customList: RelayItem.CustomList) -> Unit =
+            { _, _ ->
+            },
+        onRemoveLocationFromList:
+            (location: RelayItem.Location, customListId: CustomListId) -> Unit =
+            { _, _ ->
+            },
+        onEditCustomListName: (RelayItem.CustomList) -> Unit = {},
+        onEditLocationsCustomList: (RelayItem.CustomList) -> Unit = {},
+        onDeleteCustomList: (RelayItem.CustomList) -> Unit = {},
+        onRemoveOwnershipFilter: () -> Unit = {},
+        onRemoveProviderFilter: () -> Unit = {},
+        onGoBack: () -> Unit = {},
+    ) {
+        setContentWithTheme {
+            SearchLocationScreen(
+                state = state,
+                onSelectRelay = onSelectRelay,
+                onToggleExpand = onToggleExpand,
+                onSearchInputChanged = onSearchInputChanged,
+                onCreateCustomList = onCreateCustomList,
+                onEditCustomLists = onEditCustomLists,
+                onAddLocationToList = onAddLocationToList,
+                onRemoveLocationFromList = onRemoveLocationFromList,
+                onEditCustomListName = onEditCustomListName,
+                onEditLocationsCustomList = onEditLocationsCustomList,
+                onDeleteCustomList = onDeleteCustomList,
+                onRemoveOwnershipFilter = onRemoveOwnershipFilter,
+                onRemoveProviderFilter = onRemoveProviderFilter,
+                onGoBack = onGoBack,
+            )
+        }
+    }
+
     @Test
     fun testSearchInput() =
         composeExtension.use {
             // Arrange
             val mockedSearchTermInput: (String) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                SearchLocationScreen(
-                    state =
-                        SearchLocationUiState.NoQuery(searchTerm = "", filterChips = emptyList()),
-                    onSearchInputChanged = mockedSearchTermInput,
-                )
-            }
+            initScreen(
+                state = SearchLocationUiState.NoQuery(searchTerm = "", filterChips = emptyList()),
+                onSearchInputChanged = mockedSearchTermInput,
+            )
             val mockSearchString = "SEARCH"
 
             // Act
@@ -59,18 +102,15 @@ class SearchLocationScreenTest {
         composeExtension.use {
             // Arrange
             val mockSearchString = "SEARCH"
-            setContentWithTheme {
-                SearchLocationScreen(
-                    state =
-                        SearchLocationUiState.Content(
-                            searchTerm = mockSearchString,
-                            filterChips = emptyList(),
-                            relayListItems =
-                                listOf(RelayListItem.LocationsEmptyText(mockSearchString)),
-                            customLists = emptyList(),
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    SearchLocationUiState.Content(
+                        searchTerm = mockSearchString,
+                        filterChips = emptyList(),
+                        relayListItems = listOf(RelayListItem.LocationsEmptyText(mockSearchString)),
+                        customLists = emptyList(),
+                    )
+            )
 
             // Assert
             onNodeWithText("No result for \"$mockSearchString\", please try a different search")
@@ -82,17 +122,15 @@ class SearchLocationScreenTest {
         composeExtension.use {
             // Arrange
             val mockSearchString = "SEARCH"
-            setContentWithTheme {
-                SearchLocationScreen(
-                    state =
-                        SearchLocationUiState.Content(
-                            searchTerm = mockSearchString,
-                            filterChips = emptyList(),
-                            relayListItems = emptyList(),
-                            customLists = DUMMY_RELAY_ITEM_CUSTOM_LISTS,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    SearchLocationUiState.Content(
+                        searchTerm = mockSearchString,
+                        filterChips = emptyList(),
+                        relayListItems = emptyList(),
+                        customLists = DUMMY_RELAY_ITEM_CUSTOM_LISTS,
+                    )
+            )
 
             // Assert
             onNodeWithText(CUSTOM_LISTS_EMPTY_TEXT).assertDoesNotExist()
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/location/SelectLocationScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/location/SelectLocationScreenTest.kt
index cf3380e97a..d0dbb9155d 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/location/SelectLocationScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/location/SelectLocationScreenTest.kt
@@ -4,6 +4,7 @@ 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.ComposeContext
 import io.mockk.MockKAnnotations
 import io.mockk.every
 import io.mockk.mockk
@@ -20,6 +21,7 @@ import net.mullvad.mullvadvpn.compose.state.SelectLocationListUiState
 import net.mullvad.mullvadvpn.compose.state.SelectLocationUiState
 import net.mullvad.mullvadvpn.compose.test.SELECT_LOCATION_CUSTOM_LIST_BOTTOM_SHEET_TEST_TAG
 import net.mullvad.mullvadvpn.compose.test.SELECT_LOCATION_LOCATION_BOTTOM_SHEET_TEST_TAG
+import net.mullvad.mullvadvpn.lib.model.CustomListId
 import net.mullvad.mullvadvpn.lib.model.RelayItem
 import net.mullvad.mullvadvpn.performLongClick
 import net.mullvad.mullvadvpn.viewmodel.location.SelectLocationListViewModel
@@ -49,6 +51,53 @@ class SelectLocationScreenTest {
         unmockkAll()
     }
 
+    private fun ComposeContext.initScreen(
+        state: SelectLocationUiState = SelectLocationUiState.Loading,
+        onSelectRelay: (item: RelayItem) -> Unit = {},
+        onSearchClick: (RelayListType) -> Unit = {},
+        onBackClick: () -> Unit = {},
+        onFilterClick: () -> Unit = {},
+        onCreateCustomList: (location: RelayItem.Location?) -> Unit = {},
+        onEditCustomLists: () -> Unit = {},
+        removeOwnershipFilter: () -> Unit = {},
+        removeProviderFilter: () -> Unit = {},
+        onAddLocationToList:
+            (location: RelayItem.Location, customList: RelayItem.CustomList) -> Unit =
+            { _, _ ->
+            },
+        onRemoveLocationFromList:
+            (location: RelayItem.Location, customListId: CustomListId) -> Unit =
+            { _, _ ->
+            },
+        onEditCustomListName: (RelayItem.CustomList) -> Unit = {},
+        onEditLocationsCustomList: (RelayItem.CustomList) -> Unit = {},
+        onDeleteCustomList: (RelayItem.CustomList) -> Unit = {},
+        onSelectRelayList: (RelayListType) -> Unit = {},
+        openDaitaSettings: () -> Unit = {},
+    ) {
+
+        setContentWithTheme {
+            SelectLocationScreen(
+                state = state,
+                onSelectRelay = onSelectRelay,
+                onSearchClick = onSearchClick,
+                onBackClick = onBackClick,
+                onFilterClick = onFilterClick,
+                onCreateCustomList = onCreateCustomList,
+                onEditCustomLists = onEditCustomLists,
+                removeOwnershipFilter = removeOwnershipFilter,
+                removeProviderFilter = removeProviderFilter,
+                onAddLocationToList = onAddLocationToList,
+                onRemoveLocationFromList = onRemoveLocationFromList,
+                onEditCustomListName = onEditCustomListName,
+                onEditLocationsCustomList = onEditLocationsCustomList,
+                onDeleteCustomList = onDeleteCustomList,
+                onSelectRelayList = onSelectRelayList,
+                openDaitaSettings = openDaitaSettings,
+            )
+        }
+    }
+
     @Test
     fun testShowRelayListState() =
         composeExtension.use {
@@ -61,17 +110,15 @@ class SelectLocationScreenTest {
                         customLists = emptyList(),
                     )
                 )
-            setContentWithTheme {
-                SelectLocationScreen(
-                    state =
-                        SelectLocationUiState.Data(
-                            // searchTerm = "",
-                            filterChips = emptyList(),
-                            multihopEnabled = false,
-                            relayListType = RelayListType.EXIT,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    SelectLocationUiState.Data(
+                        // searchTerm = "",
+                        filterChips = emptyList(),
+                        multihopEnabled = false,
+                        relayListType = RelayListType.EXIT,
+                    )
+            )
 
             // Assert
             onNodeWithText("Relay Country 1").assertExists()
@@ -93,16 +140,14 @@ class SelectLocationScreenTest {
                         customLists = emptyList(),
                     )
                 )
-            setContentWithTheme {
-                SelectLocationScreen(
-                    state =
-                        SelectLocationUiState.Data(
-                            filterChips = emptyList(),
-                            multihopEnabled = false,
-                            relayListType = RelayListType.EXIT,
-                        )
-                )
-            }
+            initScreen(
+                state =
+                    SelectLocationUiState.Data(
+                        filterChips = emptyList(),
+                        multihopEnabled = false,
+                        relayListType = RelayListType.EXIT,
+                    )
+            )
 
             // Assert
             onNodeWithText(CUSTOM_LISTS_EMPTY_TEXT).assertExists()
@@ -121,17 +166,15 @@ class SelectLocationScreenTest {
                     )
                 )
             val mockedOnSelectRelay: (RelayItem) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                SelectLocationScreen(
-                    state =
-                        SelectLocationUiState.Data(
-                            filterChips = emptyList(),
-                            multihopEnabled = false,
-                            relayListType = RelayListType.EXIT,
-                        ),
-                    onSelectRelay = mockedOnSelectRelay,
-                )
-            }
+            initScreen(
+                state =
+                    SelectLocationUiState.Data(
+                        filterChips = emptyList(),
+                        multihopEnabled = false,
+                        relayListType = RelayListType.EXIT,
+                    ),
+                onSelectRelay = mockedOnSelectRelay,
+            )
 
             // Act
             onNodeWithText(customList.name).performClick()
@@ -153,18 +196,16 @@ class SelectLocationScreenTest {
                     )
                 )
             val mockedOnSelectRelay: (RelayItem) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                SelectLocationScreen(
-                    state =
-                        SelectLocationUiState.Data(
-                            // searchTerm = "",
-                            filterChips = emptyList(),
-                            multihopEnabled = false,
-                            relayListType = RelayListType.EXIT,
-                        ),
-                    onSelectRelay = mockedOnSelectRelay,
-                )
-            }
+            initScreen(
+                state =
+                    SelectLocationUiState.Data(
+                        // searchTerm = "",
+                        filterChips = emptyList(),
+                        multihopEnabled = false,
+                        relayListType = RelayListType.EXIT,
+                    ),
+                onSelectRelay = mockedOnSelectRelay,
+            )
 
             // Act
             onNodeWithText(customList.name).performLongClick()
@@ -186,17 +227,15 @@ class SelectLocationScreenTest {
                     )
                 )
             val mockedOnSelectRelay: (RelayItem) -> Unit = mockk(relaxed = true)
-            setContentWithTheme {
-                SelectLocationScreen(
-                    state =
-                        SelectLocationUiState.Data(
-                            filterChips = emptyList(),
-                            multihopEnabled = false,
-                            relayListType = RelayListType.EXIT,
-                        ),
-                    onSelectRelay = mockedOnSelectRelay,
-                )
-            }
+            initScreen(
+                state =
+                    SelectLocationUiState.Data(
+                        filterChips = emptyList(),
+                        multihopEnabled = false,
+                        relayListType = RelayListType.EXIT,
+                    ),
+                onSelectRelay = mockedOnSelectRelay,
+            )
 
             // Act
             onNodeWithText(relayItem.name).performLongClick()