Add custom lists unit tests

9 files changed, 1208 insertions, 35 deletions

diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/repository/CustomListsRepositoryTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/repository/CustomListsRepositoryTest.kt
new file mode 100644
index 0000000000..129d921c36
--- /dev/null
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/repository/CustomListsRepositoryTest.kt
@@ -0,0 +1,268 @@
+package net.mullvad.mullvadvpn.repository
+
+import io.mockk.every
+import io.mockk.mockk
+import io.mockk.mockkStatic
+import io.mockk.verify
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.test.runTest
+import net.mullvad.mullvadvpn.lib.ipc.Event
+import net.mullvad.mullvadvpn.lib.ipc.MessageHandler
+import net.mullvad.mullvadvpn.lib.ipc.Request
+import net.mullvad.mullvadvpn.lib.ipc.events
+import net.mullvad.mullvadvpn.model.CreateCustomListResult
+import net.mullvad.mullvadvpn.model.CustomList
+import net.mullvad.mullvadvpn.model.CustomListsError
+import net.mullvad.mullvadvpn.model.GeographicLocationConstraint
+import net.mullvad.mullvadvpn.model.RelayList
+import net.mullvad.mullvadvpn.model.Settings
+import net.mullvad.mullvadvpn.model.UpdateCustomListResult
+import net.mullvad.mullvadvpn.relaylist.getGeographicLocationConstraintByCode
+import net.mullvad.mullvadvpn.ui.serviceconnection.RelayListListener
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertNull
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+
+class CustomListsRepositoryTest {
+    private val mockMessageHandler: MessageHandler = mockk()
+    private val mockSettingsRepository: SettingsRepository = mockk()
+    private val mockRelayListListener: RelayListListener = mockk()
+    private val customListsRepository =
+        CustomListsRepository(
+            messageHandler = mockMessageHandler,
+            settingsRepository = mockSettingsRepository,
+            relayListListener = mockRelayListListener
+        )
+
+    private val settingsFlow: MutableStateFlow<Settings?> = MutableStateFlow(null)
+    private val relayListFlow: MutableStateFlow<RelayList> = MutableStateFlow(mockk())
+
+    @BeforeEach
+    fun setup() {
+        mockkStatic(RELAY_LIST_EXTENSIONS)
+        every { mockSettingsRepository.settingsUpdates } returns settingsFlow
+        every { mockRelayListListener.relayListEvents } returns relayListFlow
+    }
+
+    @Test
+    fun `get custom list by id should return custom list when id matches custom list in settings`() {
+        // Arrange
+        val mockCustomList: CustomList = mockk()
+        val mockSettings: Settings = mockk()
+        val customListId = "1"
+        settingsFlow.value = mockSettings
+        every { mockSettings.customLists.customLists } returns arrayListOf(mockCustomList)
+        every { mockCustomList.id } returns customListId
+
+        // Act
+        val result = customListsRepository.getCustomListById(customListId)
+
+        // Assert
+        assertEquals(mockCustomList, result)
+    }
+
+    @Test
+    fun `get custom list by id should return null when id does not matches custom list in settings`() {
+        // Arrange
+        val mockCustomList: CustomList = mockk()
+        val mockSettings: Settings = mockk()
+        val customListId = "1"
+        val otherCustomListId = "2"
+        settingsFlow.value = mockSettings
+        every { mockSettings.customLists.customLists } returns arrayListOf(mockCustomList)
+        every { mockCustomList.id } returns customListId
+
+        // Act
+        val result = customListsRepository.getCustomListById(otherCustomListId)
+
+        // Assert
+        assertNull(result)
+    }
+
+    @Test
+    fun `create custom list should return Ok when creation is successful`() = runTest {
+        // Arrange
+        val customListId = "1"
+        val expectedResult = CreateCustomListResult.Ok(customListId)
+        val customListName = "CUSTOM"
+        every {
+            mockMessageHandler.trySendRequest(Request.CreateCustomList(customListName))
+        } returns true
+        every { mockMessageHandler.events<Event.CreateCustomListResultEvent>() } returns
+            flowOf(Event.CreateCustomListResultEvent(expectedResult))
+
+        // Act
+        val result = customListsRepository.createCustomList(customListName)
+
+        // Assert
+        assertEquals(expectedResult, result)
+    }
+
+    @Test
+    fun `create custom list should return lists exists when lists exists error event is received`() =
+        runTest {
+            // Arrange
+            val expectedResult = CreateCustomListResult.Error(CustomListsError.CustomListExists)
+            val customListName = "CUSTOM"
+            every {
+                mockMessageHandler.trySendRequest(Request.CreateCustomList(customListName))
+            } returns true
+            every { mockMessageHandler.events<Event.CreateCustomListResultEvent>() } returns
+                flowOf(Event.CreateCustomListResultEvent(expectedResult))
+
+            // Act
+            val result = customListsRepository.createCustomList(customListName)
+
+            // Assert
+            assertEquals(expectedResult, result)
+        }
+
+    @Test
+    fun `update custom list name should return ok when list updated event is received`() = runTest {
+        // Arrange
+        val customListId = "1"
+        val expectedResult = UpdateCustomListResult.Ok
+        val customListName = "CUSTOM"
+        val mockSettings: Settings = mockk()
+        val mockCustomList: CustomList = mockk()
+        val updatedCustomList: CustomList = mockk()
+        settingsFlow.value = mockSettings
+        every { mockCustomList.id } returns customListId
+        every { mockCustomList.copy(customListId, customListName, any()) } returns updatedCustomList
+        every {
+            mockMessageHandler.trySendRequest(Request.UpdateCustomList(updatedCustomList))
+        } returns true
+        every { mockMessageHandler.events<Event.UpdateCustomListResultEvent>() } returns
+            flowOf(Event.UpdateCustomListResultEvent(expectedResult))
+        every { mockSettings.customLists.customLists } returns arrayListOf(mockCustomList)
+
+        // Act
+        val result = customListsRepository.updateCustomListName(customListId, customListName)
+
+        // Assert
+        assertEquals(expectedResult, result)
+    }
+
+    @Test
+    fun `update custom list name should return list exists error when list exists error is received`() =
+        runTest {
+            // Arrange
+            val customListId = "1"
+            val expectedResult = UpdateCustomListResult.Error(CustomListsError.CustomListExists)
+            val customListName = "CUSTOM"
+            val mockSettings: Settings = mockk()
+            val mockCustomList: CustomList = mockk()
+            val updatedCustomList: CustomList = mockk()
+            settingsFlow.value = mockSettings
+            every { mockCustomList.id } returns customListId
+            every { mockCustomList.copy(customListId, customListName, any()) } returns
+                updatedCustomList
+            every {
+                mockMessageHandler.trySendRequest(Request.UpdateCustomList(updatedCustomList))
+            } returns true
+            every { mockMessageHandler.events<Event.UpdateCustomListResultEvent>() } returns
+                flowOf(Event.UpdateCustomListResultEvent(expectedResult))
+            every { mockSettings.customLists.customLists } returns arrayListOf(mockCustomList)
+
+            // Act
+            val result = customListsRepository.updateCustomListName(customListId, customListName)
+
+            // Assert
+            assertEquals(expectedResult, result)
+        }
+
+    @Test
+    fun `when delete custom lists is called a delete custom event should be sent`() = runTest {
+        // Arrange
+        val customListId = "1"
+        every { mockMessageHandler.trySendRequest(Request.DeleteCustomList(customListId)) } returns
+            true
+
+        // Act
+        customListsRepository.deleteCustomList(customListId)
+
+        // Assert
+        verify { mockMessageHandler.trySendRequest(Request.DeleteCustomList(customListId)) }
+    }
+
+    @Test
+    fun `update custom list locations should return ok when list exists and ok updated list event is received`() =
+        runTest {
+            // Arrange
+            val expectedResult = UpdateCustomListResult.Ok
+            val customListId = "1"
+            val customListName = "CUSTOM"
+            val locationCode = "AB"
+            val mockSettings: Settings = mockk()
+            val mockRelayList: RelayList = mockk()
+            val mockCustomList: CustomList = mockk()
+            val updatedCustomList: CustomList = mockk()
+            val mockLocationConstraint: GeographicLocationConstraint = mockk()
+            settingsFlow.value = mockSettings
+            relayListFlow.value = mockRelayList
+            every { mockCustomList.id } returns customListId
+            every { mockCustomList.name } returns customListName
+            every {
+                mockCustomList.copy(
+                    customListId,
+                    customListName,
+                    arrayListOf(mockLocationConstraint)
+                )
+            } returns updatedCustomList
+            every {
+                mockMessageHandler.trySendRequest(Request.UpdateCustomList(updatedCustomList))
+            } returns true
+            every { mockMessageHandler.events<Event.UpdateCustomListResultEvent>() } returns
+                flowOf(Event.UpdateCustomListResultEvent(expectedResult))
+            every { mockSettings.customLists.customLists } returns arrayListOf(mockCustomList)
+            every { mockRelayList.getGeographicLocationConstraintByCode(locationCode) } returns
+                mockLocationConstraint
+
+            // Act
+            val result =
+                customListsRepository.updateCustomListLocationsFromCodes(
+                    customListId,
+                    listOf(locationCode)
+                )
+
+            // Assert
+            assertEquals(expectedResult, result)
+        }
+
+    @Test
+    fun `update custom list locations should return other error when list does not exist`() =
+        runTest {
+            // Arrange
+            val expectedResult = UpdateCustomListResult.Error(CustomListsError.OtherError)
+            val mockCustomList: CustomList = mockk()
+            val mockSettings: Settings = mockk()
+            val customListId = "1"
+            val otherCustomListId = "2"
+            val locationCode = "AB"
+            val mockRelayList: RelayList = mockk()
+            val mockLocationConstraint: GeographicLocationConstraint = mockk()
+            settingsFlow.value = mockSettings
+            relayListFlow.value = mockRelayList
+            every { mockSettings.customLists.customLists } returns arrayListOf(mockCustomList)
+            every { mockCustomList.id } returns customListId
+            every { mockRelayList.getGeographicLocationConstraintByCode(locationCode) } returns
+                mockLocationConstraint
+
+            // Act
+            val result =
+                customListsRepository.updateCustomListLocationsFromCodes(
+                    otherCustomListId,
+                    listOf(locationCode)
+                )
+
+            // Assert
+            assertEquals(expectedResult, result)
+        }
+
+    companion object {
+        private const val RELAY_LIST_EXTENSIONS =
+            "net.mullvad.mullvadvpn.relaylist.RelayListExtensionsKt"
+    }
+}
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
new file mode 100644
index 0000000000..0370f23ffb
--- /dev/null
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/CustomListActionUseCaseTest.kt
@@ -0,0 +1,220 @@
+package net.mullvad.mullvadvpn.usecase
+
+import io.mockk.coEvery
+import io.mockk.every
+import io.mockk.mockk
+import io.mockk.mockkStatic
+import kotlin.test.assertIs
+import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.test.runTest
+import net.mullvad.mullvadvpn.compose.communication.CustomListAction
+import net.mullvad.mullvadvpn.compose.communication.CustomListResult
+import net.mullvad.mullvadvpn.model.CreateCustomListResult
+import net.mullvad.mullvadvpn.model.CustomList
+import net.mullvad.mullvadvpn.model.CustomListsError
+import net.mullvad.mullvadvpn.model.GeographicLocationConstraint
+import net.mullvad.mullvadvpn.model.UpdateCustomListResult
+import net.mullvad.mullvadvpn.relaylist.RelayItem
+import net.mullvad.mullvadvpn.relaylist.getRelayItemsByCodes
+import net.mullvad.mullvadvpn.repository.CustomListsRepository
+import net.mullvad.mullvadvpn.usecase.customlists.CustomListActionUseCase
+import net.mullvad.mullvadvpn.usecase.customlists.CustomListsException
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+
+class CustomListActionUseCaseTest {
+    private val mockCustomListsRepository: CustomListsRepository = mockk()
+    private val mockRelayListUseCase: RelayListUseCase = mockk()
+    private val customListActionUseCase =
+        CustomListActionUseCase(
+            customListsRepository = mockCustomListsRepository,
+            relayListUseCase = mockRelayListUseCase
+        )
+
+    @BeforeEach
+    fun setup() {
+        mockkStatic(RELAY_LIST_EXTENSIONS)
+    }
+
+    @Test
+    fun `create action should return success when ok`() = runTest {
+        // Arrange
+        val name = "test"
+        val locationCode = "AB"
+        val locationName = "Acklaba"
+        val createdId = "1"
+        val action = CustomListAction.Create(name = name, locations = listOf(locationCode))
+        val expectedResult =
+            Result.success(
+                CustomListResult.Created(
+                    id = createdId,
+                    name = name,
+                    locationName = locationName,
+                    undo = action.not(createdId)
+                )
+            )
+        val relayItem =
+            RelayItem.Country(
+                name = locationName,
+                code = locationCode,
+                expanded = false,
+                cities = emptyList()
+            )
+        val mockLocations: List<RelayItem.Country> = listOf(relayItem)
+        coEvery { mockCustomListsRepository.createCustomList(name) } returns
+            CreateCustomListResult.Ok(createdId)
+        coEvery {
+            mockCustomListsRepository.updateCustomListLocationsFromCodes(
+                createdId,
+                listOf(locationCode)
+            )
+        } returns UpdateCustomListResult.Ok
+        coEvery { mockRelayListUseCase.relayList() } returns flowOf(mockLocations)
+        every { mockLocations.getRelayItemsByCodes(listOf(locationCode)) } returns mockLocations
+
+        // Act
+        val result = customListActionUseCase.performAction(action)
+
+        // Assert
+        assertEquals(expectedResult, result)
+    }
+
+    @Test
+    fun `create action should return error when name already exists`() = runTest {
+        // Arrange
+        val name = "test"
+        val locationCode = "AB"
+        val action = CustomListAction.Create(name = name, locations = listOf(locationCode))
+        val expectedError = CustomListsError.CustomListExists
+        coEvery { mockCustomListsRepository.createCustomList(name) } returns
+            CreateCustomListResult.Error(CustomListsError.CustomListExists)
+
+        // Act
+        val result = customListActionUseCase.performAction(action)
+
+        // Assert
+        assertIs<Result<CustomListsException>>(result)
+        val exception = result.exceptionOrNull()
+        assertIs<CustomListsException>(exception)
+        assertEquals(expectedError, exception.error)
+    }
+
+    @Test
+    fun `rename action should return success when ok`() = runTest {
+        // Arrange
+        val name = "test"
+        val newName = "test2"
+        val customListId = "1"
+        val action =
+            CustomListAction.Rename(customListId = customListId, name = name, newName = newName)
+        val expectedResult = Result.success(CustomListResult.Renamed(undo = action.not()))
+        coEvery {
+            mockCustomListsRepository.updateCustomListName(id = customListId, name = newName)
+        } returns UpdateCustomListResult.Ok
+
+        // Act
+        val result = customListActionUseCase.performAction(action)
+
+        // Assert
+        assertEquals(expectedResult, result)
+    }
+
+    @Test
+    fun `rename action should return error when name already exists`() = runTest {
+        // Arrange
+        val name = "test"
+        val newName = "test2"
+        val customListId = "1"
+        val action =
+            CustomListAction.Rename(customListId = customListId, name = name, newName = newName)
+        val expectedError = CustomListsError.CustomListExists
+        coEvery {
+            mockCustomListsRepository.updateCustomListName(id = customListId, name = newName)
+        } returns UpdateCustomListResult.Error(expectedError)
+
+        // Act
+        val result = customListActionUseCase.performAction(action)
+
+        // Assert
+        assertIs<Result<CustomListsException>>(result)
+        val exception = result.exceptionOrNull()
+        assertIs<CustomListsException>(exception)
+        assertEquals(expectedError, exception.error)
+    }
+
+    @Test
+    fun `delete action should return successful with deleted list`() = runTest {
+        // Arrange
+        val mockCustomList: CustomList = mockk()
+        val mockLocation: GeographicLocationConstraint.Country = mockk()
+        val mockLocations: ArrayList<GeographicLocationConstraint> = arrayListOf(mockLocation)
+        val name = "test"
+        val customListId = "1"
+        val locationCode = "AB"
+        val action = CustomListAction.Delete(customListId = customListId)
+        val expectedResult =
+            Result.success(
+                CustomListResult.Deleted(
+                    undo = action.not(name = name, locations = listOf(locationCode))
+                )
+            )
+        every { mockCustomList.locations } returns mockLocations
+        every { mockCustomList.name } returns name
+        every { mockLocation.countryCode } returns locationCode
+        coEvery { mockCustomListsRepository.deleteCustomList(id = customListId) } returns true
+        every { mockCustomListsRepository.getCustomListById(customListId) } returns mockCustomList
+
+        // Act
+        val result = customListActionUseCase.performAction(action)
+
+        // Assert
+        assertEquals(expectedResult, result)
+    }
+
+    @Test
+    fun `update locations action should return success with changed locations`() = runTest {
+        // Arrange
+        val name = "test"
+        val oldLocationCodes = listOf("AB", "CD")
+        val newLocationCodes = listOf("EF", "GH")
+        val oldLocations: ArrayList<GeographicLocationConstraint> =
+            arrayListOf(
+                GeographicLocationConstraint.Country("AB"),
+                GeographicLocationConstraint.Country("CD")
+            )
+        val customListId = "1"
+        val customList = CustomList(id = customListId, name = name, locations = oldLocations)
+        val action =
+            CustomListAction.UpdateLocations(
+                customListId = customListId,
+                locations = newLocationCodes
+            )
+        val expectedResult =
+            Result.success(
+                CustomListResult.LocationsChanged(
+                    name = name,
+                    undo = action.not(locations = oldLocationCodes)
+                )
+            )
+        coEvery { mockCustomListsRepository.getCustomListById(customListId) } returns customList
+
+        coEvery {
+            mockCustomListsRepository.updateCustomListLocationsFromCodes(
+                customListId,
+                newLocationCodes
+            )
+        } returns UpdateCustomListResult.Ok
+
+        // Act
+        val result = customListActionUseCase.performAction(action)
+
+        // Assert
+        assertEquals(expectedResult, result)
+    }
+
+    companion object {
+        private const val RELAY_LIST_EXTENSIONS =
+            "net.mullvad.mullvadvpn.relaylist.RelayListExtensionsKt"
+    }
+}
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
new file mode 100644
index 0000000000..7b14db3ffb
--- /dev/null
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/CreateCustomListDialogViewModelTest.kt
@@ -0,0 +1,114 @@
+package net.mullvad.mullvadvpn.viewmodel
+
+import app.cash.turbine.test
+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.CustomListResult
+import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
+import net.mullvad.mullvadvpn.model.CustomListsError
+import net.mullvad.mullvadvpn.usecase.customlists.CustomListActionUseCase
+import net.mullvad.mullvadvpn.usecase.customlists.CustomListsException
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertNull
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.ExtendWith
+
+@ExtendWith(TestCoroutineRule::class)
+class CreateCustomListDialogViewModelTest {
+    private val mockCustomListActionUseCase: CustomListActionUseCase = mockk()
+
+    @Test
+    fun `when successfully creating a list with locations should emit return with result side effect`() =
+        runTest {
+            // Arrange
+            val expectedResult: CustomListResult.Created = mockk()
+            val customListName = "list"
+            val viewModel = createViewModelWithLocationCode("AB")
+            coEvery {
+                mockCustomListActionUseCase.performAction(any<CustomListAction.Create>())
+            } returns Result.success(expectedResult)
+            every { expectedResult.locationName } returns "locationName"
+
+            // Act, Assert
+            viewModel.uiSideEffect.test {
+                viewModel.createCustomList(customListName)
+                val sideEffect = awaitItem()
+                assertIs<CreateCustomListDialogSideEffect.ReturnWithResult>(sideEffect)
+                assertEquals(expectedResult, sideEffect.result)
+            }
+        }
+
+    @Test
+    fun `when successfully creating a list without locations should emit with navigate to location screen`() =
+        runTest {
+            // Arrange
+            val expectedResult: CustomListResult.Created = mockk()
+            val customListName = "list"
+            val createdId = "1"
+            val viewModel = createViewModelWithLocationCode("")
+            coEvery {
+                mockCustomListActionUseCase.performAction(any<CustomListAction.Create>())
+            } returns Result.success(expectedResult)
+            every { expectedResult.locationName } returns null
+            every { expectedResult.id } returns createdId
+
+            // Act, Assert
+            viewModel.uiSideEffect.test {
+                viewModel.createCustomList(customListName)
+                val sideEffect = awaitItem()
+                assertIs<CreateCustomListDialogSideEffect.NavigateToCustomListLocationsScreen>(
+                    sideEffect
+                )
+                assertEquals(createdId, sideEffect.customListId)
+            }
+        }
+
+    @Test
+    fun `when failing to creating a list should update ui state with error`() = runTest {
+        // Arrange
+        val expectedError = CustomListsError.CustomListExists
+        val customListName = "list"
+        val viewModel = createViewModelWithLocationCode("")
+        coEvery {
+            mockCustomListActionUseCase.performAction(any<CustomListAction.Create>())
+        } returns Result.failure(CustomListsException(expectedError))
+
+        // Act, Assert
+        viewModel.uiState.test {
+            awaitItem() // Default state
+            viewModel.createCustomList(customListName)
+            assertEquals(expectedError, awaitItem().error)
+        }
+    }
+
+    @Test
+    fun `given error state when calling clear error then should update to state without error`() =
+        runTest {
+            // Arrange
+            val expectedError = CustomListsError.CustomListExists
+            val customListName = "list"
+            val viewModel = createViewModelWithLocationCode("")
+            coEvery {
+                mockCustomListActionUseCase.performAction(any<CustomListAction.Create>())
+            } returns Result.failure(CustomListsException(expectedError))
+
+            // Act, Assert
+            viewModel.uiState.test {
+                awaitItem() // Default state
+                viewModel.createCustomList(customListName)
+                assertEquals(expectedError, awaitItem().error) // Showing error
+                viewModel.clearError()
+                assertNull(awaitItem().error)
+            }
+        }
+
+    private fun createViewModelWithLocationCode(locationCode: String) =
+        CreateCustomListDialogViewModel(
+            locationCode = locationCode,
+            customListActionUseCase = mockCustomListActionUseCase
+        )
+}
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
new file mode 100644
index 0000000000..df10ba96c4
--- /dev/null
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListLocationsViewModelTest.kt
@@ -0,0 +1,294 @@
+package net.mullvad.mullvadvpn.viewmodel
+
+import app.cash.turbine.test
+import io.mockk.coEvery
+import io.mockk.every
+import io.mockk.mockk
+import kotlin.test.assertIs
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.test.runTest
+import net.mullvad.mullvadvpn.compose.communication.CustomListAction
+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.relaylist.RelayItem
+import net.mullvad.mullvadvpn.relaylist.descendants
+import net.mullvad.mullvadvpn.usecase.RelayListUseCase
+import net.mullvad.mullvadvpn.usecase.customlists.CustomListActionUseCase
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.ExtendWith
+
+@ExtendWith(TestCoroutineRule::class)
+class CustomListLocationsViewModelTest {
+    private val mockRelayListUseCase: RelayListUseCase = mockk()
+    private val mockCustomListUseCase: CustomListActionUseCase = mockk()
+
+    private val relayListFlow = MutableStateFlow<List<RelayItem.Country>>(emptyList())
+    private val customListFlow = MutableStateFlow<List<RelayItem.CustomList>>(emptyList())
+
+    @BeforeEach
+    fun setup() {
+        every { mockRelayListUseCase.relayList() } returns relayListFlow
+        every { mockRelayListUseCase.customLists() } returns customListFlow
+    }
+
+    @Test
+    fun `given new list false state should return new list false`() = runTest {
+        // Arrange
+        val newList = false
+        val viewModel = createViewModel("id", newList)
+
+        // Act, Assert
+        viewModel.uiState.test { assertEquals(newList, awaitItem().newList) }
+    }
+
+    @Test
+    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 customListId = "id"
+            val customListName = "name"
+            val customList: RelayItem.CustomList = mockk {
+                every { id } returns customListId
+                every { name } returns customListName
+                every { locations } returns emptyList()
+            }
+            customListFlow.value = listOf(customList)
+            val expectedState =
+                CustomListLocationsUiState.Content.Data(
+                    newList = true,
+                    availableLocations = expectedList
+                )
+            val viewModel = createViewModel(customListId, true)
+            relayListFlow.value = expectedList
+
+            // Act, Assert
+            viewModel.uiState.test { assertEquals(expectedState, awaitItem()) }
+        }
+
+    @Test
+    fun `when selecting parent should select children`() = runTest {
+        // Arrange
+        val expectedList = DUMMY_COUNTRIES
+        val customListId = "id"
+        val customListName = "name"
+        val customList: RelayItem.CustomList = mockk {
+            every { id } returns customListId
+            every { name } returns customListName
+            every { locations } returns emptyList()
+        }
+        customListFlow.value = listOf(customList)
+        val expectedSelection =
+            (DUMMY_COUNTRIES + DUMMY_COUNTRIES.flatMap { it.descendants() }).toSet()
+        val viewModel = createViewModel(customListId, true)
+        relayListFlow.value = expectedList
+
+        // Act, Assert
+        viewModel.uiState.test {
+            // Check no selected
+            val firstState = awaitItem()
+            assertIs<CustomListLocationsUiState.Content.Data>(firstState)
+            assertEquals(emptySet<RelayItem>(), firstState.selectedLocations)
+            viewModel.onRelaySelectionClick(DUMMY_COUNTRIES[0], true)
+            // Check all items selected
+            val secondState = awaitItem()
+            assertIs<CustomListLocationsUiState.Content.Data>(secondState)
+            assertEquals(expectedSelection, secondState.selectedLocations)
+        }
+    }
+
+    @Test
+    fun `when deselecting child should deselect parent`() = runTest {
+        // Arrange
+        val expectedList = DUMMY_COUNTRIES
+        val initialSelection =
+            (DUMMY_COUNTRIES + DUMMY_COUNTRIES.flatMap { it.descendants() }).toSet()
+        val customListId = "id"
+        val customListName = "name"
+        val customList: RelayItem.CustomList = mockk {
+            every { id } returns customListId
+            every { name } returns customListName
+            every { locations } returns initialSelection.toList()
+        }
+        customListFlow.value = listOf(customList)
+        val expectedSelection = emptySet<RelayItem>()
+        val viewModel = createViewModel(customListId, true)
+        relayListFlow.value = expectedList
+
+        // Act, Assert
+        viewModel.uiState.test {
+            // Check initial selected
+            val firstState = awaitItem()
+            assertIs<CustomListLocationsUiState.Content.Data>(firstState)
+            assertEquals(initialSelection, 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)
+        }
+    }
+
+    @Test
+    fun `when deselecting parent should deselect child`() = runTest {
+        // Arrange
+        val expectedList = DUMMY_COUNTRIES
+        val initialSelection =
+            (DUMMY_COUNTRIES + DUMMY_COUNTRIES.flatMap { it.descendants() }).toSet()
+        val customListId = "id"
+        val customListName = "name"
+        val customList: RelayItem.CustomList = mockk {
+            every { id } returns customListId
+            every { name } returns customListName
+            every { locations } returns initialSelection.toList()
+        }
+        customListFlow.value = listOf(customList)
+        val expectedSelection = emptySet<RelayItem>()
+        val viewModel = createViewModel(customListId, true)
+        relayListFlow.value = expectedList
+
+        // Act, Assert
+        viewModel.uiState.test {
+            // Check initial selected
+            val firstState = awaitItem()
+            assertIs<CustomListLocationsUiState.Content.Data>(firstState)
+            assertEquals(initialSelection, firstState.selectedLocations)
+            viewModel.onRelaySelectionClick(DUMMY_COUNTRIES[0], false)
+            // Check all items selected
+            val secondState = awaitItem()
+            assertIs<CustomListLocationsUiState.Content.Data>(secondState)
+            assertEquals(expectedSelection, secondState.selectedLocations)
+        }
+    }
+
+    @Test
+    fun `when selecting child should not select parent`() = runTest {
+        // Arrange
+        val expectedList = DUMMY_COUNTRIES
+        val customListId = "id"
+        val customListName = "name"
+        val customList: RelayItem.CustomList = mockk {
+            every { id } returns customListId
+            every { name } returns customListName
+            every { locations } returns emptyList()
+        }
+        customListFlow.value = listOf(customList)
+        val expectedSelection = DUMMY_COUNTRIES[0].cities[0].relays.toSet()
+        val viewModel = createViewModel(customListId, true)
+        relayListFlow.value = expectedList
+
+        // Act, Assert
+        viewModel.uiState.test {
+            // Check no selected
+            val firstState = awaitItem()
+            assertIs<CustomListLocationsUiState.Content.Data>(firstState)
+            assertEquals(emptySet<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)
+        }
+    }
+
+    @Test
+    fun `given new list true when saving successfully should emit close screen side effect`() =
+        runTest {
+            // Arrange
+            val customListId = "1"
+            val customListName = "name"
+            val newList = true
+            val expectedResult: CustomListResult.LocationsChanged = mockk()
+            val customList: RelayItem.CustomList = mockk {
+                every { id } returns customListId
+                every { name } returns customListName
+                every { locations } returns DUMMY_COUNTRIES
+            }
+            customListFlow.value = listOf(customList)
+            coEvery {
+                mockCustomListUseCase.performAction(any<CustomListAction.UpdateLocations>())
+            } returns Result.success(expectedResult)
+            val viewModel = createViewModel(customListId, newList)
+
+            // Act, Assert
+            viewModel.uiSideEffect.test {
+                viewModel.save()
+                val sideEffect = awaitItem()
+                assertIs<CustomListLocationsSideEffect.CloseScreen>(sideEffect)
+            }
+        }
+
+    @Test
+    fun `given new list false when saving successfully should emit return with result side effect`() =
+        runTest {
+            // Arrange
+            val customListId = "1"
+            val customListName = "name"
+            val newList = false
+            val expectedResult: CustomListResult.LocationsChanged = mockk()
+            val customList: RelayItem.CustomList = mockk {
+                every { id } returns customListId
+                every { name } returns customListName
+                every { locations } returns DUMMY_COUNTRIES
+            }
+            customListFlow.value = listOf(customList)
+            coEvery {
+                mockCustomListUseCase.performAction(any<CustomListAction.UpdateLocations>())
+            } returns Result.success(expectedResult)
+            val viewModel = createViewModel(customListId, newList)
+
+            // Act, Assert
+            viewModel.uiSideEffect.test {
+                viewModel.save()
+                val sideEffect = awaitItem()
+                assertIs<CustomListLocationsSideEffect.ReturnWithResult>(sideEffect)
+                assertEquals(expectedResult, sideEffect.result)
+            }
+        }
+
+    private fun createViewModel(customListId: String, newList: Boolean) =
+        CustomListLocationsViewModel(
+            customListId = customListId,
+            newList = newList,
+            relayListUseCase = mockRelayListUseCase,
+            customListActionUseCase = mockCustomListUseCase
+        )
+
+    companion object {
+        private val DUMMY_COUNTRIES =
+            listOf(
+                RelayItem.Country(
+                    name = "Sweden",
+                    code = "SE",
+                    expanded = false,
+                    cities =
+                        listOf(
+                            RelayItem.City(
+                                name = "Gothenburg",
+                                code = "GBG",
+                                expanded = false,
+                                location = GeographicLocationConstraint.City("SE", "GBG"),
+                                relays =
+                                    listOf(
+                                        RelayItem.Relay(
+                                            name = "gbg-1",
+                                            locationName = "GBG gbg-1",
+                                            active = true,
+                                            location =
+                                                GeographicLocationConstraint.Hostname(
+                                                    "SE",
+                                                    "GBG",
+                                                    "gbg-1"
+                                                )
+                                        )
+                                    )
+                            )
+                        )
+                )
+            )
+    }
+}
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListsViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListsViewModelTest.kt
new file mode 100644
index 0000000000..612ae38a3a
--- /dev/null
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListsViewModelTest.kt
@@ -0,0 +1,54 @@
+package net.mullvad.mullvadvpn.viewmodel
+
+import app.cash.turbine.test
+import io.mockk.coVerify
+import io.mockk.every
+import io.mockk.mockk
+import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.test.runTest
+import net.mullvad.mullvadvpn.compose.communication.CustomListAction
+import net.mullvad.mullvadvpn.compose.state.CustomListsUiState
+import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
+import net.mullvad.mullvadvpn.relaylist.RelayItem
+import net.mullvad.mullvadvpn.usecase.RelayListUseCase
+import net.mullvad.mullvadvpn.usecase.customlists.CustomListActionUseCase
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.ExtendWith
+
+@ExtendWith(TestCoroutineRule::class)
+class CustomListsViewModelTest {
+    private val mockRelayListUseCase: RelayListUseCase = mockk(relaxed = true)
+    private val mockCustomListsActionUseCase: CustomListActionUseCase = mockk(relaxed = true)
+
+    @Test
+    fun `given custom list from relay list use case should be in state`() = runTest {
+        // Arrange
+        val customLists: List<RelayItem.CustomList> = mockk()
+        val expectedState = CustomListsUiState.Content(customLists)
+        every { mockRelayListUseCase.customLists() } returns flowOf(customLists)
+        val viewModel = createViewModel()
+
+        // Act, Assert
+        viewModel.uiState.test { assertEquals(expectedState, awaitItem()) }
+    }
+
+    @Test
+    fun `undo delete action should call custom list use case`() = runTest {
+        // Arrange
+        val viewModel = createViewModel()
+        val action: CustomListAction.Create = mockk()
+
+        // Act
+        viewModel.undoDeleteCustomList(action)
+
+        // Assert
+        coVerify { mockCustomListsActionUseCase.performAction(action) }
+    }
+
+    private fun createViewModel() =
+        CustomListsViewModel(
+            relayListUseCase = mockRelayListUseCase,
+            customListActionUseCase = mockCustomListsActionUseCase
+        )
+}
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
new file mode 100644
index 0000000000..9f7f3f1f0b
--- /dev/null
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/DeleteCustomListConfirmationViewModelTest.kt
@@ -0,0 +1,43 @@
+package net.mullvad.mullvadvpn.viewmodel
+
+import app.cash.turbine.test
+import io.mockk.coEvery
+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.CustomListResult
+import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
+import net.mullvad.mullvadvpn.usecase.customlists.CustomListActionUseCase
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.ExtendWith
+
+@ExtendWith(TestCoroutineRule::class)
+class DeleteCustomListConfirmationViewModelTest {
+    private val mockCustomListActionUseCase: CustomListActionUseCase = mockk()
+
+    @Test
+    fun `when successfully deleting a list should emit return with result side effect`() = runTest {
+        // Arrange
+        val expectedResult: CustomListResult.Deleted = mockk()
+        val viewModel = createViewModel()
+        coEvery {
+            mockCustomListActionUseCase.performAction(any<CustomListAction.Delete>())
+        } returns Result.success(expectedResult)
+
+        // Act, Assert
+        viewModel.uiSideEffect.test {
+            viewModel.deleteCustomList()
+            val sideEffect = awaitItem()
+            assertIs<DeleteCustomListConfirmationSideEffect.ReturnWithResult>(sideEffect)
+            assertEquals(expectedResult, sideEffect.result)
+        }
+    }
+
+    private fun createViewModel() =
+        DeleteCustomListConfirmationViewModel(
+            customListId = "1",
+            customListActionUseCase = mockCustomListActionUseCase
+        )
+}
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
new file mode 100644
index 0000000000..e9592d0336
--- /dev/null
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/EditCustomListNameDialogViewModelTest.kt
@@ -0,0 +1,90 @@
+package net.mullvad.mullvadvpn.viewmodel
+
+import app.cash.turbine.test
+import io.mockk.coEvery
+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.CustomListResult
+import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
+import net.mullvad.mullvadvpn.model.CustomListsError
+import net.mullvad.mullvadvpn.usecase.customlists.CustomListActionUseCase
+import net.mullvad.mullvadvpn.usecase.customlists.CustomListsException
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertNull
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.ExtendWith
+
+@ExtendWith(TestCoroutineRule::class)
+class EditCustomListNameDialogViewModelTest {
+    private val mockCustomListActionUseCase: CustomListActionUseCase = mockk()
+
+    @Test
+    fun `when successfully renamed list should emit return with result side effect`() = runTest {
+        // Arrange
+        val expectedResult: CustomListResult.Renamed = mockk()
+        val customListId = "id"
+        val customListName = "list"
+        val viewModel = createViewModel(customListId, customListName)
+        coEvery {
+            mockCustomListActionUseCase.performAction(any<CustomListAction.Rename>())
+        } returns Result.success(expectedResult)
+
+        // Act, Assert
+        viewModel.uiSideEffect.test {
+            viewModel.updateCustomListName(customListName)
+            val sideEffect = awaitItem()
+            assertIs<EditCustomListNameDialogSideEffect.ReturnWithResult>(sideEffect)
+            assertEquals(expectedResult, sideEffect.result)
+        }
+    }
+
+    @Test
+    fun `when failing to creating a list should update ui state with error`() = runTest {
+        // Arrange
+        val expectedError = CustomListsError.CustomListExists
+        val customListId = "id2"
+        val customListName = "list2"
+        val viewModel = createViewModel(customListId, customListName)
+        coEvery {
+            mockCustomListActionUseCase.performAction(any<CustomListAction.Rename>())
+        } returns Result.failure(CustomListsException(expectedError))
+
+        // Act, Assert
+        viewModel.uiState.test {
+            awaitItem() // Default state
+            viewModel.updateCustomListName(customListName)
+            assertEquals(expectedError, awaitItem().error)
+        }
+    }
+
+    @Test
+    fun `given error state when calling clear error then should update to state without error`() =
+        runTest {
+            // Arrange
+            val expectedError = CustomListsError.CustomListExists
+            val customListId = "id"
+            val customListName = "list"
+            val viewModel = createViewModel(customListId, customListName)
+            coEvery {
+                mockCustomListActionUseCase.performAction(any<CustomListAction.Rename>())
+            } returns Result.failure(CustomListsException(expectedError))
+
+            // Act, Assert
+            viewModel.uiState.test {
+                awaitItem() // Default state
+                viewModel.updateCustomListName(customListName)
+                assertEquals(expectedError, awaitItem().error) // Showing error
+                viewModel.clearError()
+                assertNull(awaitItem().error)
+            }
+        }
+
+    private fun createViewModel(customListId: String, initialName: String) =
+        EditCustomListNameDialogViewModel(
+            customListId = customListId,
+            initialName = initialName,
+            customListActionUseCase = mockCustomListActionUseCase
+        )
+}
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/EditCustomListViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/EditCustomListViewModelTest.kt
new file mode 100644
index 0000000000..33986961b3
--- /dev/null
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/EditCustomListViewModelTest.kt
@@ -0,0 +1,66 @@
+package net.mullvad.mullvadvpn.viewmodel
+
+import app.cash.turbine.test
+import io.mockk.every
+import io.mockk.mockk
+import kotlin.test.assertIs
+import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.test.runTest
+import net.mullvad.mullvadvpn.compose.state.EditCustomListState
+import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
+import net.mullvad.mullvadvpn.relaylist.RelayItem
+import net.mullvad.mullvadvpn.usecase.RelayListUseCase
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.ExtendWith
+
+@ExtendWith(TestCoroutineRule::class)
+class EditCustomListViewModelTest {
+    private val mockRelayListUseCase: RelayListUseCase = mockk(relaxed = true)
+
+    @Test
+    fun `given a custom list id that does not exists should return not found ui state`() = runTest {
+        // Arrange
+        val customListId = "2"
+        val customList =
+            RelayItem.CustomList(id = "1", name = "test", expanded = false, locations = emptyList())
+        every { mockRelayListUseCase.customLists() } returns flowOf(listOf(customList))
+        val viewModel = createViewModel(customListId)
+
+        // Act, Assert
+        viewModel.uiState.test {
+            val item = awaitItem()
+            assertIs<EditCustomListState.NotFound>(item)
+        }
+    }
+
+    @Test
+    fun `given a custom list id that exists should return content ui state`() = runTest {
+        // Arrange
+        val customListId = "1"
+        val customList =
+            RelayItem.CustomList(
+                id = customListId,
+                name = "test",
+                expanded = false,
+                locations = emptyList()
+            )
+        every { mockRelayListUseCase.customLists() } returns flowOf(listOf(customList))
+        val viewModel = createViewModel(customListId)
+
+        // Act, Assert
+        viewModel.uiState.test {
+            val item = awaitItem()
+            assertIs<EditCustomListState.Content>(item)
+            assertEquals(item.id, customList.id)
+            assertEquals(item.name, customList.name)
+            assertEquals(item.locations, customList.locations)
+        }
+    }
+
+    private fun createViewModel(customListId: String) =
+        EditCustomListViewModel(
+            customListId = customListId,
+            relayListUseCase = mockRelayListUseCase
+        )
+}
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 d8bcc7080f..41bff94ccd 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
@@ -2,6 +2,8 @@ package net.mullvad.mullvadvpn.viewmodel
 
 import androidx.lifecycle.viewModelScope
 import app.cash.turbine.test
