Add and fix unit tests

6 files changed, 164 insertions, 19 deletions

diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/FilterChipUseCaseTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/FilterChipUseCaseTest.kt
index 8b3d6d68a2..221d89cf40 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/FilterChipUseCaseTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/FilterChipUseCaseTest.kt
@@ -108,10 +108,14 @@ class FilterChipUseCaseTest {
         }
 
     @Test
-    fun `when Daita is enabled and multihop is disabled should return Daita filter chip`() =
+    fun `when Daita with direct only is enabled and multihop is disabled should return Daita filter chip`() =
         runTest {
             // Arrange
-            settings.value = mockk(relaxed = true) { every { isDaitaEnabled() } returns true }
+            settings.value =
+                mockk<Settings>(relaxed = true) {
+                    every { this@mockk.tunnelOptions.wireguard.daitaSettings.enabled } returns true
+                    every { tunnelOptions.wireguard.daitaSettings.directOnly } returns true
+                }
             wireguardConstraints.value =
                 mockk(relaxed = true) { every { isMultihopEnabled } returns false }
 
@@ -121,10 +125,29 @@ class FilterChipUseCaseTest {
         }
 
     @Test
-    fun `when Daita is enabled and multihop is enabled and relay list type is entry should return Daita filter chip`() =
+    fun `when Daita without direct only is enabled and multihop is disabled should return no filter chip`() =
         runTest {
             // Arrange
-            settings.value = mockk(relaxed = true) { every { isDaitaEnabled() } returns true }
+            settings.value =
+                mockk<Settings>(relaxed = true) {
+                    every { tunnelOptions.wireguard.daitaSettings.enabled } returns true
+                    every { tunnelOptions.wireguard.daitaSettings.directOnly } returns false
+                }
+            wireguardConstraints.value =
+                mockk(relaxed = true) { every { isMultihopEnabled } returns false }
+
+            filterChipUseCase(RelayListType.EXIT).test { assertLists(emptyList(), awaitItem()) }
+        }
+
+    @Test
+    fun `when Daita with direct only is enabled and multihop is enabled and relay list type is entry should return Daita filter chip`() =
+        runTest {
+            // Arrange
+            settings.value =
+                mockk<Settings>(relaxed = true) {
+                    every { tunnelOptions.wireguard.daitaSettings.enabled } returns true
+                    every { tunnelOptions.wireguard.daitaSettings.directOnly } returns true
+                }
             wireguardConstraints.value =
                 mockk(relaxed = true) { every { isMultihopEnabled } returns true }
 
@@ -134,10 +157,29 @@ class FilterChipUseCaseTest {
         }
 
     @Test
-    fun `when Daita is enabled and multihop is enabled and relay list type is exit should return no filter`() =
+    fun `when Daita with direct only is enabled and multihop is enabled and relay list type is exit should return no filter`() =
+        runTest {
+            // Arrange
+            settings.value =
+                mockk<Settings>(relaxed = true) {
+                    every { tunnelOptions.wireguard.daitaSettings.enabled } returns true
+                    every { tunnelOptions.wireguard.daitaSettings.directOnly } returns true
+                }
+            wireguardConstraints.value =
+                mockk(relaxed = true) { every { isMultihopEnabled } returns true }
+
+            filterChipUseCase(RelayListType.EXIT).test { assertLists(emptyList(), awaitItem()) }
+        }
+
+    @Test
+    fun `when Daita without direct only is enabled and multihop is enabled and relay list type is exit should return no filter`() =
         runTest {
             // Arrange
-            settings.value = mockk(relaxed = true) { every { isDaitaEnabled() } returns true }
+            settings.value =
+                mockk<Settings>(relaxed = true) {
+                    every { tunnelOptions.wireguard.daitaSettings.enabled } returns true
+                    every { tunnelOptions.wireguard.daitaSettings.directOnly } returns false
+                }
             wireguardConstraints.value =
                 mockk(relaxed = true) { every { isMultihopEnabled } returns true }
 
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/DaitaViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/DaitaViewModelTest.kt
new file mode 100644
index 0000000000..8eb9770826
--- /dev/null
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/DaitaViewModelTest.kt
@@ -0,0 +1,88 @@
+package net.mullvad.mullvadvpn.viewmodel
+
+import app.cash.turbine.test
+import arrow.core.right
+import io.mockk.coEvery
+import io.mockk.coVerify
+import io.mockk.every
+import io.mockk.mockk
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.test.runTest
+import net.mullvad.mullvadvpn.compose.state.DaitaUiState
+import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
+import net.mullvad.mullvadvpn.lib.model.Settings
+import net.mullvad.mullvadvpn.repository.SettingsRepository
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.ExtendWith
+
+@ExtendWith(TestCoroutineRule::class)
+class DaitaViewModelTest {
+    private val mockSettingsRepository: SettingsRepository = mockk()
+    private val settings = MutableStateFlow<Settings>(mockk(relaxed = true))
+
+    private lateinit var viewModel: DaitaViewModel
+
+    @BeforeEach
+    fun setUp() {
+        every { mockSettingsRepository.settingsUpdates } returns settings
+        viewModel = DaitaViewModel(mockSettingsRepository)
+    }
+
+    @Test
+    fun `given daita enabled ui state should be daita enabled`() = runTest {
+        // Arrange
+        val expectedState = DaitaUiState(daitaEnabled = true, directOnly = false)
+        settings.value = mockk {
+            every { tunnelOptions.wireguard.daitaSettings } returns
+                mockk {
+                    every { enabled } returns true
+                    every { directOnly } returns false
+                }
+        }
+
+        // Act, Assert
+        viewModel.uiState.test { assertEquals(expectedState, awaitItem()) }
+    }
+
+    @Test
+    fun `given direct only enabled ui state should be direct only enabled`() = runTest {
+        // Arrange
+        val expectedState = DaitaUiState(daitaEnabled = false, directOnly = true)
+        settings.value = mockk {
+            every { tunnelOptions.wireguard.daitaSettings } returns
+                mockk {
+                    every { enabled } returns false
+                    every { directOnly } returns true
+                }
+        }
+
+        // Act, Assert
+        viewModel.uiState.test { assertEquals(expectedState, awaitItem()) }
+    }
+
+    @Test
+    fun `set daita should call settings repository set daita enabled`() {
+        // Arrange
+        coEvery { mockSettingsRepository.setDaitaEnabled(any()) } returns Unit.right()
+
+        // Act
+        viewModel.setDaita(true)
+
+        // Assert
+        coVerify { mockSettingsRepository.setDaitaEnabled(true) }
+    }
+
+    @Test
+    fun `set direct only should call settings repository set daita direct only`() {
+        // Arrange
+        coEvery { mockSettingsRepository.setDaitaDirectOnly(any()) } returns Unit.right()
+
+        // Act
+        viewModel.setDirectOnly(true)
+
+        // Assert
+        coVerify { mockSettingsRepository.setDaitaDirectOnly(true) }
+    }
+}
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/SettingsViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/SettingsViewModelTest.kt
index f2468cbb11..7bc05df05c 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/SettingsViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/SettingsViewModelTest.kt
@@ -12,8 +12,10 @@ import kotlinx.coroutines.test.runTest
 import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
 import net.mullvad.mullvadvpn.lib.model.Constraint
 import net.mullvad.mullvadvpn.lib.model.DeviceState
+import net.mullvad.mullvadvpn.lib.model.Settings
 import net.mullvad.mullvadvpn.lib.model.WireguardConstraints
 import net.mullvad.mullvadvpn.lib.shared.DeviceRepository
+import net.mullvad.mullvadvpn.repository.SettingsRepository
 import net.mullvad.mullvadvpn.repository.WireguardConstraintsRepository
 import net.mullvad.mullvadvpn.ui.VersionInfo
 import net.mullvad.mullvadvpn.ui.serviceconnection.AppVersionInfoRepository
@@ -28,10 +30,12 @@ class SettingsViewModelTest {
     private val mockDeviceRepository: DeviceRepository = mockk()
     private val mockAppVersionInfoRepository: AppVersionInfoRepository = mockk()
     private val mockWireguardConstraintsRepository: WireguardConstraintsRepository = mockk()
+    private val mockSettingsRepository: SettingsRepository = mockk()
 
     private val versionInfo =
         MutableStateFlow(VersionInfo(currentVersion = "", isSupported = false))
     private val wireguardConstraints = MutableStateFlow<WireguardConstraints>(mockk(relaxed = true))
+    private val settings = MutableStateFlow(mockk<Settings>(relaxed = true))
 
     private lateinit var viewModel: SettingsViewModel
 
@@ -43,12 +47,14 @@ class SettingsViewModelTest {
         every { mockAppVersionInfoRepository.versionInfo } returns versionInfo
         every { mockWireguardConstraintsRepository.wireguardConstraints } returns
             wireguardConstraints
+        every { mockSettingsRepository.settingsUpdates } returns settings
 
         viewModel =
             SettingsViewModel(
                 deviceRepository = mockDeviceRepository,
                 appVersionInfoRepository = mockAppVersionInfoRepository,
                 wireguardConstraintsRepository = mockWireguardConstraintsRepository,
+                settingsRepository = mockSettingsRepository,
                 isPlayBuild = false,
             )
     }
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 427b003d33..aae283e91e 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
@@ -18,8 +18,10 @@ import kotlinx.coroutines.cancel
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.runTest
+import mullvad_daemon.management_interface.daitaSettings
 import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
 import net.mullvad.mullvadvpn.lib.model.Constraint
+import net.mullvad.mullvadvpn.lib.model.DaitaSettings
 import net.mullvad.mullvadvpn.lib.model.Mtu
 import net.mullvad.mullvadvpn.lib.model.Port
 import net.mullvad.mullvadvpn.lib.model.PortRange
@@ -131,7 +133,7 @@ class VpnSettingsViewModelTest {
                 WireguardTunnelOptions(
                     mtu = Mtu(0),
                     quantumResistant = expectedResistantState,
-                    daita = false,
+                    daitaSettings = DaitaSettings(enabled = false, directOnly = false),
                 )
 
             every { mockSettings.tunnelOptions } returns mockTunnelOptions
@@ -167,7 +169,7 @@ class VpnSettingsViewModelTest {
                         WireguardTunnelOptions(
                             mtu = null,
                             quantumResistant = QuantumResistantState.Off,
-                            daita = false,
+                            daitaSettings = DaitaSettings(enabled = false, directOnly = false),
                         ),
                     dnsOptions = mockk(relaxed = true),
                 )
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/location/SelectLocationListViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/location/SelectLocationListViewModelTest.kt
index 3584877170..acdf3f5c95 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/location/SelectLocationListViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/location/SelectLocationListViewModelTest.kt
@@ -15,7 +15,9 @@ import net.mullvad.mullvadvpn.lib.model.Constraint
 import net.mullvad.mullvadvpn.lib.model.GeoLocationId
 import net.mullvad.mullvadvpn.lib.model.RelayItem
 import net.mullvad.mullvadvpn.lib.model.RelayItemSelection
+import net.mullvad.mullvadvpn.lib.model.Settings
 import net.mullvad.mullvadvpn.repository.RelayListRepository
+import net.mullvad.mullvadvpn.repository.SettingsRepository
 import net.mullvad.mullvadvpn.repository.WireguardConstraintsRepository
 import net.mullvad.mullvadvpn.usecase.FilteredRelayListUseCase
 import net.mullvad.mullvadvpn.usecase.SelectedLocationUseCase
@@ -36,12 +38,14 @@ class SelectLocationListViewModelTest {
     private val mockWireguardConstraintsRepository: WireguardConstraintsRepository = mockk()
     private val mockRelayListRepository: RelayListRepository = mockk()
     private val mockCustomListRelayItemsUseCase: CustomListsRelayItemUseCase = mockk()
+    private val mockSettingsRepository: SettingsRepository = mockk()
 
     private val filteredRelayList = MutableStateFlow<List<RelayItem.Location.Country>>(emptyList())
     private val selectedLocationFlow = MutableStateFlow<RelayItemSelection>(mockk(relaxed = true))
     private val filteredCustomListRelayItems =
         MutableStateFlow<List<RelayItem.CustomList>>(emptyList())
     private val customListRelayItems = MutableStateFlow<List<RelayItem.CustomList>>(emptyList())
+    private val settings = MutableStateFlow(mockk<Settings>(relaxed = true))
 
     private lateinit var viewModel: SelectLocationListViewModel
 
@@ -57,6 +61,7 @@ class SelectLocationListViewModelTest {
         every { mockFilteredCustomListRelayItemsUseCase(any()) } returns
             filteredCustomListRelayItems
         every { mockCustomListRelayItemsUseCase() } returns customListRelayItems
+        every { mockSettingsRepository.settingsUpdates } returns settings
     }
 
     @Test
@@ -125,6 +130,7 @@ class SelectLocationListViewModelTest {
             wireguardConstraintsRepository = mockWireguardConstraintsRepository,
             relayListRepository = mockRelayListRepository,
             customListsRelayItemUseCase = mockCustomListRelayItemsUseCase,
+            settingsRepository = mockSettingsRepository,
         )
 
     private fun RelayListItem.relayItemId() =
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/location/SelectLocationViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/location/SelectLocationViewModelTest.kt
index ef21eac139..77a5e948eb 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/location/SelectLocationViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/location/SelectLocationViewModelTest.kt
@@ -34,6 +34,7 @@ import net.mullvad.mullvadvpn.relaylist.descendants
 import net.mullvad.mullvadvpn.repository.CustomListsRepository
 import net.mullvad.mullvadvpn.repository.RelayListFilterRepository
 import net.mullvad.mullvadvpn.repository.RelayListRepository
+import net.mullvad.mullvadvpn.repository.SettingsRepository
 import net.mullvad.mullvadvpn.repository.WireguardConstraintsRepository
 import net.mullvad.mullvadvpn.usecase.FilterChip
 import net.mullvad.mullvadvpn.usecase.FilterChipUseCase
@@ -53,6 +54,7 @@ class SelectLocationViewModelTest {
     private val mockCustomListsRepository: CustomListsRepository = mockk()
     private val mockWireguardConstraintsRepository: WireguardConstraintsRepository = mockk()
     private val mockFilterChipUseCase: FilterChipUseCase = mockk()
+    private val mocksSettingsRepository: SettingsRepository = mockk()
 
     private lateinit var viewModel: SelectLocationViewModel
 
@@ -67,6 +69,7 @@ class SelectLocationViewModelTest {
         every { mockWireguardConstraintsRepository.wireguardConstraints } returns
             wireguardConstraints
         every { mockFilterChipUseCase(any()) } returns filterChips
+        every { mocksSettingsRepository.settingsUpdates } returns MutableStateFlow(null)
 
         mockkStatic(RELAY_LIST_EXTENSIONS)
         mockkStatic(RELAY_ITEM_EXTENSIONS)
@@ -79,6 +82,7 @@ class SelectLocationViewModelTest {
                 customListsRepository = mockCustomListsRepository,
                 filterChipUseCase = mockFilterChipUseCase,
                 wireguardConstraintsRepository = mockWireguardConstraintsRepository,
+                settingsRepository = mocksSettingsRepository,
             )
     }
 
@@ -90,14 +94,7 @@ class SelectLocationViewModelTest {
 
     @Test
     fun `initial state should be correct`() = runTest {
-        Assertions.assertEquals(
-            SelectLocationUiState(
-                filterChips = emptyList(),
-                multihopEnabled = false,
-                relayListType = RelayListType.EXIT,
-            ),
-            viewModel.uiState.value,
-        )
+        Assertions.assertEquals(SelectLocationUiState.Loading, viewModel.uiState.value)
     }
 
     @Test
@@ -134,11 +131,15 @@ class SelectLocationViewModelTest {
                 awaitItem() // Default value
                 viewModel.selectRelayList(RelayListType.ENTRY)
                 // Assert relay list type is entry
-                assertEquals(RelayListType.ENTRY, awaitItem().relayListType)
+                val firstState = awaitItem()
+                assertIs<SelectLocationUiState.Data>(firstState)
+                assertEquals(RelayListType.ENTRY, firstState.relayListType)
                 // Select entry
                 viewModel.selectRelay(mockRelayItem)
-                // Await an empty item
-                assertEquals(RelayListType.EXIT, awaitItem().relayListType)
+                // Assert relay list type is exit
+                val secondState = awaitItem()
+                assertIs<SelectLocationUiState.Data>(secondState)
+                assertEquals(RelayListType.EXIT, secondState.relayListType)
                 coVerify { mockWireguardConstraintsRepository.setEntryLocation(relayItemId) }
             }
         }