Merge branch 'custom-list-added-animation-causes-custom-list-description-droid-901'

3 files changed, 166 insertions, 78 deletions

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 d98292ccd9..d72d183202 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,6 @@ class CustomListActionUseCaseTest {
                 RelayItem.Location.Country(
                     id = locationId,
                     name = locationName,
-                    expanded = false,
                     cities = emptyList()
                 )
             )
@@ -151,7 +150,7 @@ class CustomListActionUseCaseTest {
         val action = CustomListAction.Delete(id = customListId)
         val expectedResult =
             Deleted(undo = action.not(name = name, locations = listOf(location))).right()
-        every { mockLocation.countryCode } returns location.countryCode
+        every { mockLocation.code } returns location.code
         coEvery { mockCustomListsRepository.deleteCustomList(id = customListId) } returns
             Unit.right()
         coEvery { mockCustomListsRepository.getCustomListById(customListId) } returns
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 4879031ce7..d2eaedd8c2 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
@@ -13,7 +13,9 @@ import net.mullvad.mullvadvpn.compose.communication.CustomListAction
 import net.mullvad.mullvadvpn.compose.communication.LocationsChanged
 import net.mullvad.mullvadvpn.compose.screen.CustomListLocationsNavArgs
 import net.mullvad.mullvadvpn.compose.state.CustomListLocationsUiState
+import net.mullvad.mullvadvpn.compose.state.RelayLocationListItem
 import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
+import net.mullvad.mullvadvpn.lib.common.test.assertLists
 import net.mullvad.mullvadvpn.lib.model.CustomList
 import net.mullvad.mullvadvpn.lib.model.CustomListId
 import net.mullvad.mullvadvpn.lib.model.CustomListName
@@ -23,6 +25,7 @@ import net.mullvad.mullvadvpn.lib.model.Provider
 import net.mullvad.mullvadvpn.lib.model.ProviderId
 import net.mullvad.mullvadvpn.lib.model.RelayItem
 import net.mullvad.mullvadvpn.relaylist.descendants
+import net.mullvad.mullvadvpn.relaylist.withDescendants
 import net.mullvad.mullvadvpn.repository.RelayListRepository
 import net.mullvad.mullvadvpn.usecase.customlists.CustomListActionUseCase
 import net.mullvad.mullvadvpn.usecase.customlists.CustomListRelayItemsUseCase
@@ -66,15 +69,20 @@ class CustomListLocationsViewModelTest {
     fun `when selected locations is not null and relay countries is not empty should return ui state content`() =
         runTest {
             // Arrange
-            val expectedList = DUMMY_COUNTRIES
+            val expectedList =
+                DUMMY_COUNTRIES.map {
+                    RelayLocationListItem(
+                        item = it,
+                        depth = it.toDepth(),
+                        checked = false,
+                        expanded = false
+                    )
+                }
             val customListId = CustomListId("id")
             val expectedState =
-                CustomListLocationsUiState.Content.Data(
-                    newList = true,
-                    availableLocations = expectedList
-                )
+                CustomListLocationsUiState.Content.Data(newList = true, locations = expectedList)
             val viewModel = createViewModel(customListId, true)
-            relayListFlow.value = expectedList
+            relayListFlow.value = DUMMY_COUNTRIES
 
             // Act, Assert
             viewModel.uiState.test { assertEquals(expectedState, awaitItem()) }
@@ -85,8 +93,8 @@ class CustomListLocationsViewModelTest {
         // Arrange
         val expectedList = DUMMY_COUNTRIES
         val customListId = CustomListId("id")
-        val expectedSelection =
-            (DUMMY_COUNTRIES + DUMMY_COUNTRIES.flatMap { it.descendants() }).toSet()
+        val expectedSelection = (DUMMY_COUNTRIES.take(1).withDescendants()).map { it.id }
+
         val viewModel = createViewModel(customListId, true)
         relayListFlow.value = expectedList
 
@@ -95,12 +103,19 @@ class CustomListLocationsViewModelTest {
             // Check no selected
             val firstState = awaitItem()
             assertIs<CustomListLocationsUiState.Content.Data>(firstState)
-            assertEquals(emptySet<RelayItem>(), firstState.selectedLocations)
+            assertEquals(emptyList<RelayItem>(), firstState.selectedLocations())
+            // Expand country
+            viewModel.onExpand(DUMMY_COUNTRIES[0], true)
+            awaitItem()
+            // Expand city
+            viewModel.onExpand(DUMMY_COUNTRIES[0].cities[0], expand = true)
+            awaitItem()
+            // Select country
             viewModel.onRelaySelectionClick(DUMMY_COUNTRIES[0], true)
             // Check all items selected
             val secondState = awaitItem()
             assertIs<CustomListLocationsUiState.Content.Data>(secondState)
-            assertEquals(expectedSelection, secondState.selectedLocations)
+            assertLists(expectedSelection, secondState.selectedLocations())
         }
     }
 
@@ -108,25 +123,29 @@ class CustomListLocationsViewModelTest {
     fun `when deselecting child should deselect parent`() = runTest {
         // Arrange
         val expectedList = DUMMY_COUNTRIES
-        val initialSelection =
-            (DUMMY_COUNTRIES + DUMMY_COUNTRIES.flatMap { it.descendants() }).toSet()
+        val initialSelection = DUMMY_COUNTRIES.withDescendants()
+        val initialSelectionIds = initialSelection.map { it.id }
         val customListId = CustomListId("id")
-        val expectedSelection = emptySet<RelayItem>()
+        val expectedSelection = emptyList<RelayItem>()
         relayListFlow.value = expectedList
-        selectedLocationsFlow.value = initialSelection.toList()
+        selectedLocationsFlow.value = initialSelection
         val viewModel = createViewModel(customListId, true)
 
         // Act, Assert
         viewModel.uiState.test {
+            // Expand country
+            viewModel.onExpand(DUMMY_COUNTRIES[0], true)
+            // Expand city
+            viewModel.onExpand(DUMMY_COUNTRIES[0].cities[0], true)
             // Check initial selected
             val firstState = awaitItem()
             assertIs<CustomListLocationsUiState.Content.Data>(firstState)
-            assertEquals(initialSelection, firstState.selectedLocations)
+            assertEquals(initialSelectionIds, firstState.selectedLocations())
             viewModel.onRelaySelectionClick(DUMMY_COUNTRIES[0].cities[0].relays[0], false)
             // Check all items selected
             val secondState = awaitItem()
             assertIs<CustomListLocationsUiState.Content.Data>(secondState)
-            assertEquals(expectedSelection, secondState.selectedLocations)
+            assertEquals(expectedSelection, secondState.selectedLocations())
         }
     }
 
@@ -136,23 +155,28 @@ class CustomListLocationsViewModelTest {
         val expectedList = DUMMY_COUNTRIES
         val initialSelection =
             (DUMMY_COUNTRIES + DUMMY_COUNTRIES.flatMap { it.descendants() }).toSet()
+        val initialSelectionIds = initialSelection.map { it.id }
         val customListId = CustomListId("id")
-        val expectedSelection = emptySet<RelayItem>()
+        val expectedSelection = emptyList<RelayItem>()
         relayListFlow.value = expectedList
         selectedLocationsFlow.value = initialSelection.toList()
         val viewModel = createViewModel(customListId, true)
 
         // Act, Assert
         viewModel.uiState.test {
+            // Expand country
+            viewModel.onExpand(DUMMY_COUNTRIES[0], true)
+            // Expand city
+            viewModel.onExpand(DUMMY_COUNTRIES[0].cities[0], true)
             // Check initial selected
             val firstState = awaitItem()
             assertIs<CustomListLocationsUiState.Content.Data>(firstState)
-            assertEquals(initialSelection, firstState.selectedLocations)
+            assertEquals(initialSelectionIds, firstState.selectedLocations())
             viewModel.onRelaySelectionClick(DUMMY_COUNTRIES[0], false)
             // Check all items selected
             val secondState = awaitItem()
             assertIs<CustomListLocationsUiState.Content.Data>(secondState)
-            assertEquals(expectedSelection, secondState.selectedLocations)
+            assertEquals(expectedSelection, secondState.selectedLocations())
         }
     }
 
@@ -161,21 +185,27 @@ class CustomListLocationsViewModelTest {
         // Arrange
         val expectedList = DUMMY_COUNTRIES
         val customListId = CustomListId("id")
-        val expectedSelection = DUMMY_COUNTRIES[0].cities[0].relays.toSet()
+        val expectedSelection = DUMMY_COUNTRIES[0].cities[0].relays.map { it.id }
         val viewModel = createViewModel(customListId, true)
         relayListFlow.value = expectedList
 
         // Act, Assert
         viewModel.uiState.test {
+            awaitItem() // Initial item
+            // Expand country
+            viewModel.onExpand(DUMMY_COUNTRIES[0], true)
+            awaitItem()
+            // Expand city
+            viewModel.onExpand(DUMMY_COUNTRIES[0].cities[0], true)
             // Check no selected
             val firstState = awaitItem()
             assertIs<CustomListLocationsUiState.Content.Data>(firstState)
-            assertEquals(emptySet<RelayItem>(), firstState.selectedLocations)
+            assertEquals(emptyList<RelayItem>(), firstState.selectedLocations())
             viewModel.onRelaySelectionClick(DUMMY_COUNTRIES[0].cities[0].relays[0], true)
             // Check all items selected
             val secondState = awaitItem()
             assertIs<CustomListLocationsUiState.Content.Data>(secondState)
-            assertEquals(expectedSelection, secondState.selectedLocations)
+            assertEquals(expectedSelection, secondState.selectedLocations())
         }
     }
 
@@ -235,18 +265,26 @@ class CustomListLocationsViewModelTest {
         )
     }
 
+    private fun CustomListLocationsUiState.Content.Data.selectedLocations() =
+        this.locations.filter { it.checked }.map { it.item.id }
+
+    private fun RelayItem.Location.toDepth() =
+        when (this) {
+            is RelayItem.Location.Country -> 0
+            is RelayItem.Location.City -> 1
+            is RelayItem.Location.Relay -> 2
+        }
+
     companion object {
         private val DUMMY_COUNTRIES =
             listOf(
                 RelayItem.Location.Country(
                     name = "Sweden",
                     id = GeoLocationId.Country("SE"),
-                    expanded = false,
                     cities =
                         listOf(
                             RelayItem.Location.City(
                                 name = "Gothenburg",
-                                expanded = false,
                                 id = GeoLocationId.City(GeoLocationId.Country("SE"), "GBG"),
                                 relays =
                                     listOf(
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 28f52ba261..d15e460e5d 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
@@ -11,15 +11,18 @@ 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.LocationsChanged
+import net.mullvad.mullvadvpn.compose.state.RelayListItem
 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
@@ -29,13 +32,14 @@ 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.relaylist.descendants
-import net.mullvad.mullvadvpn.relaylist.filterOnSearchTerm
+import net.mullvad.mullvadvpn.repository.CustomListsRepository
 import net.mullvad.mullvadvpn.repository.RelayListFilterRepository
 import net.mullvad.mullvadvpn.repository.RelayListRepository
 import net.mullvad.mullvadvpn.usecase.AvailableProvidersUseCase
 import net.mullvad.mullvadvpn.usecase.FilteredRelayListUseCase
 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.BeforeEach
 import org.junit.jupiter.api.Test
@@ -47,9 +51,11 @@ class SelectLocationViewModelTest {
     private val mockRelayListFilterRepository: RelayListFilterRepository = mockk()
     private val mockAvailableProvidersUseCase: AvailableProvidersUseCase = mockk(relaxed = true)
     private val mockCustomListActionUseCase: CustomListActionUseCase = mockk(relaxed = true)
-    private val mockCustomListsRelayItemUseCase: CustomListsRelayItemUseCase = mockk()
+    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 lateinit var viewModel: SelectLocationViewModel
 
@@ -58,7 +64,9 @@ class SelectLocationViewModelTest {
     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 customRelayListItems = MutableStateFlow<List<RelayItem.CustomList>>(emptyList())
+    private val filteredCustomRelayListItems =
+        MutableStateFlow<List<RelayItem.CustomList>>(emptyList())
+    private val customListsRelayItem = MutableStateFlow<List<RelayItem.CustomList>>(emptyList())
 
     @BeforeEach
     fun setup() {
@@ -68,7 +76,8 @@ class SelectLocationViewModelTest {
         every { mockAvailableProvidersUseCase() } returns allProviders
         every { mockRelayListRepository.selectedLocation } returns selectedRelayItemFlow
         every { mockFilteredRelayListUseCase() } returns filteredRelayList
-        every { mockCustomListsRelayItemUseCase() } returns customRelayListItems
+        every { mockFilteredCustomListRelayItemsUseCase() } returns filteredCustomRelayListItems
+        every { mockCustomListsRelayItemUseCase() } returns customListsRelayItem
 
         mockkStatic(RELAY_LIST_EXTENSIONS)
         mockkStatic(RELAY_ITEM_EXTENSIONS)
@@ -77,10 +86,12 @@ class SelectLocationViewModelTest {
             SelectLocationViewModel(
                 relayListFilterRepository = mockRelayListFilterRepository,
                 availableProvidersUseCase = mockAvailableProvidersUseCase,
-                customListsRelayItemUseCase = mockCustomListsRelayItemUseCase,
+                filteredCustomListRelayItemsUseCase = mockFilteredCustomListRelayItemsUseCase,
                 customListActionUseCase = mockCustomListActionUseCase,
                 filteredRelayListUseCase = mockFilteredRelayListUseCase,
-                relayListRepository = mockRelayListRepository
+                relayListRepository = mockRelayListRepository,
+                customListsRepository = mockCustomListsRepository,
+                customListsRelayItemUseCase = mockCustomListsRelayItemUseCase,
             )
     }
 
@@ -96,41 +107,50 @@ class SelectLocationViewModelTest {
     }
 
     @Test
-    fun `given relayListWithSelection emits update uiState should contain new update`() = runTest {
+    fun `given filteredRelayList emits update uiState should contain new update`() = runTest {
         // Arrange
-        val mockCountries = listOf<RelayItem.Location.Country>(mockk(), mockk())
-        val selectedItem: RelayItemId = mockk()
-        every { mockCountries.filterOnSearchTerm(any(), selectedItem) } returns mockCountries
-        filteredRelayList.value = mockCountries
-        selectedRelayItemFlow.value = Constraint.Only(selectedItem)
+        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(mockCountries, actualState.countries)
-            assertEquals(selectedItem, actualState.selectedItem)
+            assertLists(
+                testCountries.map { it.id },
+                actualState.relayListItems.mapNotNull { it.relayItemId() }
+            )
+            assertTrue(
+                actualState.relayListItems
+                    .filterIsInstance<RelayListItem.SelectableItem>()
+                    .first { it.relayItemId() == selectedId }
+                    .isSelected
+            )
         }
     }
 
     @Test
-    fun `given relayListWithSelection emits update with no selections selectedItem should be null`() =
-        runTest {
-            // Arrange
-            val mockCountries = listOf<RelayItem.Location.Country>(mockk(), mockk())
-            val selectedItem: RelayItemId? = null
-            every { mockCountries.filterOnSearchTerm(any(), selectedItem) } returns mockCountries
-            filteredRelayList.value = mockCountries
-            selectedRelayItemFlow.value = Constraint.Any
+    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(mockCountries, actualState.countries)
-                assertEquals(selectedItem, actualState.selectedItem)
-            }
+        // 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 {
@@ -153,15 +173,8 @@ class SelectLocationViewModelTest {
     @Test
     fun `on onSearchTermInput call uiState should emit with filtered countries`() = runTest {
         // Arrange
-        val mockCustomList = listOf<RelayItem.CustomList>(mockk(relaxed = true))
-        val mockCountries = listOf<RelayItem.Location.Country>(mockk(), mockk())
-        val selectedItem: RelayItemId? = null
-        val mockRelayList: List<RelayItem.Location.Country> = mockk(relaxed = true)
-        val mockSearchString = "SEARCH"
-        every { mockRelayList.filterOnSearchTerm(mockSearchString, selectedItem) } returns
-            mockCountries
-        every { mockCustomList.filterOnSearchTerm(mockSearchString) } returns mockCustomList
-        filteredRelayList.value = mockRelayList
+        val mockSearchString = "got"
+        filteredRelayList.value = testCountries
         selectedRelayItemFlow.value = Constraint.Any
 
         // Act, Assert
@@ -172,25 +185,26 @@ class SelectLocationViewModelTest {
             // Update search string
             viewModel.onSearchTermInput(mockSearchString)
 
-            // Assert
+            // We get some unnecessary emissions for now
+            awaitItem()
+            awaitItem()
+            awaitItem()
+
             val actualState = awaitItem()
             assertIs<SelectLocationUiState.Content>(actualState)
-            assertLists(mockCountries, actualState.countries)
-            assertEquals(selectedItem, actualState.selectedItem)
+            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
-        val mockCustomList = listOf<RelayItem.CustomList>(mockk(relaxed = true))
-        val mockCountries = emptyList<RelayItem.Location.Country>()
-        val selectedItem: RelayItemId? = null
-        val mockRelayList: List<RelayItem.Location.Country> = mockk(relaxed = true)
+        filteredRelayList.value = testCountries
         val mockSearchString = "SEARCH"
-        every { mockRelayList.filterOnSearchTerm(mockSearchString, selectedItem) } returns
-            mockCountries
-        every { mockCustomList.filterOnSearchTerm(mockSearchString) } returns mockCustomList
 
         // Act, Assert
         viewModel.uiState.test {
@@ -200,10 +214,17 @@ class SelectLocationViewModelTest {
             // Update search string
             viewModel.onSearchTermInput(mockSearchString)
 
+            // We get some unnecessary emissions for now
+            awaitItem()
+            awaitItem()
+
             // Assert
             val actualState = awaitItem()
             assertIs<SelectLocationUiState.Content>(actualState)
-            assertEquals(mockSearchString, actualState.searchTerm)
+            assertEquals(
+                listOf(RelayListItem.LocationsEmptyText(mockSearchString)),
+                actualState.relayListItems
+            )
         }
     }
 
@@ -259,10 +280,13 @@ class SelectLocationViewModelTest {
         }
         val customList =
             RelayItem.CustomList(
-                id = CustomListId("1"),
-                customListName = CustomListName.fromString("custom"),
+                customList =
+                    CustomList(
+                        id = CustomListId("1"),
+                        name = CustomListName.fromString("custom"),
+                        locations = emptyList()
+                    ),
                 locations = emptyList(),
-                expanded = false
             )
         coEvery { mockCustomListActionUseCase(any<CustomListAction.UpdateLocations>()) } returns
             expectedResult.right()
@@ -276,6 +300,17 @@ class SelectLocationViewModelTest {
         }
     }
 
+    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"
@@ -283,5 +318,21 @@ 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>(
+                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())
+            )
     }
 }