Update unit tests for new relay items

3 files changed, 126 insertions, 58 deletions

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 dc9c141f28..6fe71ae443 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
@@ -16,9 +16,14 @@ class RelayNameComparatorTest {
     @Test
     fun test_compare_respect_numbers_in_name() {
         val relay9 =
-            Relay(name = "se9-wireguard", location = mockk(), locationName = "mock", active = false)
+            RelayItem.Relay(
+                name = "se9-wireguard",
+                location = mockk(),
+                locationName = "mock",
+                active = false
+            )
         val relay10 =
-            Relay(
+            RelayItem.Relay(
                 name = "se10-wireguard",
                 location = mockk(),
                 locationName = "mock",
@@ -31,9 +36,19 @@ class RelayNameComparatorTest {
     @Test
     fun test_compare_same_name() {
         val relay9a =
-            Relay(name = "se9-wireguard", location = mockk(), locationName = "mock", active = false)
+            RelayItem.Relay(
+                name = "se9-wireguard",
+                location = mockk(),
+                locationName = "mock",
+                active = false
+            )
         val relay9b =
-            Relay(name = "se9-wireguard", location = mockk(), locationName = "mock", active = false)
+            RelayItem.Relay(
+                name = "se9-wireguard",
+                location = mockk(),
+                locationName = "mock",
+                active = false
+            )
 
         assertTrue(RelayNameComparator.compare(relay9a, relay9b) == 0)
         assertTrue(RelayNameComparator.compare(relay9b, relay9a) == 0)
@@ -42,11 +57,13 @@ class RelayNameComparatorTest {
     @Test
     fun test_compare_only_numbers_in_name() {
         val relay001 =
-            Relay(name = "001", location = mockk(), locationName = "mock", active = false)
-        val relay1 = Relay(name = "1", location = mockk(), locationName = "mock", active = false)
-        val relay3 = Relay(name = "3", location = mockk(), locationName = "mock", active = false)
+            RelayItem.Relay(name = "001", location = mockk(), locationName = "mock", active = false)
+        val relay1 =
+            RelayItem.Relay(name = "1", location = mockk(), locationName = "mock", active = false)
+        val relay3 =
+            RelayItem.Relay(name = "3", location = mockk(), locationName = "mock", active = false)
         val relay100 =
-            Relay(name = "100", location = mockk(), locationName = "mock", active = false)
+            RelayItem.Relay(name = "100", location = mockk(), locationName = "mock", active = false)
 
         relay001 assertOrderBothDirection relay1
         relay001 assertOrderBothDirection relay3
@@ -57,9 +74,19 @@ class RelayNameComparatorTest {
     @Test
     fun test_compare_without_numbers_in_name() {
         val relay9a =
-            Relay(name = "se-wireguard", location = mockk(), locationName = "mock", active = false)
+            RelayItem.Relay(
+                name = "se-wireguard",
+                location = mockk(),
+                locationName = "mock",
+                active = false
+            )
         val relay9b =
-            Relay(name = "se-wireguard", location = mockk(), locationName = "mock", active = false)
+            RelayItem.Relay(
+                name = "se-wireguard",
+                location = mockk(),
+                locationName = "mock",
+                active = false
+            )
 
         assertTrue(RelayNameComparator.compare(relay9a, relay9b) == 0)
         assertTrue(RelayNameComparator.compare(relay9b, relay9a) == 0)
@@ -68,14 +95,14 @@ class RelayNameComparatorTest {
     @Test
     fun test_compare_with_trailing_zeros_in_name() {
         val relay001 =
-            Relay(
+            RelayItem.Relay(
                 name = "se001-wireguard",
                 location = mockk(),
                 locationName = "mock",
                 active = false
             )
         val relay005 =
-            Relay(
+            RelayItem.Relay(
                 name = "se005-wireguard",
                 location = mockk(),
                 locationName = "mock",
@@ -88,13 +115,28 @@ class RelayNameComparatorTest {
     @Test
     fun test_compare_prefix_and_numbers() {
         val relayAr2 =
-            Relay(name = "ar2-wireguard", location = mockk(), locationName = "mock", active = false)
+            RelayItem.Relay(
+                name = "ar2-wireguard",
+                location = mockk(),
+                locationName = "mock",
+                active = false
+            )
         val relayAr8 =
-            Relay(name = "ar8-wireguard", location = mockk(), locationName = "mock", active = false)
+            RelayItem.Relay(
+                name = "ar8-wireguard",
+                location = mockk(),
+                locationName = "mock",
+                active = false
+            )
         val relaySe5 =
-            Relay(name = "se5-wireguard", location = mockk(), locationName = "mock", active = false)
+            RelayItem.Relay(
+                name = "se5-wireguard",
+                location = mockk(),
+                locationName = "mock",
+                active = false
+            )
         val relaySe10 =
-            Relay(
+            RelayItem.Relay(
                 name = "se10-wireguard",
                 location = mockk(),
                 locationName = "mock",
@@ -109,9 +151,19 @@ class RelayNameComparatorTest {
     @Test
     fun test_compare_suffix_and_numbers() {
         val relay2c =
-            Relay(name = "se2-cloud", location = mockk(), locationName = "mock", active = false)
+            RelayItem.Relay(
+                name = "se2-cloud",
+                location = mockk(),
+                locationName = "mock",
+                active = false
+            )
         val relay2w =
-            Relay(name = "se2-wireguard", location = mockk(), locationName = "mock", active = false)
+            RelayItem.Relay(
+                name = "se2-wireguard",
+                location = mockk(),
+                locationName = "mock",
+                active = false
+            )
 
         relay2c assertOrderBothDirection relay2w
     }
@@ -119,9 +171,14 @@ class RelayNameComparatorTest {
     @Test
     fun test_compare_different_length() {
         val relay22a =
-            Relay(name = "se22", location = mockk(), locationName = "mock", active = false)
+            RelayItem.Relay(
+                name = "se22",
+                location = mockk(),
+                locationName = "mock",
+                active = false
+            )
         val relay22b =
-            Relay(
+            RelayItem.Relay(
                 name = "se22-wireguard",
                 location = mockk(),
                 locationName = "mock",
@@ -131,7 +188,7 @@ class RelayNameComparatorTest {
         relay22a assertOrderBothDirection relay22b
     }
 
-    private infix fun Relay.assertOrderBothDirection(other: Relay) {
+    private infix fun RelayItem.Relay.assertOrderBothDirection(other: RelayItem.Relay) {
         assertTrue(RelayNameComparator.compare(this, other) < 0)
         assertTrue(RelayNameComparator.compare(other, this) > 0)
     }
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/ConnectViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/ConnectViewModelTest.kt
index a1e3ce57e5..57eeef5ca4 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/ConnectViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/ConnectViewModelTest.kt
@@ -22,7 +22,6 @@ import net.mullvad.mullvadvpn.model.AccountExpiry
 import net.mullvad.mullvadvpn.model.DeviceState
 import net.mullvad.mullvadvpn.model.GeoIpLocation
 import net.mullvad.mullvadvpn.model.TunnelState
-import net.mullvad.mullvadvpn.relaylist.RelayCountry
 import net.mullvad.mullvadvpn.relaylist.RelayItem
 import net.mullvad.mullvadvpn.repository.AccountRepository
 import net.mullvad.mullvadvpn.repository.DeviceRepository
@@ -96,7 +95,7 @@ class ConnectViewModelTest {
         EventNotifier<TunnelState>(TunnelState.Disconnected())
 
     // Flows
-    private val selectedRelayFlow = MutableStateFlow<RelayItem?>(null)
+    private val selectedRelayItemFlow = MutableStateFlow<RelayItem?>(null)
 
     // Out Of Time Use Case
     private val outOfTimeUseCase: OutOfTimeUseCase = mockk()
@@ -131,7 +130,7 @@ class ConnectViewModelTest {
         every { mockAppVersionInfoCache.onUpdate = any() } answers {}
 
         // Flows
-        every { mockRelayListUseCase.selectedRelayItem() } returns selectedRelayFlow
+        every { mockRelayListUseCase.selectedRelayItem() } returns selectedRelayItemFlow
 
         every { outOfTimeUseCase.isOutOfTime() } returns outOfTimeViewFlow
         viewModel =
@@ -188,17 +187,17 @@ class ConnectViewModelTest {
     }
 
     @Test
-    fun testRelayItemUpdate() = runTest {
-        val relayTestItem =
-            RelayCountry(name = "Name", code = "Code", expanded = false, cities = emptyList())
-        selectedRelayFlow.value = relayTestItem
+    fun testSelectedLocationUpdate() = runTest {
+        val selectedRelayItem =
+            RelayItem.Country(name = "Name", code = "Code", expanded = false, cities = emptyList())
+        selectedRelayItemFlow.value = selectedRelayItem
 
         viewModel.uiState.test {
             assertEquals(ConnectUiState.INITIAL, awaitItem())
             serviceConnectionState.value =
                 ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
             val result = awaitItem()
-            assertEquals(relayTestItem, result.relayLocation)
+            assertEquals(selectedRelayItem, result.selectedRelayItem)
         }
     }
 
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 fc6408d8ab..4224563d3b 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
@@ -20,13 +20,14 @@ import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
 import net.mullvad.mullvadvpn.lib.common.test.assertLists
 import net.mullvad.mullvadvpn.model.Constraint
 import net.mullvad.mullvadvpn.model.GeographicLocationConstraint
+import net.mullvad.mullvadvpn.model.LocationConstraint
 import net.mullvad.mullvadvpn.model.Ownership
 import net.mullvad.mullvadvpn.model.Providers
 import net.mullvad.mullvadvpn.relaylist.Provider
-import net.mullvad.mullvadvpn.relaylist.RelayCountry
 import net.mullvad.mullvadvpn.relaylist.RelayItem
 import net.mullvad.mullvadvpn.relaylist.RelayList
 import net.mullvad.mullvadvpn.relaylist.filterOnSearchTerm
+import net.mullvad.mullvadvpn.relaylist.toLocationConstraint
 import net.mullvad.mullvadvpn.ui.serviceconnection.ConnectionProxy
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManager
 import net.mullvad.mullvadvpn.ui.serviceconnection.connectionProxy
@@ -43,7 +44,8 @@ class SelectLocationViewModelTest {
     private val mockRelayListFilterUseCase: RelayListFilterUseCase = mockk(relaxed = true)
     private val mockServiceConnectionManager: ServiceConnectionManager = mockk()
     private lateinit var viewModel: SelectLocationViewModel
-    private val relayListWithSelectionFlow = MutableStateFlow(RelayList(emptyList(), null))
+    private val relayListWithSelectionFlow =
+        MutableStateFlow(RelayList(emptyList(), emptyList(), null))
     private val mockRelayListUseCase: RelayListUseCase = mockk()
     private val selectedOwnership = MutableStateFlow<Constraint<Ownership>>(Constraint.Any())
     private val selectedProvider = MutableStateFlow<Constraint<Providers>>(Constraint.Any())
@@ -60,6 +62,7 @@ class SelectLocationViewModelTest {
 
         mockkStatic(SERVICE_CONNECTION_MANAGER_EXTENSIONS)
         mockkStatic(RELAY_LIST_EXTENSIONS)
+        mockkStatic(RELAY_ITEM_EXTENSIONS)
         viewModel =
             SelectLocationViewModel(
                 mockServiceConnectionManager,
@@ -82,10 +85,11 @@ class SelectLocationViewModelTest {
     @Test
     fun testUpdateLocations() = runTest {
         // Arrange
-        val mockCountries = listOf<RelayCountry>(mockk(), mockk())
-        val selectedRelay: RelayItem = mockk()
-        every { mockCountries.filterOnSearchTerm(any(), selectedRelay) } returns mockCountries
-        relayListWithSelectionFlow.value = RelayList(mockCountries, selectedRelay)
+        val mockCountries = listOf<RelayItem.Country>(mockk(), mockk())
+        val mockCustomList = listOf<RelayItem.CustomList>(mockk())
+        val selectedItem: RelayItem = mockk()
+        every { mockCountries.filterOnSearchTerm(any(), selectedItem) } returns mockCountries
+        relayListWithSelectionFlow.value = RelayList(mockCustomList, mockCountries, selectedItem)
 
         // Act, Assert
         viewModel.uiState.test {
@@ -97,8 +101,8 @@ class SelectLocationViewModelTest {
                 (actualState.relayListState as RelayListState.RelayList).countries
             )
             assertEquals(
-                selectedRelay,
-                (actualState.relayListState as RelayListState.RelayList).selectedRelay
+                selectedItem,
+                (actualState.relayListState as RelayListState.RelayList).selectedItem
             )
         }
     }
@@ -106,10 +110,11 @@ class SelectLocationViewModelTest {
     @Test
     fun testUpdateLocationsNoSelectedRelay() = runTest {
         // Arrange
-        val mockCountries = listOf<RelayCountry>(mockk(), mockk())
-        val selectedRelay: RelayItem? = null
-        every { mockCountries.filterOnSearchTerm(any(), selectedRelay) } returns mockCountries
-        relayListWithSelectionFlow.value = RelayList(mockCountries, selectedRelay)
+        val mockCustomList = listOf<RelayItem.CustomList>(mockk())
+        val mockCountries = listOf<RelayItem.Country>(mockk(), mockk())
+        val selectedItem: RelayItem? = null
+        every { mockCountries.filterOnSearchTerm(any(), selectedItem) } returns mockCountries
+        relayListWithSelectionFlow.value = RelayList(mockCustomList, mockCountries, selectedItem)
 
         // Act, Assert
         viewModel.uiState.test {
@@ -121,8 +126,8 @@ class SelectLocationViewModelTest {
                 (actualState.relayListState as RelayListState.RelayList).countries
             )
             assertEquals(
-                selectedRelay,
-                (actualState.relayListState as RelayListState.RelayList).selectedRelay
+                selectedItem,
+                (actualState.relayListState as RelayListState.RelayList).selectedItem
             )
         }
     }
@@ -130,12 +135,15 @@ class SelectLocationViewModelTest {
     @Test
     fun testSelectRelayAndClose() = runTest {
         // Arrange
-        val mockRelayItem: RelayItem = mockk()
+        val mockRelayItem: RelayItem.Country = mockk()
         val mockLocation: GeographicLocationConstraint.Country = mockk(relaxed = true)
+        val mockLocationConstraint: LocationConstraint = mockk()
         val connectionProxyMock: ConnectionProxy = mockk(relaxUnitFun = true)
         every { mockRelayItem.location } returns mockLocation
         every { mockServiceConnectionManager.connectionProxy() } returns connectionProxyMock
-        every { mockRelayListUseCase.updateSelectedRelayLocation(mockLocation) } returns Unit
+        every { mockRelayListUseCase.updateSelectedRelayLocation(mockLocationConstraint) } returns
+            Unit
+        every { mockRelayItem.toLocationConstraint() } returns mockLocationConstraint
 
         // Act, Assert
         viewModel.uiSideEffect.test {
@@ -144,7 +152,7 @@ class SelectLocationViewModelTest {
             assertEquals(SelectLocationSideEffect.CloseScreen, awaitItem())
             verify {
                 connectionProxyMock.connect()
-                mockRelayListUseCase.updateSelectedRelayLocation(mockLocation)
+                mockRelayListUseCase.updateSelectedRelayLocation(mockLocationConstraint)
             }
         }
     }
@@ -152,13 +160,14 @@ class SelectLocationViewModelTest {
     @Test
     fun testFilterRelay() = runTest {
         // Arrange
-        val mockCountries = listOf<RelayCountry>(mockk(), mockk())
-        val selectedRelay: RelayItem? = null
-        val mockRelayList: List<RelayCountry> = mockk(relaxed = true)
+        val mockCustomList = listOf<RelayItem.CustomList>(mockk())
+        val mockCountries = listOf<RelayItem.Country>(mockk(), mockk())
+        val selectedItem: RelayItem? = null
+        val mockRelayList: List<RelayItem.Country> = mockk(relaxed = true)
         val mockSearchString = "SEARCH"
-        every { mockRelayList.filterOnSearchTerm(mockSearchString, selectedRelay) } returns
+        every { mockRelayList.filterOnSearchTerm(mockSearchString, selectedItem) } returns
             mockCountries
-        relayListWithSelectionFlow.value = RelayList(mockRelayList, selectedRelay)
+        relayListWithSelectionFlow.value = RelayList(mockCustomList, mockRelayList, selectedItem)
 
         // Act, Assert
         viewModel.uiState.test {
@@ -177,8 +186,8 @@ class SelectLocationViewModelTest {
                 (actualState.relayListState as RelayListState.RelayList).countries
             )
             assertEquals(
-                selectedRelay,
-                (actualState.relayListState as RelayListState.RelayList).selectedRelay
+                selectedItem,
+                (actualState.relayListState as RelayListState.RelayList).selectedItem
             )
         }
     }
@@ -186,13 +195,14 @@ class SelectLocationViewModelTest {
     @Test
     fun testFilterNotFound() = runTest {
         // Arrange
-        val mockCountries = emptyList<RelayCountry>()
-        val selectedRelay: RelayItem? = null
-        val mockRelayList: List<RelayCountry> = mockk(relaxed = true)
+        val mockCustomList = listOf<RelayItem.CustomList>(mockk())
+        val mockCountries = emptyList<RelayItem.Country>()
+        val selectedItem: RelayItem? = null
+        val mockRelayList: List<RelayItem.Country> = mockk(relaxed = true)
         val mockSearchString = "SEARCH"
-        every { mockRelayList.filterOnSearchTerm(mockSearchString, selectedRelay) } returns
+        every { mockRelayList.filterOnSearchTerm(mockSearchString, selectedItem) } returns
             mockCountries
-        relayListWithSelectionFlow.value = RelayList(mockRelayList, selectedRelay)
+        relayListWithSelectionFlow.value = RelayList(mockCustomList, mockRelayList, selectedItem)
 
         // Act, Assert
         viewModel.uiState.test {
@@ -250,5 +260,7 @@ class SelectLocationViewModelTest {
             "net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManagerExtensionsKt"
         private const val RELAY_LIST_EXTENSIONS =
             "net.mullvad.mullvadvpn.relaylist.RelayListExtensionsKt"
+        private const val RELAY_ITEM_EXTENSIONS =
+            "net.mullvad.mullvadvpn.relaylist.RelayItemExtensionsKt"
     }
 }