Merge branch 'edit-list-view-gets-affected-by-filters-droid-903'

14 files changed, 157 insertions, 50 deletions

diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/data/DummyRelayItems.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/data/DummyRelayItems.kt
index fd4a97a1d2..2adfa22220 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/data/DummyRelayItems.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/data/DummyRelayItems.kt
@@ -1,6 +1,5 @@
 package net.mullvad.mullvadvpn.compose.data
 
-import net.mullvad.mullvadvpn.model.Constraint
 import net.mullvad.mullvadvpn.model.CustomListName
 import net.mullvad.mullvadvpn.model.PortRange
 import net.mullvad.mullvadvpn.model.RelayEndpointData
@@ -43,7 +42,7 @@ val DUMMY_RELAY_COUNTRIES =
             arrayListOf(DUMMY_RELAY_COUNTRY_1, DUMMY_RELAY_COUNTRY_2),
             DUMMY_WIREGUARD_ENDPOINT_DATA,
         )
-        .toRelayCountries(ownership = Constraint.Any(), providers = Constraint.Any())
+        .toRelayCountries()
 
 val DUMMY_CUSTOM_LISTS =
     listOf(
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/EditCustomListScreen.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/EditCustomListScreen.kt
index 87ed88d263..9186e639c5 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/EditCustomListScreen.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/EditCustomListScreen.kt
@@ -45,6 +45,7 @@ import net.mullvad.mullvadvpn.compose.transitions.SlideInFromRightTransition
 import net.mullvad.mullvadvpn.lib.theme.AppTheme
 import net.mullvad.mullvadvpn.lib.theme.Dimens
 import net.mullvad.mullvadvpn.model.GeographicLocationConstraint
+import net.mullvad.mullvadvpn.model.Ownership
 import net.mullvad.mullvadvpn.relaylist.RelayItem
 import net.mullvad.mullvadvpn.viewmodel.EditCustomListViewModel
 import org.koin.androidx.compose.koinViewModel
@@ -69,7 +70,9 @@ private fun PreviewEditCustomListScreen() {
                                     "hostname",
                                     "hostname",
                                     "hostname"
-                                )
+                                ),
+                                "Provider",
+                                Ownership.MullvadOwned
                             )
                         )
                 )
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/util/PreviewData.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/util/PreviewData.kt
index 61b563564c..3581d1d0b4 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/util/PreviewData.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/util/PreviewData.kt
@@ -1,6 +1,7 @@
 package net.mullvad.mullvadvpn.compose.util
 
 import net.mullvad.mullvadvpn.model.GeographicLocationConstraint
+import net.mullvad.mullvadvpn.model.Ownership
 import net.mullvad.mullvadvpn.relaylist.RelayItem
 
 fun generateRelayItemCountry(
@@ -65,7 +66,9 @@ fun generateRelayItemRelay(
                 hostname = hostName,
             ),
         locationName = "$cityCode $hostName",
-        active = active
+        active = active,
+        providerName = "Provider",
+        ownership = Ownership.MullvadOwned,
     )
 
 private fun String.generateCountryCode() = (take(1) + takeLast(1)).lowercase()
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItem.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItem.kt
index ce4be395b6..af4a0084d2 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItem.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItem.kt
@@ -3,6 +3,7 @@ package net.mullvad.mullvadvpn.relaylist
 import net.mullvad.mullvadvpn.model.CustomListName
 import net.mullvad.mullvadvpn.model.GeoIpLocation
 import net.mullvad.mullvadvpn.model.GeographicLocationConstraint