+import io.mockk.coEvery
+import io.mockk.coVerify
 import io.mockk.every
 import io.mockk.just
 import io.mockk.mockk
@@ -14,7 +16,8 @@ import kotlin.test.assertIs
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.test.runTest
-import net.mullvad.mullvadvpn.compose.state.RelayListState
+import net.mullvad.mullvadvpn.compose.communication.CustomListAction
+import net.mullvad.mullvadvpn.compose.communication.CustomListResult
 import net.mullvad.mullvadvpn.compose.state.SelectLocationUiState
 import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
 import net.mullvad.mullvadvpn.lib.common.test.assertLists
@@ -33,6 +36,7 @@ import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManager
 import net.mullvad.mullvadvpn.ui.serviceconnection.connectionProxy
 import net.mullvad.mullvadvpn.usecase.RelayListFilterUseCase
 import net.mullvad.mullvadvpn.usecase.RelayListUseCase
+import net.mullvad.mullvadvpn.usecase.customlists.CustomListActionUseCase
 import org.junit.jupiter.api.AfterEach
 import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
@@ -47,6 +51,7 @@ class SelectLocationViewModelTest {
     private val relayListWithSelectionFlow =
         MutableStateFlow(RelayList(emptyList(), emptyList(), null))
     private val mockRelayListUseCase: RelayListUseCase = mockk()
+    private val mockCustomListActionUseCase: CustomListActionUseCase = mockk(relaxed = true)
     private val selectedOwnership = MutableStateFlow<Constraint<Ownership>>(Constraint.Any())
     private val selectedProvider = MutableStateFlow<Constraint<Providers>>(Constraint.Any())
     private val allProvider = MutableStateFlow<List<Provider>>(emptyList())
@@ -63,11 +68,13 @@ class SelectLocationViewModelTest {
         mockkStatic(SERVICE_CONNECTION_MANAGER_EXTENSIONS)
         mockkStatic(RELAY_LIST_EXTENSIONS)
         mockkStatic(RELAY_ITEM_EXTENSIONS)
+        mockkStatic(CUSTOM_LIST_EXTENSIONS)
         viewModel =
             SelectLocationViewModel(
                 mockServiceConnectionManager,
                 mockRelayListUseCase,
-                mockRelayListFilterUseCase
+                mockRelayListFilterUseCase,
+                mockCustomListActionUseCase
             )
     }
 
@@ -94,16 +101,9 @@ class SelectLocationViewModelTest {
         // Act, Assert
         viewModel.uiState.test {
             val actualState = awaitItem()
-            assertIs<SelectLocationUiState.Data>(actualState)
-            assertIs<RelayListState.RelayList>(actualState.relayListState)
-            assertLists(
-                mockCountries,
-                (actualState.relayListState as RelayListState.RelayList).countries
-            )
-            assertEquals(
-                selectedItem,
-                (actualState.relayListState as RelayListState.RelayList).selectedItem
-            )
+            assertIs<SelectLocationUiState.Content>(actualState)
+            assertLists(mockCountries, actualState.countries)
+            assertEquals(selectedItem, actualState.selectedItem)
         }
     }
 
@@ -121,16 +121,9 @@ class SelectLocationViewModelTest {
             // Act, Assert
             viewModel.uiState.test {
                 val actualState = awaitItem()
-                assertIs<SelectLocationUiState.Data>(actualState)
-                assertIs<RelayListState.RelayList>(actualState.relayListState)
-                assertLists(
-                    mockCountries,
-                    (actualState.relayListState as RelayListState.RelayList).countries
-                )
-                assertEquals(
-                    selectedItem,
-                    (actualState.relayListState as RelayListState.RelayList).selectedItem
-                )
+                assertIs<SelectLocationUiState.Content>(actualState)
+                assertLists(mockCountries, actualState.countries)
+                assertEquals(selectedItem, actualState.selectedItem)
             }
         }
 
@@ -169,28 +162,22 @@ class SelectLocationViewModelTest {
         val mockSearchString = "SEARCH"
         every { mockRelayList.filterOnSearchTerm(mockSearchString, selectedItem) } returns
             mockCountries
+        every { mockCustomList.filterOnSearchTerm(mockSearchString) } returns mockCustomList
         relayListWithSelectionFlow.value = RelayList(mockCustomList, mockRelayList, selectedItem)
 
         // Act, Assert
         viewModel.uiState.test {
             // Wait for first data
-            assertIs<SelectLocationUiState.Data>(awaitItem())
+            assertIs<SelectLocationUiState.Content>(awaitItem())
 
             // Update search string
             viewModel.onSearchTermInput(mockSearchString)
 
             // Assert
             val actualState = awaitItem()
-            assertIs<SelectLocationUiState.Data>(actualState)
-            assertIs<RelayListState.RelayList>(actualState.relayListState)
-            assertLists(
-                mockCountries,
-                (actualState.relayListState as RelayListState.RelayList).countries
-            )
-            assertEquals(
-                selectedItem,
-                (actualState.relayListState as RelayListState.RelayList).selectedItem
-            )
+            assertIs<SelectLocationUiState.Content>(actualState)
+            assertLists(mockCountries, actualState.countries)
+            assertEquals(selectedItem, actualState.selectedItem)
         }
     }
 
@@ -204,19 +191,20 @@ class SelectLocationViewModelTest {
         val mockSearchString = "SEARCH"
         every { mockRelayList.filterOnSearchTerm(mockSearchString, selectedItem) } returns
             mockCountries
+        every { mockCustomList.filterOnSearchTerm(mockSearchString) } returns mockCustomList
         relayListWithSelectionFlow.value = RelayList(mockCustomList, mockRelayList, selectedItem)
 
         // Act, Assert
         viewModel.uiState.test {
             // Wait for first data
-            assertIs<SelectLocationUiState.Data>(awaitItem())
+            assertIs<SelectLocationUiState.Content>(awaitItem())
 
             // Update search string
             viewModel.onSearchTermInput(mockSearchString)
 
             // Assert
             val actualState = awaitItem()
-            assertIs<SelectLocationUiState.Data>(actualState)
+            assertIs<SelectLocationUiState.Content>(actualState)
             assertEquals(mockSearchString, actualState.searchTerm)
         }
     }
@@ -257,6 +245,40 @@ class SelectLocationViewModelTest {
         }
     }
 
