Add and update unit tests

8 files changed, 691 insertions, 181 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
new file mode 100644
index 0000000000..8b3d6d68a2
--- /dev/null
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/FilterChipUseCaseTest.kt
@@ -0,0 +1,146 @@
+package net.mullvad.mullvadvpn.usecase
+
+import app.cash.turbine.test
+import io.mockk.every
+import io.mockk.mockk
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.test.runTest
+import net.mullvad.mullvadvpn.compose.state.RelayListType
+import net.mullvad.mullvadvpn.lib.common.test.assertLists
+import net.mullvad.mullvadvpn.lib.model.Constraint
+import net.mullvad.mullvadvpn.lib.model.Ownership
+import net.mullvad.mullvadvpn.lib.model.Provider
+import net.mullvad.mullvadvpn.lib.model.ProviderId
+import net.mullvad.mullvadvpn.lib.model.Providers
+import net.mullvad.mullvadvpn.lib.model.Settings
+import net.mullvad.mullvadvpn.lib.model.WireguardConstraints
+import net.mullvad.mullvadvpn.repository.RelayListFilterRepository
+import net.mullvad.mullvadvpn.repository.SettingsRepository
+import net.mullvad.mullvadvpn.repository.WireguardConstraintsRepository
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+
+class FilterChipUseCaseTest {
+
+    private val mockRelayListFilterRepository: RelayListFilterRepository = mockk()
+    private val mockAvailableProvidersUseCase: AvailableProvidersUseCase = mockk()
+    private val mockSettingRepository: SettingsRepository = mockk()
+    private val mockWireguardConstraintsRepository: WireguardConstraintsRepository = mockk()
+
+    private val selectedOwnership = MutableStateFlow<Constraint<Ownership>>(Constraint.Any)
+    private val selectedProviders = MutableStateFlow<Constraint<Providers>>(Constraint.Any)
+    private val availableProviders = MutableStateFlow<List<Provider>>(emptyList())
+    private val settings = MutableStateFlow<Settings>(mockk(relaxed = true))
+    private val wireguardConstraints = MutableStateFlow<WireguardConstraints>(mockk(relaxed = true))
+
+    private lateinit var filterChipUseCase: FilterChipUseCase
+
+    @BeforeEach
+    fun setUp() {
+        every { mockRelayListFilterRepository.selectedOwnership } returns selectedOwnership
+        every { mockRelayListFilterRepository.selectedProviders } returns selectedProviders
+        every { mockAvailableProvidersUseCase() } returns availableProviders
+        every { mockSettingRepository.settingsUpdates } returns settings
+        every { mockWireguardConstraintsRepository.wireguardConstraints } returns
+            wireguardConstraints
+
+        filterChipUseCase =
+            FilterChipUseCase(
+                relayListFilterRepository = mockRelayListFilterRepository,
+                availableProvidersUseCase = mockAvailableProvidersUseCase,
+                settingsRepository = mockSettingRepository,
+                wireguardConstraintsRepository = mockWireguardConstraintsRepository,
+            )
+    }
+
+    @Test
+    fun `when no filters are applied should return empty list`() = runTest {
+        filterChipUseCase(RelayListType.EXIT).test { assertLists(emptyList(), awaitItem()) }
+    }
+
+    @Test
+    fun `when ownership filter is applied should return correct ownership`() = runTest {
+        // Arrange
+        val expectedOwnership = Ownership.MullvadOwned
+        selectedOwnership.value = Constraint.Only(expectedOwnership)
+
+        filterChipUseCase(RelayListType.EXIT).test {
+            assertLists(listOf(FilterChip.Ownership(expectedOwnership)), awaitItem())
+        }
+    }
+
+    @Test
+    fun `when provider filter is applied should return correct number of providers`() = runTest {
+        // Arrange
+        val expectedProviders = Providers(providers = setOf(ProviderId("1"), ProviderId("2")))
+        selectedProviders.value = Constraint.Only(expectedProviders)
+        availableProviders.value =
+            listOf(
+                Provider(ProviderId("1"), Ownership.MullvadOwned),
+                Provider(ProviderId("2"), Ownership.Rented),
+            )
+
+        filterChipUseCase(RelayListType.EXIT).test {
+            assertLists(listOf(FilterChip.Provider(2)), awaitItem())
+        }
+    }
+
+    @Test
+    fun `when provider and ownership filter is applied should return correct filter chips`() =
+        runTest {
+            // Arrange
+            val expectedProviders = Providers(providers = setOf(ProviderId("1")))
+            val expectedOwnership = Ownership.MullvadOwned
+            selectedProviders.value = Constraint.Only(expectedProviders)
+            selectedOwnership.value = Constraint.Only(expectedOwnership)
+            availableProviders.value =
+                listOf(
+                    Provider(ProviderId("1"), Ownership.MullvadOwned),
+                    Provider(ProviderId("2"), Ownership.Rented),
+                )
+
+            filterChipUseCase(RelayListType.EXIT).test {
+                assertLists(
+                    listOf(FilterChip.Ownership(expectedOwnership), FilterChip.Provider(1)),
+                    awaitItem(),
+                )
+            }
+        }
+
+    @Test
+    fun `when Daita is enabled and multihop is disabled should return Daita filter chip`() =
+        runTest {
+            // Arrange
+            settings.value = mockk(relaxed = true) { every { isDaitaEnabled() } returns true }
+            wireguardConstraints.value =
+                mockk(relaxed = true) { every { isMultihopEnabled } returns false }
+
+            filterChipUseCase(RelayListType.EXIT).test {
+                assertLists(listOf(FilterChip.Daita), awaitItem())
+            }
+        }
+
+    @Test
+    fun `when Daita is enabled and multihop is enabled and relay list type is entry should return Daita filter chip`() =
+        runTest {
+            // Arrange
+            settings.value = mockk(relaxed = true) { every { isDaitaEnabled() } returns true }
+            wireguardConstraints.value =
+                mockk(relaxed = true) { every { isMultihopEnabled } returns true }
+
+            filterChipUseCase(RelayListType.ENTRY).test {
+                assertLists(listOf(FilterChip.Daita), awaitItem())
+            }
+        }
+
+    @Test
+    fun `when Daita 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 }
+            wireguardConstraints.value =
+                mockk(relaxed = true) { every { isMultihopEnabled } returns true }
+
+            filterChipUseCase(RelayListType.EXIT).test { assertLists(emptyList(), awaitItem()) }
+        }
+}
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/SelectedLocationUseCaseTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/SelectedLocationUseCaseTest.kt
new file mode 100644
index 0000000000..deef7b7ab9
--- /dev/null
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/SelectedLocationUseCaseTest.kt
@@ -0,0 +1,71 @@
+package net.mullvad.mullvadvpn.usecase
+
+import app.cash.turbine.test
+import io.mockk.every
+import io.mockk.mockk
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.test.runTest
+import net.mullvad.mullvadvpn.lib.model.Constraint
+import net.mullvad.mullvadvpn.lib.model.GeoLocationId
+import net.mullvad.mullvadvpn.lib.model.RelayItemId
+import net.mullvad.mullvadvpn.lib.model.RelayItemSelection
+import net.mullvad.mullvadvpn.lib.model.WireguardConstraints
+import net.mullvad.mullvadvpn.repository.RelayListRepository
+import net.mullvad.mullvadvpn.repository.WireguardConstraintsRepository
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+
+class SelectedLocationUseCaseTest {
+    private val mockRelayListRepository: RelayListRepository = mockk()
+    private val mockWireguardConstraintsRepository: WireguardConstraintsRepository = mockk()
+
+    private val selectedLocation = MutableStateFlow<Constraint<RelayItemId>>(Constraint.Any)
+    private val wireguardConstraints = MutableStateFlow<WireguardConstraints>(mockk(relaxed = true))
+
+    private lateinit var selectLocationUseCase: SelectedLocationUseCase
+
+    @BeforeEach
+    fun setup() {
+        every { mockRelayListRepository.selectedLocation } returns selectedLocation
+        every { mockWireguardConstraintsRepository.wireguardConstraints } returns
+            wireguardConstraints
+
+        selectLocationUseCase =
+            SelectedLocationUseCase(
+                relayListRepository = mockRelayListRepository,
+                wireguardConstraintsRepository = mockWireguardConstraintsRepository,
+            )
+    }
+
+    @Test
+    fun `when wireguard constraints is multihop enabled should return Multiple`() = runTest {
+        // Arrange
+        val entryLocation: Constraint<RelayItemId> = Constraint.Only(GeoLocationId.Country("se"))
+        val exitLocation = Constraint.Only(GeoLocationId.Country("us"))
+        wireguardConstraints.value =
+            WireguardConstraints(
+                isMultihopEnabled = true,
+                entryLocation = entryLocation,
+                port = Constraint.Any,
+            )
+        selectedLocation.value = exitLocation
+
+        // Act, Assert
+        selectLocationUseCase().test {
+            assertEquals(RelayItemSelection.Multiple(entryLocation, exitLocation), awaitItem())
+        }
+    }
+
+    @Test
+    fun `when wireguard constraints is multihop disabled should return Single`() = runTest {
+        // Arrange
+        val exitLocation = Constraint.Only(GeoLocationId.Country("us"))
+        selectedLocation.value = exitLocation
+
+        // Act, Assert
+        selectLocationUseCase().test {
+            assertEquals(RelayItemSelection.Single(exitLocation), awaitItem())
+        }
+    }
+}
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/MultihopViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/MultihopViewModelTest.kt
new file mode 100644
index 0000000000..34cb1353bb
--- /dev/null
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/MultihopViewModelTest.kt
@@ -0,0 +1,68 @@
+package net.mullvad.mullvadvpn.viewmodel
+
+import app.cash.turbine.test
+import arrow.core.Either
+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.lib.common.test.TestCoroutineRule
+import net.mullvad.mullvadvpn.lib.model.Constraint
+import net.mullvad.mullvadvpn.lib.model.WireguardConstraints
+import net.mullvad.mullvadvpn.repository.WireguardConstraintsRepository
+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 MultihopViewModelTest {
+
+    private val mockWireguardConstraintsRepository: WireguardConstraintsRepository = mockk()
+
+    private val wireguardConstraints = MutableStateFlow<WireguardConstraints>(mockk(relaxed = true))
+
+    private lateinit var multihopViewModel: MultihopViewModel
+
+    @BeforeEach
+    fun setUp() {
+        every { mockWireguardConstraintsRepository.wireguardConstraints } returns
+            wireguardConstraints
+
+        multihopViewModel =
+            MultihopViewModel(wireguardConstraintsRepository = mockWireguardConstraintsRepository)
+    }
+
+    @Test
+    fun `default state should be multihop disabled`() {
+        assertEquals(false, multihopViewModel.uiState.value.enable)
+    }
+
+    @Test
+    fun `when multihop enabled is true state should return multihop enabled true`() = runTest {
+        // Arrange
+        wireguardConstraints.value =
+            WireguardConstraints(
+                isMultihopEnabled = true,
+                entryLocation = Constraint.Any,
+                port = Constraint.Any,
+            )
+
+        // Act, Assert
+        multihopViewModel.uiState.test { assertEquals(MultihopUiState(true), awaitItem()) }
+    }
+
+    @Test
+    fun `when set multihop is called should call repository set multihop`() = runTest {
+        // Arrange
+        coEvery { mockWireguardConstraintsRepository.setMultihop(any()) } returns Either.Right(Unit)
+
+        // Act
+        multihopViewModel.setMultihop(true)
+
+        // Assert
+        coVerify { mockWireguardConstraintsRepository.setMultihop(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 8857eb364a..f2468cbb11 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
@@ -10,8 +10,11 @@ import kotlinx.coroutines.cancel
 import kotlinx.coroutines.flow.MutableStateFlow
 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.WireguardConstraints
 import net.mullvad.mullvadvpn.lib.shared.DeviceRepository
+import net.mullvad.mullvadvpn.repository.WireguardConstraintsRepository
 import net.mullvad.mullvadvpn.ui.VersionInfo
 import net.mullvad.mullvadvpn.ui.serviceconnection.AppVersionInfoRepository
 import org.junit.jupiter.api.AfterEach
@@ -24,9 +27,11 @@ class SettingsViewModelTest {
 
     private val mockDeviceRepository: DeviceRepository = mockk()
     private val mockAppVersionInfoRepository: AppVersionInfoRepository = mockk()
+    private val mockWireguardConstraintsRepository: WireguardConstraintsRepository = mockk()
 
     private val versionInfo =
         MutableStateFlow(VersionInfo(currentVersion = "", isSupported = false))
+    private val wireguardConstraints = MutableStateFlow<WireguardConstraints>(mockk(relaxed = true))
 
     private lateinit var viewModel: SettingsViewModel
 
@@ -36,11 +41,14 @@ class SettingsViewModelTest {
 
         every { mockDeviceRepository.deviceState } returns deviceState
         every { mockAppVersionInfoRepository.versionInfo } returns versionInfo
+        every { mockWireguardConstraintsRepository.wireguardConstraints } returns
+            wireguardConstraints
 
         viewModel =
             SettingsViewModel(
                 deviceRepository = mockDeviceRepository,
                 appVersionInfoRepository = mockAppVersionInfoRepository,
+                wireguardConstraintsRepository = mockWireguardConstraintsRepository,
                 isPlayBuild = false,
             )
     }
@@ -84,4 +92,22 @@ class SettingsViewModelTest {
                 assertEquals(false, result.isSupportedVersion)
             }
         }
+
+    @Test
+    fun `when WireguardConstraintsRepository return multihop enabled uiState should return multihop enabled true`() =
+        runTest {
+            // Arrange
+            wireguardConstraints.value =
+                WireguardConstraints(
+                    isMultihopEnabled = true,
+                    entryLocation = Constraint.Any,
+                    port = Constraint.Any,
+                )
+
+            // Act, Assert
+            viewModel.uiState.test {
+                val result = awaitItem()
+                assertEquals(true, result.multihopEnabled)
+            }
+        }
 }
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 340809fbb3..427b003d33 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
@@ -189,7 +189,7 @@ class VpnSettingsViewModelTest {
             val wireguardConstraints =
                 WireguardConstraints(
                     port = wireguardPort,
-                    useMultihop = false,
+                    isMultihopEnabled = false,
                     entryLocation = Constraint.Any,
                 )
             coEvery { mockWireguardConstraintsRepository.setWireguardPort(any()) } returns
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/location/SearchLocationViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/location/SearchLocationViewModelTest.kt
new file mode 100644
index 0000000000..be60f9d723
--- /dev/null
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/location/SearchLocationViewModelTest.kt
@@ -0,0 +1,161 @@
+package net.mullvad.mullvadvpn.viewmodel.location
+
+import app.cash.turbine.test
+import com.ramcosta.composedestinations.generated.navargs.toSavedStateHandle
+import io.mockk.every
+import io.mockk.mockk
+import kotlin.test.assertIs
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.test.runTest
+import net.mullvad.mullvadvpn.compose.screen.location.SearchLocationNavArgs
+import net.mullvad.mullvadvpn.compose.state.RelayListItem
+import net.mullvad.mullvadvpn.compose.state.RelayListType
+import net.mullvad.mullvadvpn.compose.state.SearchLocationUiState
+import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
+import net.mullvad.mullvadvpn.lib.common.test.assertLists
+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.WireguardConstraints
+import net.mullvad.mullvadvpn.repository.CustomListsRepository
+import net.mullvad.mullvadvpn.repository.RelayListFilterRepository
+import net.mullvad.mullvadvpn.repository.RelayListRepository
+import net.mullvad.mullvadvpn.repository.WireguardConstraintsRepository
+import net.mullvad.mullvadvpn.usecase.FilterChip
+import net.mullvad.mullvadvpn.usecase.FilterChipUseCase
+import net.mullvad.mullvadvpn.usecase.FilteredRelayListUseCase
+import net.mullvad.mullvadvpn.usecase.SelectedLocationUseCase
+import net.mullvad.mullvadvpn.usecase.customlists.CustomListActionUseCase
+import net.mullvad.mullvadvpn.usecase.customlists.CustomListsRelayItemUseCase
+import net.mullvad.mullvadvpn.usecase.customlists.FilterCustomListsRelayItemUseCase
+import org.junit.jupiter.api.Assertions.assertTrue
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.ExtendWith
+
+@ExtendWith(TestCoroutineRule::class)
+class SearchLocationViewModelTest {
+
+    private val mockWireguardConstraintsRepository: WireguardConstraintsRepository = mockk()
+    private val mockRelayListRepository: RelayListRepository = mockk()
+    private val mockFilteredRelayListUseCase: FilteredRelayListUseCase = mockk()
+    private val mockCustomListActionUseCase: CustomListActionUseCase = mockk()
+    private val mockCustomListsRepository: CustomListsRepository = mockk()
+    private val mockRelayListFilterRepository: RelayListFilterRepository = mockk()
+    private val mockFilterChipUseCase: FilterChipUseCase = mockk()
+    private val mockFilteredCustomListRelayItemsUseCase: FilterCustomListsRelayItemUseCase = mockk()
+    private val mockSelectedLocationUseCase: SelectedLocationUseCase = mockk()
+    private val mockCustomListsRelayItemUseCase: CustomListsRelayItemUseCase = mockk()
+
+    private val filteredRelayList = MutableStateFlow<List<RelayItem.Location.Country>>(emptyList())
+    private val selectedLocation =
+        MutableStateFlow<RelayItemSelection>(RelayItemSelection.Single(Constraint.Any))
+    private val filteredCustomListRelayItems =
+        MutableStateFlow<List<RelayItem.CustomList>>(emptyList())
+    private val customListRelayItems = MutableStateFlow<List<RelayItem.CustomList>>(emptyList())
+    private val filterChips = MutableStateFlow<List<FilterChip>>(emptyList())
+    private val wireguardConstraints = MutableStateFlow<WireguardConstraints>(mockk(relaxed = true))
+
+    private lateinit var viewModel: SearchLocationViewModel
+
+    @BeforeEach
+    fun setup() {
+        every { mockFilteredRelayListUseCase(any()) } returns filteredRelayList
+        every { mockSelectedLocationUseCase() } returns selectedLocation
+        every { mockFilteredCustomListRelayItemsUseCase(any()) } returns
+            filteredCustomListRelayItems
+        every { mockCustomListsRelayItemUseCase() } returns customListRelayItems
+        every { mockFilterChipUseCase(any()) } returns filterChips
+        every { mockWireguardConstraintsRepository.wireguardConstraints } returns
+            wireguardConstraints
+
+        viewModel =
+            SearchLocationViewModel(
+                wireguardConstraintsRepository = mockWireguardConstraintsRepository,
+                relayListRepository = mockRelayListRepository,
+                filteredRelayListUseCase = mockFilteredRelayListUseCase,
+                customListActionUseCase = mockCustomListActionUseCase,
+                customListsRepository = mockCustomListsRepository,
+                relayListFilterRepository = mockRelayListFilterRepository,
+                filterChipUseCase = mockFilterChipUseCase,
+                filteredCustomListRelayItemsUseCase = mockFilteredCustomListRelayItemsUseCase,
+                selectedLocationUseCase = mockSelectedLocationUseCase,
+                customListsRelayItemUseCase = mockCustomListsRelayItemUseCase,
+                savedStateHandle =
+                    SearchLocationNavArgs(relayListType = RelayListType.ENTRY).toSavedStateHandle(),
+            )
+    }
+
+    @Test
+    fun `on onSearchTermInput call uiState should emit with filtered countries`() = runTest {
+        // Arrange
+        val mockSearchString = "got"
+        filteredRelayList.value = testCountries
+
+        // Act, Assert
+        viewModel.uiState.test() {
+            // Wait for first data
+            assertIs<SearchLocationUiState.NoQuery>(awaitItem())
+
+            // Update search string
+            viewModel.onSearchInputUpdated(mockSearchString)
+
+            // We get some unnecessary emissions for now
+            awaitItem()
+
+            val actualState = awaitItem()
+            assertIs<SearchLocationUiState.Content>(actualState)
+            assertTrue(
+                actualState.relayListItems.filterIsInstance<RelayListItem.GeoLocationItem>().any {
+                    it.item is RelayItem.Location.City && it.item.name == "Gothenburg"
+                }
+            )
+        }
+    }
+
+    @Test
+    fun `when onSearchTermInput returns empty result uiState should return empty list`() = runTest {
+        // Arrange
+        filteredRelayList.value = testCountries
+        val mockSearchString = "SEARCH"
+
+        // Act, Assert
+        viewModel.uiState.test {
+            // Wait for first data
+            assertIs<SearchLocationUiState.NoQuery>(awaitItem())
+
+            // Update search string
+            viewModel.onSearchInputUpdated(mockSearchString)
+
+            // We get some unnecessary emissions for now
+            awaitItem()
+
+            // Assert
+            val actualState = awaitItem()
+            assertIs<SearchLocationUiState.Content>(actualState)
+            assertLists(
+                listOf(RelayListItem.LocationsEmptyText(mockSearchString)),
+                actualState.relayListItems,
+            )
+        }
+    }
+
+    companion object {
+        private val testCountries =
+            listOf(
+                RelayItem.Location.Country(
+                    id = GeoLocationId.Country("se"),
+                    "Sweden",
+                    listOf(
+                        RelayItem.Location.City(
+                            id = GeoLocationId.City(GeoLocationId.Country("se"), "got"),
+                            "Gothenburg",
+                            emptyList(),
+                        )
+                    ),
+                ),
+                RelayItem.Location.Country(id = GeoLocationId.Country("no"), "Norway", emptyList()),
+            )
+    }
+}
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
new file mode 100644
index 0000000000..3584877170
--- /dev/null
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/location/SelectLocationListViewModelTest.kt
@@ -0,0 +1,158 @@
+package net.mullvad.mullvadvpn.viewmodel.location
+
+import app.cash.turbine.test
+import io.mockk.every
+import io.mockk.mockk
+import kotlin.test.assertIs
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.test.runTest
+import net.mullvad.mullvadvpn.compose.state.RelayListItem
+import net.mullvad.mullvadvpn.compose.state.RelayListType
+import net.mullvad.mullvadvpn.compose.state.SelectLocationListUiState
+import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
+import net.mullvad.mullvadvpn.lib.common.test.assertLists
+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.repository.RelayListRepository
+import net.mullvad.mullvadvpn.repository.WireguardConstraintsRepository
+import net.mullvad.mullvadvpn.usecase.FilteredRelayListUseCase
+import net.mullvad.mullvadvpn.usecase.SelectedLocationUseCase
+import net.mullvad.mullvadvpn.usecase.customlists.CustomListsRelayItemUseCase
+import net.mullvad.mullvadvpn.usecase.customlists.FilterCustomListsRelayItemUseCase
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertTrue
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.ExtendWith
+
+@ExtendWith(TestCoroutineRule::class)
+class SelectLocationListViewModelTest {
+
+    private val mockFilteredRelayListUseCase: FilteredRelayListUseCase = mockk()
+    private val mockFilteredCustomListRelayItemsUseCase: FilterCustomListsRelayItemUseCase = mockk()
+    private val mockSelectedLocationUseCase: SelectedLocationUseCase = mockk()
+    private val mockWireguardConstraintsRepository: WireguardConstraintsRepository = mockk()
+    private val mockRelayListRepository: RelayListRepository = mockk()
+    private val mockCustomListRelayItemsUseCase: CustomListsRelayItemUseCase = 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 lateinit var viewModel: SelectLocationListViewModel
+
+    @BeforeEach
+    fun setUp() {
+        // Used for initial selection
+        every { mockRelayListRepository.selectedLocation } returns MutableStateFlow(Constraint.Any)
+        every { mockWireguardConstraintsRepository.wireguardConstraints } returns
+            MutableStateFlow(null)
+
+        every { mockSelectedLocationUseCase() } returns selectedLocationFlow
+        every { mockFilteredRelayListUseCase(any()) } returns filteredRelayList
+        every { mockFilteredCustomListRelayItemsUseCase(any()) } returns
+            filteredCustomListRelayItems
+        every { mockCustomListRelayItemsUseCase() } returns customListRelayItems
+    }
+
+    @Test
+    fun `initial state should be loading`() = runTest {
+        // Arrange
+        viewModel = createSelectLocationListViewModel(relayListType = RelayListType.ENTRY)
+
+        // Assert
+        assertEquals(SelectLocationListUiState.Loading, viewModel.uiState.value)
+    }
+
+    @Test
+    fun `given filteredRelayList emits update uiState should contain new update`() = runTest {
+        // Arrange
+        viewModel = createSelectLocationListViewModel(RelayListType.EXIT)
+        filteredRelayList.value = testCountries
+        val selectedId = testCountries.first().id
+        selectedLocationFlow.value = RelayItemSelection.Single(Constraint.Only(selectedId))
+
+        // Act, Assert
+        viewModel.uiState.test {
+            val actualState = awaitItem()
+            assertIs<SelectLocationListUiState.Content>(actualState)
+            assertLists(
+                testCountries.map { it.id },
+                actualState.relayListItems.mapNotNull { it.relayItemId() },
+            )
+            assertTrue(
+                actualState.relayListItems
+                    .filterIsInstance<RelayListItem.SelectableItem>()
+                    .first { it.relayItemId() == selectedId }
+                    .isSelected
+            )
+        }
+    }
+
+    @Test
+    fun `given relay is not selected all relay items should not be selected`() = runTest {
+        // Arrange
+        viewModel = createSelectLocationListViewModel(RelayListType.EXIT)
+        filteredRelayList.value = testCountries
+        selectedLocationFlow.value = RelayItemSelection.Single(Constraint.Any)
+
+        // Act, Assert
+        viewModel.uiState.test {
+            val actualState = awaitItem()
+            assertIs<SelectLocationListUiState.Content>(actualState)
+            assertLists(
+                testCountries.map { it.id },
+                actualState.relayListItems.mapNotNull { it.relayItemId() },
+            )
+            assertTrue(
+                actualState.relayListItems.filterIsInstance<RelayListItem.SelectableItem>().all {
+                    !it.isSelected
+                }
+            )
+        }
+    }
+
+    private fun createSelectLocationListViewModel(relayListType: RelayListType) =
+        SelectLocationListViewModel(
+            relayListType = relayListType,
+            filteredRelayListUseCase = mockFilteredRelayListUseCase,
+            filteredCustomListRelayItemsUseCase = mockFilteredCustomListRelayItemsUseCase,
+            selectedLocationUseCase = mockSelectedLocationUseCase,
+            wireguardConstraintsRepository = mockWireguardConstraintsRepository,
+            relayListRepository = mockRelayListRepository,
+            customListsRelayItemUseCase = mockCustomListRelayItemsUseCase,
+        )
+
+    private fun RelayListItem.relayItemId() =
+        when (this) {
+            is RelayListItem.CustomListFooter -> null
+            RelayListItem.CustomListHeader -> null
+            RelayListItem.LocationHeader -> null
+            is RelayListItem.LocationsEmptyText -> null
+            is RelayListItem.CustomListEntryItem -> item.id
+            is RelayListItem.CustomListItem -> item.id
+            is RelayListItem.GeoLocationItem -> item.id
+        }
+
+    companion object {
+        private val testCountries =
+            listOf(
+                RelayItem.Location.Country(
+                    id = GeoLocationId.Country("se"),
+                    "Sweden",
+                    listOf(
+                        RelayItem.Location.City(
+                            id = GeoLocationId.City(GeoLocationId.Country("se"), "got"),
+                            "Gothenburg",
+                            emptyList(),
+                        )
+                    ),
+                ),
+                RelayItem.Location.Country(id = GeoLocationId.Country("no"), "Norway", emptyList()),
+            )
+    }
+}
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/location/SelectLocationViewModelTest.kt
index bee888d279..ef21eac139 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/location/SelectLocationViewModelTest.kt
@@ -1,4 +1,4 @@
-package net.mullvad.mullvadvpn.viewmodel
+package net.mullvad.mullvadvpn.viewmodel.location
 
 import androidx.lifecycle.viewModelScope
 import app.cash.turbine.test
@@ -11,39 +11,35 @@ import io.mockk.mockkStatic
 import io.mockk.unmockkAll
 import kotlin.test.assertEquals
 import kotlin.test.assertIs
-import kotlin.test.assertTrue
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.test.runTest
 import net.mullvad.mullvadvpn.compose.communication.CustomListAction
 import net.mullvad.mullvadvpn.compose.communication.CustomListActionResultData
 import net.mullvad.mullvadvpn.compose.communication.LocationsChanged
-import net.mullvad.mullvadvpn.compose.state.RelayListItem
+import net.mullvad.mullvadvpn.compose.state.RelayListType
 import net.mullvad.mullvadvpn.compose.state.SelectLocationUiState
 import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
-import net.mullvad.mullvadvpn.lib.common.test.assertLists
 import net.mullvad.mullvadvpn.lib.model.Constraint
 import net.mullvad.mullvadvpn.lib.model.CustomList
 import net.mullvad.mullvadvpn.lib.model.CustomListId
 import net.mullvad.mullvadvpn.lib.model.CustomListName
 import net.mullvad.mullvadvpn.lib.model.GeoLocationId
 import net.mullvad.mullvadvpn.lib.model.Ownership
-import net.mullvad.mullvadvpn.lib.model.Provider
 import net.mullvad.mullvadvpn.lib.model.Providers
 import net.mullvad.mullvadvpn.lib.model.RelayItem
 import net.mullvad.mullvadvpn.lib.model.RelayItemId
-import net.mullvad.mullvadvpn.lib.model.Settings
+import net.mullvad.mullvadvpn.lib.model.WireguardConstraints
 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.usecase.AvailableProvidersUseCase
-import net.mullvad.mullvadvpn.usecase.FilteredRelayListUseCase
+import net.mullvad.mullvadvpn.repository.WireguardConstraintsRepository
+import net.mullvad.mullvadvpn.usecase.FilterChip
+import net.mullvad.mullvadvpn.usecase.FilterChipUseCase
 import net.mullvad.mullvadvpn.usecase.customlists.CustomListActionUseCase
-import net.mullvad.mullvadvpn.usecase.customlists.CustomListsRelayItemUseCase
-import net.mullvad.mullvadvpn.usecase.customlists.FilterCustomListsRelayItemUseCase
 import org.junit.jupiter.api.AfterEach
+import org.junit.jupiter.api.Assertions
 import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.extension.ExtendWith
@@ -52,39 +48,25 @@ import org.junit.jupiter.api.extension.ExtendWith
 class SelectLocationViewModelTest {
 
     private val mockRelayListFilterRepository: RelayListFilterRepository = mockk()
-    private val mockAvailableProvidersUseCase: AvailableProvidersUseCase = mockk(relaxed = true)
     private val mockCustomListActionUseCase: CustomListActionUseCase = mockk(relaxed = true)
-    private val mockFilteredCustomListRelayItemsUseCase: FilterCustomListsRelayItemUseCase = mockk()
-    private val mockFilteredRelayListUseCase: FilteredRelayListUseCase = mockk()
     private val mockRelayListRepository: RelayListRepository = mockk()
     private val mockCustomListsRepository: CustomListsRepository = mockk()
-    private val mockCustomListsRelayItemUseCase: CustomListsRelayItemUseCase = mockk()
-
-    private val mockSettingsRepository: SettingsRepository = mockk()
-    private val settingsFlow = MutableStateFlow(mockk<Settings>(relaxed = true))
+    private val mockWireguardConstraintsRepository: WireguardConstraintsRepository = mockk()
+    private val mockFilterChipUseCase: FilterChipUseCase = mockk()
 
     private lateinit var viewModel: SelectLocationViewModel
 
-    private val allProviders = MutableStateFlow<List<Provider>>(emptyList())
-    private val selectedOwnership = MutableStateFlow<Constraint<Ownership>>(Constraint.Any)
-    private val selectedProviders = MutableStateFlow<Constraint<Providers>>(Constraint.Any)
     private val selectedRelayItemFlow = MutableStateFlow<Constraint<RelayItemId>>(Constraint.Any)
-    private val filteredRelayList = MutableStateFlow<List<RelayItem.Location.Country>>(emptyList())
-    private val filteredCustomRelayListItems =
-        MutableStateFlow<List<RelayItem.CustomList>>(emptyList())
-    private val customListsRelayItem = MutableStateFlow<List<RelayItem.CustomList>>(emptyList())
+    private val wireguardConstraints = MutableStateFlow<WireguardConstraints>(mockk(relaxed = true))
+    private val filterChips = MutableStateFlow<List<FilterChip>>(emptyList())
 
     @BeforeEach
     fun setup() {
 
-        every { mockRelayListFilterRepository.selectedOwnership } returns selectedOwnership
-        every { mockRelayListFilterRepository.selectedProviders } returns selectedProviders
-        every { mockAvailableProvidersUseCase() } returns allProviders
         every { mockRelayListRepository.selectedLocation } returns selectedRelayItemFlow
-        every { mockFilteredRelayListUseCase() } returns filteredRelayList
-        every { mockFilteredCustomListRelayItemsUseCase() } returns filteredCustomRelayListItems
-        every { mockCustomListsRelayItemUseCase() } returns customListsRelayItem
-        every { mockSettingsRepository.settingsUpdates } returns settingsFlow
+        every { mockWireguardConstraintsRepository.wireguardConstraints } returns
+            wireguardConstraints
+        every { mockFilterChipUseCase(any()) } returns filterChips
 
         mockkStatic(RELAY_LIST_EXTENSIONS)
         mockkStatic(RELAY_ITEM_EXTENSIONS)
@@ -92,14 +74,11 @@ class SelectLocationViewModelTest {
         viewModel =
             SelectLocationViewModel(
                 relayListFilterRepository = mockRelayListFilterRepository,
-                availableProvidersUseCase = mockAvailableProvidersUseCase,
-                filteredCustomListRelayItemsUseCase = mockFilteredCustomListRelayItemsUseCase,
                 customListActionUseCase = mockCustomListActionUseCase,
-                filteredRelayListUseCase = mockFilteredRelayListUseCase,
                 relayListRepository = mockRelayListRepository,
                 customListsRepository = mockCustomListsRepository,
-                customListsRelayItemUseCase = mockCustomListsRelayItemUseCase,
-                settingsRepository = mockSettingsRepository,
+                filterChipUseCase = mockFilterChipUseCase,
+                wireguardConstraintsRepository = mockWireguardConstraintsRepository,
             )
     }
 
@@ -110,131 +89,59 @@ class SelectLocationViewModelTest {
     }
 
     @Test
-    fun `initial state should be loading`() = runTest {
-        assertEquals(SelectLocationUiState.Loading, viewModel.uiState.value)
-    }
-
-    @Test
-    fun `given filteredRelayList emits update uiState should contain new update`() = runTest {
-        // Arrange
-        filteredRelayList.value = testCountries
-        val selectedId = testCountries.first().id
-        selectedRelayItemFlow.value = Constraint.Only(selectedId)
-
-        // Act, Assert
-        viewModel.uiState.test {
-            val actualState = awaitItem()
-            assertIs<SelectLocationUiState.Content>(actualState)
-            assertLists(
-                testCountries.map { it.id },
-                actualState.relayListItems.mapNotNull { it.relayItemId() },
-            )
-            assertTrue(
-                actualState.relayListItems
-                    .filterIsInstance<RelayListItem.SelectableItem>()
-                    .first { it.relayItemId() == selectedId }
-                    .isSelected
-            )
-        }
-    }
-
-    @Test
-    fun `given relay is selected all relay items should not be selected`() = runTest {
-        // Arrange
-        filteredRelayList.value = testCountries
-        selectedRelayItemFlow.value = Constraint.Any
-
-        // Act, Assert
-        viewModel.uiState.test {
-            val actualState = awaitItem()
-            assertIs<SelectLocationUiState.Content>(actualState)
-            assertLists(
-                testCountries.map { it.id },
-                actualState.relayListItems.mapNotNull { it.relayItemId() },
-            )
-            assertTrue(
-                actualState.relayListItems.filterIsInstance<RelayListItem.SelectableItem>().all {
-                    !it.isSelected
-                }
-            )
-        }
-    }
-
-    @Test
-    fun `on selectRelay call uiSideEffect should emit CloseScreen and connect`() = runTest {
-        // Arrange
-        val mockRelayItem: RelayItem.Location.Country = mockk()
-        val relayItemId: GeoLocationId.Country = mockk(relaxed = true)
-        every { mockRelayItem.id } returns relayItemId
-        coEvery { mockRelayListRepository.updateSelectedRelayLocation(relayItemId) } returns
-            Unit.right()
-
-        // Act, Assert
-        viewModel.uiSideEffect.test {
-            viewModel.selectRelay(mockRelayItem)
-            // Await an empty item
-            assertEquals(SelectLocationSideEffect.CloseScreen, awaitItem())
-            coVerify { mockRelayListRepository.updateSelectedRelayLocation(relayItemId) }
-        }
+    fun `initial state should be correct`() = runTest {
+        Assertions.assertEquals(
+            SelectLocationUiState(
+                filterChips = emptyList(),
+                multihopEnabled = false,
+                relayListType = RelayListType.EXIT,
+            ),
+            viewModel.uiState.value,
+        )
     }
 
     @Test
-    fun `on onSearchTermInput call uiState should emit with filtered countries`() = runTest {
-        // Arrange
-        val mockSearchString = "got"
-        filteredRelayList.value = testCountries
-        selectedRelayItemFlow.value = Constraint.Any
-
-        // Act, Assert
-        viewModel.uiState.test {
-            // Wait for first data
-            assertIs<SelectLocationUiState.Content>(awaitItem())
-
-            // Update search string
-            viewModel.onSearchTermInput(mockSearchString)
-
-            // We get some unnecessary emissions for now
-            awaitItem()
-            awaitItem()
-            awaitItem()
+    fun `on selectRelay when relay list type is exit call uiSideEffect should emit CloseScreen and connect`() =
+        runTest {
+            // Arrange
+            val mockRelayItem: RelayItem.Location.Country = mockk()
+            val relayItemId: GeoLocationId.Country = mockk(relaxed = true)
+            every { mockRelayItem.id } returns relayItemId
+            coEvery { mockRelayListRepository.updateSelectedRelayLocation(relayItemId) } returns
+                Unit.right()
 
-            val actualState = awaitItem()
-            assertIs<SelectLocationUiState.Content>(actualState)
-            assertTrue(
-                actualState.relayListItems.filterIsInstance<RelayListItem.GeoLocationItem>().any {
-                    it.item is RelayItem.Location.City && it.item.name == "Gothenburg"
-                }
-            )
+            // Act, Assert
+            viewModel.uiSideEffect.test {
+                viewModel.selectRelay(mockRelayItem)
+                // Await an empty item
+                assertEquals(SelectLocationSideEffect.CloseScreen, awaitItem())
+                coVerify { mockRelayListRepository.updateSelectedRelayLocation(relayItemId) }
+            }
         }
-    }
 
     @Test
-    fun `when onSearchTermInput returns empty result uiState should return empty list`() = runTest {
-        // Arrange
-        filteredRelayList.value = testCountries
-        val mockSearchString = "SEARCH"
-
-        // Act, Assert
-        viewModel.uiState.test {
-            // Wait for first data
-            assertIs<SelectLocationUiState.Content>(awaitItem())
-
-            // Update search string
-            viewModel.onSearchTermInput(mockSearchString)
-
-            // We get some unnecessary emissions for now
-            awaitItem()
-            awaitItem()
+    fun `on selectRelay when relay list type is entry call uiSideEffect should switch relay list type to exit`() =
+        runTest {
+            // Arrange
+            val mockRelayItem: RelayItem.Location.Country = mockk()
+            val relayItemId: GeoLocationId.Country = mockk(relaxed = true)
+            every { mockRelayItem.id } returns relayItemId
+            coEvery { mockWireguardConstraintsRepository.setEntryLocation(relayItemId) } returns
+                Unit.right()
 
-            // Assert
-            val actualState = awaitItem()
-            assertIs<SelectLocationUiState.Content>(actualState)
-            assertEquals(
-                listOf(RelayListItem.LocationsEmptyText(mockSearchString)),
-                actualState.relayListItems,
-            )
+            // Act, Assert
+            viewModel.uiState.test {
+                awaitItem() // Default value
+                viewModel.selectRelayList(RelayListType.ENTRY)
+                // Assert relay list type is entry
+                assertEquals(RelayListType.ENTRY, awaitItem().relayListType)
+                // Select entry
+                viewModel.selectRelay(mockRelayItem)
+                // Await an empty item
+                assertEquals(RelayListType.EXIT, awaitItem().relayListType)
+                coVerify { mockWireguardConstraintsRepository.setEntryLocation(relayItemId) }
+            }
         }
-    }
 
     @Test
     fun `removeOwnerFilter should invoke use case with Constraint Any Ownership`() = runTest {
@@ -372,17 +279,6 @@ class SelectLocationViewModelTest {
             }
         }
 
-    private fun RelayListItem.relayItemId() =
-        when (this) {
-            is RelayListItem.CustomListFooter -> null
-            RelayListItem.CustomListHeader -> null
-            RelayListItem.LocationHeader -> null
-            is RelayListItem.LocationsEmptyText -> null
-            is RelayListItem.CustomListEntryItem -> item.id
-            is RelayListItem.CustomListItem -> item.id
-            is RelayListItem.GeoLocationItem -> item.id
-        }
-
     companion object {
         private const val RELAY_LIST_EXTENSIONS =
             "net.mullvad.mullvadvpn.relaylist.RelayListExtensionsKt"
@@ -390,21 +286,5 @@ class SelectLocationViewModelTest {
             "net.mullvad.mullvadvpn.relaylist.RelayItemExtensionsKt"
         private const val CUSTOM_LIST_EXTENSIONS =
             "net.mullvad.mullvadvpn.relaylist.CustomListExtensionsKt"
-
-        private val testCountries =
-            listOf(
-                RelayItem.Location.Country(
-                    id = GeoLocationId.Country("se"),
-                    "Sweden",
-                    listOf(
-                        RelayItem.Location.City(
-                            id = GeoLocationId.City(GeoLocationId.Country("se"), "got"),
-                            "Gothenburg",
-                            emptyList(),
-                        )
-                    ),
-                ),
-                RelayItem.Location.Country(id = GeoLocationId.Country("no"), "Norway", emptyList()),
-            )
     }
 }