+import net.mullvad.mullvadvpn.model.Ownership
 
 sealed interface RelayItem {
     val name: String
@@ -66,6 +67,8 @@ sealed interface RelayItem {
         override val locationName: String,
         override val active: Boolean,
         val location: GeographicLocationConstraint.Hostname,
+        val providerName: String,
+        val ownership: Ownership,
     ) : RelayItem {
         override val code = name
         override val hasChildren = false
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayList.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayList.kt
index 30e4146245..e469aec118 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayList.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayList.kt
@@ -2,6 +2,7 @@ package net.mullvad.mullvadvpn.relaylist
 
 data class RelayList(
     val customLists: List<RelayItem.CustomList>,
-    val country: List<RelayItem.Country>,
+    val allCountries: List<RelayItem.Country>,
+    val filteredCountries: List<RelayItem.Country>,
     val selectedItem: RelayItem?,
 )
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayListExtensions.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayListExtensions.kt
index 882a3e42a4..8d2ea5f348 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayListExtensions.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayListExtensions.kt
@@ -12,10 +12,7 @@ import net.mullvad.mullvadvpn.model.RelayList
  * filtered out and also relays that do not fit the ownership and provider list So are also cities
  * that only contains non-wireguard relays Countries, cities and relays are ordered by name
  */
-fun RelayList.toRelayCountries(
-    ownership: Constraint<Ownership>,
-    providers: Constraint<Providers>
-): List<RelayItem.Country> {
+fun RelayList.toRelayCountries(): List<RelayItem.Country> {
     val relayCountries =
         this.countries
             .map { country ->
@@ -33,8 +30,7 @@ fun RelayList.toRelayCountries(
                             relays = relays
                         )
 
-                    val validCityRelays =
-                        city.relays.filterValidRelays(ownership = ownership, providers = providers)
+                    val validCityRelays = city.relays.filterValidRelays()
 
                     for (relay in validCityRelays) {
                         relays.add(
@@ -47,7 +43,10 @@ fun RelayList.toRelayCountries(
                                         relay.hostname
                                     ),
                                 locationName = "${city.name} (${relay.hostname})",
-                                active = relay.active
+                                active = relay.active,
+                                providerName = relay.provider,
+                                ownership =
+                                    if (relay.owned) Ownership.MullvadOwned else Ownership.Rented
                             )
                         )
                     }
@@ -172,25 +171,40 @@ fun List<RelayItem.Country>.filterOnSearchTerm(
     }
 }
 
-private fun List<DaemonRelay>.filterValidRelays(
+private fun List<DaemonRelay>.filterValidRelays(): List<DaemonRelay> = filter {
+    it.isWireguardRelay
+}
+
+fun List<RelayItem.Country>.filterOnOwnershipAndProviders(
     ownership: Constraint<Ownership>,
     providers: Constraint<Providers>
-): List<DaemonRelay> =
-    filter { it.isWireguardRelay }
-        .filter {
+): List<RelayItem.Country> {
+    return map { country ->
+            val cities =
+                country.cities.map { city ->
+                    val relays =
+                        city.relays.filterRelayByOwnershipAndProviders(ownership, providers)
+                    city.copy(relays = relays)
+                }
+            country.copy(cities = cities.filter { it.relays.isNotEmpty() })
+        }
+        .filter { it.cities.isNotEmpty() }
+}
+
+private fun List<RelayItem.Relay>.filterRelayByOwnershipAndProviders(
+    ownership: Constraint<Ownership>,
+    providers: Constraint<Providers>
+): List<RelayItem.Relay> =
+    filter {
             when (ownership) {
                 is Constraint.Any -> true
-                is Constraint.Only ->
-                    when (ownership.value) {
-                        Ownership.MullvadOwned -> it.owned
-                        Ownership.Rented -> !it.owned
-                    }
+                is Constraint.Only -> it.ownership == ownership.value
             }
         }
         .filter { relay ->
             when (providers) {
                 is Constraint.Any -> true
-                is Constraint.Only -> providers.value.providers.contains(relay.provider)
+                is Constraint.Only -> providers.value.providers.contains(relay.providerName)
             }
         }
 
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/RelayListUseCase.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/RelayListUseCase.kt
index 4957818283..bdeaac0d0c 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/RelayListUseCase.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/RelayListUseCase.kt
@@ -9,6 +9,7 @@ import net.mullvad.mullvadvpn.model.RelaySettings
 import net.mullvad.mullvadvpn.model.WireguardConstraints
 import net.mullvad.mullvadvpn.relaylist.RelayItem
 import net.mullvad.mullvadvpn.relaylist.RelayList
+import net.mullvad.mullvadvpn.relaylist.filterOnOwnershipAndProviders
 import net.mullvad.mullvadvpn.relaylist.findItemForGeographicLocationConstraint
 import net.mullvad.mullvadvpn.relaylist.toRelayCountries
 import net.mullvad.mullvadvpn.relaylist.toRelayItemLists
@@ -36,22 +37,29 @@ class RelayListUseCase(
                 settings?.relaySettings?.relayConstraints()?.ownership ?: Constraint.Any()
             val providers =
                 settings?.relaySettings?.relayConstraints()?.providers ?: Constraint.Any()
-            val relayCountries =
-                relayList.toRelayCountries(ownership = ownership, providers = providers)
+            val relayCountries = relayList.toRelayCountries()
             val customLists =
                 settings?.customLists?.customLists?.toRelayItemLists(relayCountries) ?: emptyList()
+            val relayCountriesFiltered =
+                relayCountries.filterOnOwnershipAndProviders(ownership, providers)
             val selectedItem =
                 findSelectedRelayItem(
                     relaySettings = settings?.relaySettings,
-                    relayCountries = relayCountries,
+                    relayCountries = relayCountriesFiltered,
                     customLists = customLists,
                 )
-            RelayList(customLists, relayCountries, selectedItem)
+            RelayList(
+                customLists = customLists,
+                allCountries = relayCountries,
+                filteredCountries = relayCountriesFiltered,
+                selectedItem = selectedItem
+            )
         }
 
     fun selectedRelayItem(): Flow<RelayItem?> = relayListWithSelection().map { it.selectedItem }
 
-    fun relayList(): Flow<List<RelayItem.Country>> = relayListWithSelection().map { it.country }
+    fun fullRelayList(): Flow<List<RelayItem.Country>> =
+        relayListWithSelection().map { it.allCountries }
 
     fun customLists(): Flow<List<RelayItem.CustomList>> =
         relayListWithSelection().map { it.customLists }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/customlists/CustomListActionUseCase.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/customlists/CustomListActionUseCase.kt
index 8d722325d6..16c86c0d59 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/customlists/CustomListActionUseCase.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/customlists/CustomListActionUseCase.kt
@@ -53,7 +53,7 @@ class CustomListActionUseCase(
                     )
                     val locationNames =
                         relayListUseCase
-                            .relayList()
+                            .fullRelayList()
                             .firstOrNull()
                             ?.getRelayItemsByCodes(action.locations)
                             ?.map { it.name }
@@ -61,7 +61,7 @@ class CustomListActionUseCase(
                         CustomListResult.Created(
                             id = result.id,
                             name = action.name,
-                            locationName = locationNames?.first(),
+                            locationName = locationNames?.firstOrNull(),
                             undo = action.not(result.id)
                         )
                     )
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListLocationsViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListLocationsViewModel.kt
index 5efba5321e..cdbcebbb83 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListLocationsViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListLocationsViewModel.kt
@@ -37,7 +37,7 @@ class CustomListLocationsViewModel(
     private val _searchTerm = MutableStateFlow(EMPTY_SEARCH_TERM)
 
     val uiState =
-        combine(relayListUseCase.relayList(), _searchTerm, _selectedLocations) {
+        combine(relayListUseCase.fullRelayList(), _searchTerm, _selectedLocations) {
                 relayCountries,
                 searchTerm,
                 selectedLocations ->
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModel.kt
index 15df90be9e..586bbf8bc1 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModel.kt
@@ -35,6 +35,7 @@ class SelectLocationViewModel(
 ) : ViewModel() {
     private val _searchTerm = MutableStateFlow(EMPTY_SEARCH_TERM)
 
+    @Suppress("DestructuringDeclarationWithTooManyEntries")
     val uiState =
         combine(
                 relayListUseCase.relayListWithSelection(),
@@ -43,7 +44,7 @@ class SelectLocationViewModel(
                 relayListFilterUseCase.availableProviders(),
                 relayListFilterUseCase.selectedProviders(),
             ) {
-                (customLists, relayCountries, selectedItem),
+                (customLists, _, relayCountries, selectedItem),
                 searchTerm,
                 selectedOwnership,
                 allProviders,
@@ -63,7 +64,15 @@ class SelectLocationViewModel(
                 val filteredRelayCountries =
                     relayCountries.filterOnSearchTerm(searchTerm, selectedItem)
 
-                val filteredCustomLists = customLists.filterOnSearchTerm(searchTerm)
+                val filteredCustomLists =
+                    customLists.filterOnSearchTerm(searchTerm).map { customList ->
+                        customList.copy(
+                            locations =
+                                customList.locations.filter { location ->
+                                    filteredRelayCountries.any { it.code == location.code }
+                                }
+                        )
+                    }
 
                 SelectLocationUiState.Content(
                     searchTerm = searchTerm,
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/relaylist/RelayNameComparatorTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/relaylist/RelayNameComparatorTest.kt
index 1532328729..eb66c2d4f9 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/relaylist/RelayNameComparatorTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/relaylist/RelayNameComparatorTest.kt
@@ -2,6 +2,7 @@ package net.mullvad.mullvadvpn.relaylist
 
 import io.mockk.mockk
 import io.mockk.unmockkAll
+import net.mullvad.mullvadvpn.model.Ownership
 import org.junit.jupiter.api.AfterEach
 import org.junit.jupiter.api.Assertions.assertTrue
 import org.junit.jupiter.api.Test
@@ -21,6 +22,8 @@ class RelayNameComparatorTest {
                 location = mockk(),
                 locationName = "mock",
                 active = false,
+                providerName = "Provider",
+                ownership = Ownership.MullvadOwned
             )
         val relay10 =
             RelayItem.Relay(
@@ -28,6 +31,8 @@ class RelayNameComparatorTest {
                 location = mockk(),
                 locationName = "mock",
                 active = false,
+                providerName = "Provider",
+                ownership = Ownership.MullvadOwned
             )
 
         relay9 assertOrderBothDirection relay10
@@ -41,6 +46,8 @@ class RelayNameComparatorTest {
                 location = mockk(),
                 locationName = "mock",
                 active = false,
+                providerName = "Provider",
+                ownership = Ownership.MullvadOwned
             )
         val relay9b =
             RelayItem.Relay(
@@ -48,6 +55,8 @@ class RelayNameComparatorTest {
                 location = mockk(),
                 locationName = "mock",
                 active = false,
+                providerName = "Provider",
+                ownership = Ownership.MullvadOwned
             )
 
         assertTrue(RelayNameComparator.compare(relay9a, relay9b) == 0)
@@ -57,13 +66,41 @@ class RelayNameComparatorTest {
     @Test
     fun `comparator should be able to handle name of only numbers`() {
         val relay001 =
-            RelayItem.Relay(name = "001", location = mockk(), locationName = "mock", active = false)
+            RelayItem.Relay(
+                name = "001",
+                location = mockk(),
+                locationName = "mock",
+                active = false,
+                providerName = "Provider",
+                ownership = Ownership.MullvadOwned
+            )
         val relay1 =
-            RelayItem.Relay(name = "1", location = mockk(), locationName = "mock", active = false)
+            RelayItem.Relay(
+                name = "1",
+                location = mockk(),
+                locationName = "mock",
+                active = false,
+                providerName = "Provider",
+                ownership = Ownership.MullvadOwned
+            )
         val relay3 =
-            RelayItem.Relay(name = "3", location = mockk(), locationName = "mock", active = false)
+            RelayItem.Relay(
+                name = "3",
+                location = mockk(),
+                locationName = "mock",
+                active = false,
+                providerName = "Provider",
+                ownership = Ownership.MullvadOwned
+            )
         val relay100 =
-            RelayItem.Relay(name = "100", location = mockk(), locationName = "mock", active = false)
+            RelayItem.Relay(
+                name = "100",
+                location = mockk(),
+                locationName = "mock",
+                active = false,
+                providerName = "Provider",
+                ownership = Ownership.MullvadOwned
+            )
 
         relay001 assertOrderBothDirection relay1
         relay001 assertOrderBothDirection relay3
@@ -79,6 +116,8 @@ class RelayNameComparatorTest {
                 location = mockk(),
                 locationName = "mock",
                 active = false,
+                providerName = "Provider",
+                ownership = Ownership.MullvadOwned
             )
         val relay9b =
             RelayItem.Relay(
@@ -86,6 +125,8 @@ class RelayNameComparatorTest {
                 location = mockk(),
                 locationName = "mock",
                 active = false,
+                providerName = "Provider",
+                ownership = Ownership.MullvadOwned
             )
 
         assertTrue(RelayNameComparator.compare(relay9a, relay9b) == 0)
@@ -100,6 +141,8 @@ class RelayNameComparatorTest {
                 location = mockk(),
                 locationName = "mock",
                 active = false,
+                providerName = "Provider",
+                ownership = Ownership.MullvadOwned
             )
         val relay005 =
             RelayItem.Relay(
@@ -107,6 +150,8 @@ class RelayNameComparatorTest {
                 location = mockk(),
                 locationName = "mock",
                 active = false,
+                providerName = "Provider",
+                ownership = Ownership.MullvadOwned
             )
 
         relay001 assertOrderBothDirection relay005
@@ -120,6 +165,8 @@ class RelayNameComparatorTest {
                 location = mockk(),
                 locationName = "mock",
                 active = false,
+                providerName = "Provider",
+                ownership = Ownership.MullvadOwned
             )
         val relayAr8 =
             RelayItem.Relay(
@@ -127,6 +174,8 @@ class RelayNameComparatorTest {
                 location = mockk(),
                 locationName = "mock",
                 active = false,
+                providerName = "Provider",
+                ownership = Ownership.MullvadOwned
             )
         val relaySe5 =
             RelayItem.Relay(
@@ -134,6 +183,8 @@ class RelayNameComparatorTest {
                 location = mockk(),
                 locationName = "mock",
                 active = false,
+                providerName = "Provider",
+                ownership = Ownership.MullvadOwned
             )
         val relaySe10 =
             RelayItem.Relay(
@@ -141,6 +192,8 @@ class RelayNameComparatorTest {
                 location = mockk(),
                 locationName = "mock",
                 active = false,
+                providerName = "Provider",
+                ownership = Ownership.MullvadOwned
             )
 
         relayAr2 assertOrderBothDirection relayAr8
@@ -155,7 +208,9 @@ class RelayNameComparatorTest {
                 name = "se2-cloud",
                 location = mockk(),
                 locationName = "mock",
-                active = false
+                active = false,
+                providerName = "Provider",
+                ownership = Ownership.MullvadOwned
             )
         val relay2w =
             RelayItem.Relay(
@@ -163,6 +218,8 @@ class RelayNameComparatorTest {
                 location = mockk(),
                 locationName = "mock",
                 active = false,
+                providerName = "Provider",
+                ownership = Ownership.MullvadOwned
             )
 
         relay2c assertOrderBothDirection relay2w
@@ -175,7 +232,9 @@ class RelayNameComparatorTest {
                 name = "se22",
                 location = mockk(),
                 locationName = "mock",
-                active = false
+                active = false,
+                providerName = "Provider",
+                ownership = Ownership.MullvadOwned
             )
         val relay22b =
             RelayItem.Relay(
@@ -183,6 +242,8 @@ class RelayNameComparatorTest {
                 location = mockk(),
                 locationName = "mock",
                 active = false,
+                providerName = "Provider",
+                ownership = Ownership.MullvadOwned
             )
 
         relay22a assertOrderBothDirection relay22b
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/CustomListActionUseCaseTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/CustomListActionUseCaseTest.kt
index fdcb4170f0..4dfb95768b 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/CustomListActionUseCaseTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/CustomListActionUseCaseTest.kt
@@ -71,7 +71,7 @@ class CustomListActionUseCaseTest {
                 listOf(locationCode)
             )
         } returns UpdateCustomListResult.Ok
-        coEvery { mockRelayListUseCase.relayList() } returns flowOf(mockLocations)
+        coEvery { mockRelayListUseCase.fullRelayList() } returns flowOf(mockLocations)
         every { mockLocations.getRelayItemsByCodes(listOf(locationCode)) } returns mockLocations
 
         // Act
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListLocationsViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListLocationsViewModelTest.kt
index df10ba96c4..d21789d36f 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListLocationsViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListLocationsViewModelTest.kt
@@ -12,6 +12,7 @@ import net.mullvad.mullvadvpn.compose.communication.CustomListResult
 import net.mullvad.mullvadvpn.compose.state.CustomListLocationsUiState
 import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
 import net.mullvad.mullvadvpn.model.GeographicLocationConstraint
+import net.mullvad.mullvadvpn.model.Ownership
 import net.mullvad.mullvadvpn.relaylist.RelayItem
 import net.mullvad.mullvadvpn.relaylist.descendants
 import net.mullvad.mullvadvpn.usecase.RelayListUseCase
@@ -31,7 +32,7 @@ class CustomListLocationsViewModelTest {
 
     @BeforeEach
     fun setup() {
-        every { mockRelayListUseCase.relayList() } returns relayListFlow
+        every { mockRelayListUseCase.fullRelayList() } returns relayListFlow
         every { mockRelayListUseCase.customLists() } returns customListFlow
     }
 
@@ -283,7 +284,9 @@ class CustomListLocationsViewModelTest {
                                                     "SE",
                                                     "GBG",
                                                     "gbg-1"
-                                                )
+                                                ),
+                                            providerName = "Provider",
+                                            ownership = Ownership.MullvadOwned
                                         )
                                     )
                             )
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModelTest.kt
index 41bff94ccd..6522892ec2 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModelTest.kt
@@ -49,7 +49,7 @@ class SelectLocationViewModelTest {
     private val mockServiceConnectionManager: ServiceConnectionManager = mockk()
     private lateinit var viewModel: SelectLocationViewModel
     private val relayListWithSelectionFlow =
-        MutableStateFlow(RelayList(emptyList(), emptyList(), null))
+        MutableStateFlow(RelayList(emptyList(), emptyList(), emptyList(), null))
     private val mockRelayListUseCase: RelayListUseCase = mockk()
     private val mockCustomListActionUseCase: CustomListActionUseCase = mockk(relaxed = true)
     private val selectedOwnership = MutableStateFlow<Constraint<Ownership>>(Constraint.Any())
@@ -93,10 +93,11 @@ class SelectLocationViewModelTest {
     fun `given relayListWithSelection emits update uiState should contain new update`() = runTest {
         // Arrange
         val mockCountries = listOf<RelayItem.Country>(mockk(), mockk())
-        val mockCustomList = listOf<RelayItem.CustomList>(mockk())
+        val mockCustomList = listOf<RelayItem.CustomList>(mockk(relaxed = true))
         val selectedItem: RelayItem = mockk()
         every { mockCountries.filterOnSearchTerm(any(), selectedItem) } returns mockCountries
-        relayListWithSelectionFlow.value = RelayList(mockCustomList, mockCountries, selectedItem)
+        relayListWithSelectionFlow.value =
+            RelayList(mockCustomList, mockCountries, mockCountries, selectedItem)
 
         // Act, Assert
         viewModel.uiState.test {
@@ -111,12 +112,12 @@ class SelectLocationViewModelTest {
     fun `given relayListWithSelection emits update with no selections selectedItem should be null`() =
         runTest {
             // Arrange
-            val mockCustomList = listOf<RelayItem.CustomList>(mockk())
+            val mockCustomList = listOf<RelayItem.CustomList>(mockk(relaxed = true))
             val mockCountries = listOf<RelayItem.Country>(mockk(), mockk())
             val selectedItem: RelayItem? = null
             every { mockCountries.filterOnSearchTerm(any(), selectedItem) } returns mockCountries
             relayListWithSelectionFlow.value =
-                RelayList(mockCustomList, mockCountries, selectedItem)
+                RelayList(mockCustomList, mockCountries, mockCountries, selectedItem)
 
             // Act, Assert
             viewModel.uiState.test {
@@ -155,7 +156,7 @@ class SelectLocationViewModelTest {
     @Test
     fun `on onSearchTermInput call uiState should emit with filtered countries`() = runTest {
         // Arrange
-        val mockCustomList = listOf<RelayItem.CustomList>(mockk())
+        val mockCustomList = listOf<RelayItem.CustomList>(mockk(relaxed = true))
         val mockCountries = listOf<RelayItem.Country>(mockk(), mockk())
         val selectedItem: RelayItem? = null
         val mockRelayList: List<RelayItem.Country> = mockk(relaxed = true)
@@ -163,7 +164,8 @@ class SelectLocationViewModelTest {
         every { mockRelayList.filterOnSearchTerm(mockSearchString, selectedItem) } returns
             mockCountries
         every { mockCustomList.filterOnSearchTerm(mockSearchString) } returns mockCustomList
-        relayListWithSelectionFlow.value = RelayList(mockCustomList, mockRelayList, selectedItem)
+        relayListWithSelectionFlow.value =
+            RelayList(mockCustomList, mockRelayList, mockRelayList, selectedItem)
 
         // Act, Assert
         viewModel.uiState.test {
@@ -184,7 +186,7 @@ class SelectLocationViewModelTest {
     @Test
     fun `when onSearchTermInput returns empty result uiState should return empty list`() = runTest {
         // Arrange
-        val mockCustomList = listOf<RelayItem.CustomList>(mockk())
+        val mockCustomList = listOf<RelayItem.CustomList>(mockk(relaxed = true))
         val mockCountries = emptyList<RelayItem.Country>()
         val selectedItem: RelayItem? = null
         val mockRelayList: List<RelayItem.Country> = mockk(relaxed = true)
@@ -192,7 +194,8 @@ class SelectLocationViewModelTest {
         every { mockRelayList.filterOnSearchTerm(mockSearchString, selectedItem) } returns
             mockCountries
         every { mockCustomList.filterOnSearchTerm(mockSearchString) } returns mockCustomList
-        relayListWithSelectionFlow.value = RelayList(mockCustomList, mockRelayList, selectedItem)
+        relayListWithSelectionFlow.value =
+            RelayList(mockCustomList, mockRelayList, mockRelayList, selectedItem)
 
         // Act, Assert
         viewModel.uiState.test {