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 0e8fc14441f2..7ca4d7f3d4b0 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
@@ -21,7 +21,6 @@ import net.mullvad.mullvadvpn.compose.setContentWithTheme
 import net.mullvad.mullvadvpn.compose.test.CUSTOM_PORT_DIALOG_INPUT_TEST_TAG
 import net.mullvad.mullvadvpn.lib.model.PortRange
 import net.mullvad.mullvadvpn.onNodeWithTagAndText
-import net.mullvad.mullvadvpn.viewmodel.WireguardCustomPortDialogUiState
 import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.extension.RegisterExtension
@@ -50,17 +49,12 @@ class CustomPortDialogTest {
         onResetPort: () -> Unit = {},
         onDismiss: () -> Unit = {},
     ) {
-        val state =
-            WireguardCustomPortDialogUiState(
-                portInput = portInput,
-                isValidInput = isValidInput,
-                allowedPortRanges = allowedPortRanges,
-                showResetToDefault = showResetToDefault,
-            )
-
         CustomPortDialog(
-            state,
             title = title,
+            portInput = portInput,
+            isValidInput = isValidInput,
+            showResetToDefault = showResetToDefault,
+            allowedPortRanges = allowedPortRanges,
             onInputChanged = onInputChanged,
             onSavePort = onSavePort,
             onDismiss = onDismiss,