Fix unit tests

6 files changed, 145 insertions, 28 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 d72d183202..3528d1bfa2 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
@@ -174,7 +174,13 @@ class CustomListActionUseCaseTest {
         val customList = CustomList(id = customListId, name = name, locations = oldLocations)
         val action = CustomListAction.UpdateLocations(id = customListId, locations = newLocations)
         val expectedResult =
-            LocationsChanged(name = name, undo = action.not(locations = oldLocations)).right()
+            LocationsChanged(
+                    id = customListId,
+                    name = name,
+                    locations = newLocations,
+                    oldLocations = oldLocations,
+                )
+                .right()
         coEvery { mockCustomListsRepository.getCustomListById(customListId) } returns
             customList.right()
 
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/CreateCustomListDialogViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/CreateCustomListDialogViewModelTest.kt
index 2ab6e6267b..fee63ae804 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/CreateCustomListDialogViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/CreateCustomListDialogViewModelTest.kt
@@ -11,6 +11,7 @@ import kotlin.test.assertIs
 import kotlinx.coroutines.test.runTest
 import net.mullvad.mullvadvpn.compose.communication.Created
 import net.mullvad.mullvadvpn.compose.communication.CustomListAction
+import net.mullvad.mullvadvpn.compose.communication.CustomListActionResultData
 import net.mullvad.mullvadvpn.compose.dialog.CreateCustomListNavArgs
 import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
 import net.mullvad.mullvadvpn.lib.model.CustomListAlreadyExists
@@ -32,16 +33,28 @@ class CreateCustomListDialogViewModelTest {
     fun `when successfully creating a list with locations should emit return with result side effect`() =
         runTest {
             // Arrange
-            val expectedResult: Created = mockk()
-            val customListName = "list"
+            val mockCreated: Created = mockk()
+            val mockUndo: CustomListAction.Delete = mockk()
+            val customListName = CustomListName.fromString("list")
+            val customListId = CustomListId("1")
+            val locationNames = listOf("locationName")
+            val expectedResult =
+                CustomListActionResultData.Success.CreatedWithLocations(
+                    customListName = customListName,
+                    locationNames = locationNames,
+                    undo = mockUndo
+                )
             val viewModel = createViewModelWithLocationCode(GeoLocationId.Country("AB"))
             coEvery { mockCustomListActionUseCase(any<CustomListAction.Create>()) } returns
-                expectedResult.right()
-            every { expectedResult.locationNames } returns listOf("locationName")
+                mockCreated.right()
+            every { mockCreated.locationNames } returns locationNames
+            every { mockCreated.name } returns customListName
+            every { mockCreated.id } returns customListId
+            every { mockCreated.undo } returns mockUndo
 
             // Act, Assert
             viewModel.uiSideEffect.test {
-                viewModel.createCustomList(customListName)
+                viewModel.createCustomList(customListName.value)
                 val sideEffect = awaitItem()
                 assertIs<CreateCustomListDialogSideEffect.ReturnWithResult>(sideEffect)
                 assertEquals(expectedResult, sideEffect.result)
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 f73dbee7a3..f453223ec7 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
@@ -10,6 +10,7 @@ import kotlin.test.assertIs
 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.screen.CustomListLocationsNavArgs
 import net.mullvad.mullvadvpn.compose.state.CustomListLocationsUiState
@@ -214,40 +215,57 @@ class CustomListLocationsViewModelTest {
     }
 
     @Test
-    fun `given new list true when saving successfully should emit close screen side effect`() =
+    fun `given new list true when saving successfully should emit return with result data`() =
         runTest {
             // Arrange
             val customListId = CustomListId("1")
+            val customListName = CustomListName.fromString("name")
             val newList = true
-            val expectedResult: LocationsChanged = mockk()
+            val locationChangedMock: LocationsChanged = mockk()
             coEvery { mockCustomListUseCase(any<CustomListAction.UpdateLocations>()) } returns
-                expectedResult.right()
+                locationChangedMock.right()
+            every { locationChangedMock.name } returns customListName
+            every { locationChangedMock.id } returns customListId
             val viewModel = createViewModel(customListId, newList)
 
             // Act, Assert
             viewModel.uiSideEffect.test {
                 viewModel.save()
                 val sideEffect = awaitItem()
-                assertIs<CustomListLocationsSideEffect.CloseScreen>(sideEffect)
+                assertIs<CustomListLocationsSideEffect.ReturnWithResultData>(sideEffect)
             }
         }
 
     @Test
-    fun `given new list false when saving successfully should emit return with result side effect`() =
+    fun `given new list false when saving successfully should emit return with result data`() =
         runTest {
             // Arrange
             val customListId = CustomListId("1")
+            val customListName = CustomListName.fromString("name")
+            val mockUndo: CustomListAction.UpdateLocations = mockk()
+            val addedLocations: List<GeoLocationId> = listOf(mockk())
+            val removedLocations: List<GeoLocationId> = listOf(mockk())
             val newList = false
-            val expectedResult: LocationsChanged = mockk()
+            val locationsChangedMock: LocationsChanged = mockk()
+            val expectedResult =
+                CustomListActionResultData.Success.LocationChanged(
+                    customListName = customListName,
+                    undo = mockUndo
+                )
             coEvery { mockCustomListUseCase(any<CustomListAction.UpdateLocations>()) } returns
-                expectedResult.right()
+                locationsChangedMock.right()
+            every { locationsChangedMock.id } returns customListId
+            every { locationsChangedMock.name } returns customListName
+            every { locationsChangedMock.addedLocations } returns addedLocations
+            every { locationsChangedMock.removedLocations } returns removedLocations
+            every { locationsChangedMock.undo } returns mockUndo
             val viewModel = createViewModel(customListId, newList)
 
             // Act, Assert
             viewModel.uiSideEffect.test {
                 viewModel.save()
                 val sideEffect = awaitItem()
-                assertIs<CustomListLocationsSideEffect.ReturnWithResult>(sideEffect)
+                assertIs<CustomListLocationsSideEffect.ReturnWithResultData>(sideEffect)
                 assertEquals(expectedResult, sideEffect.result)
             }
         }
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/DeleteCustomListConfirmationViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/DeleteCustomListConfirmationViewModelTest.kt
index 2656d1a4ff..d90deeb99e 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/DeleteCustomListConfirmationViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/DeleteCustomListConfirmationViewModelTest.kt
@@ -4,10 +4,12 @@ import app.cash.turbine.test
 import arrow.core.right
 import com.ramcosta.composedestinations.generated.navargs.toSavedStateHandle
 import io.mockk.coEvery
+import io.mockk.every
 import io.mockk.mockk
 import kotlin.test.assertIs
 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.Deleted
 import net.mullvad.mullvadvpn.compose.dialog.DeleteCustomListNavArgs
 import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
@@ -25,10 +27,16 @@ class DeleteCustomListConfirmationViewModelTest {
     @Test
     fun `when successfully deleting a list should emit return with result side effect`() = runTest {
         // Arrange
-        val expectedResult: Deleted = mockk()
+        val deleted: Deleted = mockk()
+        val customListName = CustomListName.fromString("name")
+        val undo: CustomListAction.Create = mockk()
+        val expectedResult =
+            CustomListActionResultData.Success.Deleted(customListName = customListName, undo = undo)
+        every { deleted.name } returns customListName
+        every { deleted.undo } returns undo
         val viewModel = createViewModel()
         coEvery { mockCustomListActionUseCase(any<CustomListAction.Delete>()) } returns
-            expectedResult.right()
+            deleted.right()
 
         // Act, Assert
         viewModel.uiSideEffect.test {
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/EditCustomListNameDialogViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/EditCustomListNameDialogViewModelTest.kt
index 7b4dcc0b83..faf83da7a8 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/EditCustomListNameDialogViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/EditCustomListNameDialogViewModelTest.kt
@@ -5,10 +5,12 @@ import arrow.core.left
 import arrow.core.right
 import com.ramcosta.composedestinations.generated.navargs.toSavedStateHandle
 import io.mockk.coEvery
+import io.mockk.every
 import io.mockk.mockk
 import kotlin.test.assertIs
 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.Renamed
 import net.mullvad.mullvadvpn.compose.dialog.EditCustomListNameNavArgs
 import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
@@ -28,16 +30,21 @@ class EditCustomListNameDialogViewModelTest {
     @Test
     fun `when successfully renamed list should emit return with result side effect`() = runTest {
         // Arrange
-        val expectedResult: Renamed = mockk()
+        val renamed: Renamed = mockk()
         val customListId = CustomListId("id")
-        val customListName = "list"
-        val viewModel = createViewModel(customListId, customListName)
+        val customListName = CustomListName.fromString("list")
+        val undo: CustomListAction.Rename = mockk()
+        val expectedResult =
+            CustomListActionResultData.Success.Renamed(newName = customListName, undo = undo)
+        every { renamed.name } returns customListName
+        every { renamed.undo } returns undo
+        val viewModel = createViewModel(customListId, customListName.value)
         coEvery { mockCustomListActionUseCase(any<CustomListAction.Rename>()) } returns
-            expectedResult.right()
+            renamed.right()
 
         // Act, Assert
         viewModel.uiSideEffect.test {
-            viewModel.updateCustomListName(customListName)
+            viewModel.updateCustomListName(customListName.value)
             val sideEffect = awaitItem()
             assertIs<EditCustomListNameDialogSideEffect.ReturnWithResult>(sideEffect)
             assertEquals(expectedResult, sideEffect.result)
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 d15e460e5d..993190ba2c 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
@@ -16,6 +16,7 @@ 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.SelectLocationUiState
@@ -273,9 +274,12 @@ class SelectLocationViewModelTest {
     @Test
     fun `after adding a location to a list should emit location added side effect`() = runTest {
         // Arrange
-        val expectedResult: LocationsChanged = mockk()
+        val customListId = CustomListId("1")
+        val addedLocationsId = GeoLocationId.Country("se")
+        val customListName = CustomListName.fromString("custom")
         val location: RelayItem.Location.Country = mockk {
             every { id } returns GeoLocationId.Country("se")
+            every { name } returns "Sweden"
             every { descendants() } returns emptyList()
         }
         val customList =
@@ -283,24 +287,85 @@ class SelectLocationViewModelTest {
                 customList =
                     CustomList(
                         id = CustomListId("1"),
-                        name = CustomListName.fromString("custom"),
+                        name = customListName,
                         locations = emptyList()
                     ),
                 locations = emptyList(),
             )
+        val expectedResult =
+            CustomListActionResultData.Success.LocationAdded(
+                customListName = customListName,
+                locationName = location.name,
+                undo = CustomListAction.UpdateLocations(id = customListId, locations = emptyList())
+            )
+
         coEvery { mockCustomListActionUseCase(any<CustomListAction.UpdateLocations>()) } returns
-            expectedResult.right()
+            LocationsChanged(
+                    id = customListId,
+                    name = customListName,
+                    locations = listOf(addedLocationsId),
+                    oldLocations = emptyList()
+                )
+                .right()
 
         // Act, Assert
         viewModel.uiSideEffect.test {
             viewModel.addLocationToList(item = location, customList = customList)
             val sideEffect = awaitItem()
-            assertIs<SelectLocationSideEffect.LocationAddedToCustomList>(sideEffect)
-            assertEquals(expectedResult, sideEffect.result)
+            assertIs<SelectLocationSideEffect.CustomListActionToast>(sideEffect)
+            assertEquals(expectedResult, sideEffect.resultData)
         }
     }
 
-    fun RelayListItem.relayItemId() =
+    @Test
+    fun `after removing a location from a list should emit location removed side effect`() =
+        runTest {
+            // Arrange
+            val locationName = "Sweden"
+            val customListId = CustomListId("1")
+            val removedLocationsId = GeoLocationId.Country("se")
+            val customListName = CustomListName.fromString("custom")
+            val location: RelayItem.Location.Country = mockk {
+                every { id } returns removedLocationsId
+                every { name } returns locationName
+                every { descendants() } returns emptyList()
+            }
+            val expectedResult =
+                CustomListActionResultData.Success.LocationRemoved(
+                    customListName = customListName,
+                    locationName = locationName,
+                    undo =
+                        CustomListAction.UpdateLocations(
+                            id = customListId,
+                            locations = listOf(location.id)
+                        )
+                )
+            coEvery { mockCustomListActionUseCase(any<CustomListAction.UpdateLocations>()) } returns
+                LocationsChanged(
+                        id = customListId,
+                        name = customListName,
+                        locations = emptyList(),
+                        oldLocations = listOf(removedLocationsId),
+                    )
+                    .right()
+            coEvery { mockCustomListsRepository.getCustomListById(customListId) } returns
+                CustomList(
+                        id = customListId,
+                        name = customListName,
+                        locations = listOf(removedLocationsId)
+                    )
+                    .right()
+
+            // Act, Assert
+            viewModel.uiSideEffect.test {
+                viewModel.removeLocationFromList(item = location, customListId = customListId)
+                val sideEffect = awaitItem()
+                assertIs<SelectLocationSideEffect.CustomListActionToast>(sideEffect)
+                assertEquals(expectedResult, sideEffect.resultData)
+            }
+        }
+
+    private fun RelayListItem.relayItemId() =
         when (this) {
             is RelayListItem.CustomListFooter -> null
             RelayListItem.CustomListHeader -> null
@@ -320,7 +385,7 @@ class SelectLocationViewModelTest {
             "net.mullvad.mullvadvpn.relaylist.CustomListExtensionsKt"
 
         private val testCountries =
-            listOf<RelayItem.Location.Country>(
+            listOf(
                 RelayItem.Location.Country(
                     id = GeoLocationId.Country("se"),
                     "Sweden",