Update tests

1 files changed, 4 insertions, 2 deletions

diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/VpnSettingsViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/VpnSettingsViewModelTest.kt
index 4e02f1f703..ef872d8ae4 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/VpnSettingsViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/VpnSettingsViewModelTest.kt
@@ -14,6 +14,7 @@ import kotlinx.coroutines.cancel
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.runTest
+import net.mullvad.mullvadvpn.compose.state.VpnSettingsDialog
 import net.mullvad.mullvadvpn.compose.state.VpnSettingsUiState
 import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
 import net.mullvad.mullvadvpn.lib.common.test.assertLists
@@ -37,6 +38,7 @@ import org.junit.After
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
+import kotlin.test.assertTrue
 
 class VpnSettingsViewModelTest {
     @get:Rule val testCoroutineRule = TestCoroutineRule()
@@ -180,14 +182,14 @@ class VpnSettingsViewModelTest {
 
         // Act, Assert
         viewModel.uiState.test {
-            assertIs<VpnSettingsUiState.DefaultUiState>(awaitItem())
+            assertIs<VpnSettingsUiState>(awaitItem())
             mockSettingsUpdate.value = mockSettings
             viewModel.onWireguardPortInfoClicked()
             mockConnectionState.value =
                 ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
             portRangeSlot.captured.invoke(expectedPortRange)
             val state = awaitItem()
-            assertIs<VpnSettingsUiState.WireguardPortInfoDialogUiState>(state)
+            assertTrue{ state.dialog is VpnSettingsDialog.WireguardPortInfo }
             assertLists(expectedPortRange, state.availablePortRanges)
         }
     }