+    @Test
+    fun `when perform action is called should call custom list use case`() {
+        // Arrange
+        val action: CustomListAction = mockk()
+
+        // Act
+        viewModel.performAction(action)
+
+        // Assert
+        coVerify { mockCustomListActionUseCase.performAction(action) }
+    }
+
+    @Test
+    fun `after adding a location to a list should emit location added side effect`() = runTest {
+        // Arrange
+        val expectedResult: CustomListResult.LocationsChanged = mockk()
+        val location: RelayItem = mockk { every { code } returns "code" }
+        val customList: RelayItem.CustomList = mockk {
+            every { id } returns "1"
+            every { locations } returns emptyList()
+        }
+        coEvery {
+            mockCustomListActionUseCase.performAction(any<CustomListAction.UpdateLocations>())
+        } returns Result.success(expectedResult)
+
+        // Act, Assert
+        viewModel.uiSideEffect.test {
+            viewModel.addLocationToList(item = location, customList = customList)
+            val sideEffect = awaitItem()
+            assertIs<SelectLocationSideEffect.LocationAddedToCustomList>(sideEffect)
+            assertEquals(expectedResult, sideEffect.result)
+        }
+    }
+
     companion object {
         private const val SERVICE_CONNECTION_MANAGER_EXTENSIONS =
             "net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManagerExtensionsKt"
@@ -264,5 +286,7 @@ class SelectLocationViewModelTest {
             "net.mullvad.mullvadvpn.relaylist.RelayListExtensionsKt"
         private const val RELAY_ITEM_EXTENSIONS =
             "net.mullvad.mullvadvpn.relaylist.RelayItemExtensionsKt"
+        private const val CUSTOM_LIST_EXTENSIONS =
+            "net.mullvad.mullvadvpn.relaylist.CustomListExtensionsKt"
     }
 }