Make UI state have constraints

15 files changed, 147 insertions, 121 deletions

diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/FilterScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/FilterScreenTest.kt
index 4b80ea0e3c..7e7c61c9de 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/FilterScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/FilterScreenTest.kt
@@ -10,6 +10,7 @@ import io.mockk.verify
 import net.mullvad.mullvadvpn.compose.createEdgeToEdgeComposeExtension
 import net.mullvad.mullvadvpn.compose.setContentWithTheme
 import net.mullvad.mullvadvpn.compose.state.RelayFilterUiState
+import net.mullvad.mullvadvpn.lib.model.Constraint
 import net.mullvad.mullvadvpn.lib.model.Ownership
 import net.mullvad.mullvadvpn.lib.model.ProviderId
 import org.junit.jupiter.api.Test
@@ -27,7 +28,7 @@ class FilterScreenTest {
         state: RelayFilterUiState = RelayFilterUiState(),
         onBackClick: () -> Unit = {},
         onApplyClick: () -> Unit = {},
-        onSelectedOwnership: (ownership: Ownership?) -> Unit = {},
+        onSelectedOwnership: (ownership: Constraint<Ownership>) -> Unit = {},
         onAllProviderCheckChange: (isChecked: Boolean) -> Unit = {},
         onSelectedProvider: (checked: Boolean, provider: ProviderId) -> Unit = { _, _ -> },
     ) {
@@ -50,8 +51,8 @@ class FilterScreenTest {
                 state =
                     RelayFilterUiState(
                         providerToOwnerships = DUMMY_RELAY_ALL_PROVIDERS,
-                        selectedOwnership = null,
-                        selectedProviders = DUMMY_SELECTED_PROVIDERS,
+                        selectedOwnership = Constraint.Any,
+                        selectedProviders = Constraint.Only(DUMMY_SELECTED_PROVIDERS),
                     )
             )
             onNodeWithText("Ownership").assertExists()
@@ -65,8 +66,8 @@ class FilterScreenTest {
                 state =
                     RelayFilterUiState(
                         providerToOwnerships = DUMMY_RELAY_ALL_PROVIDERS,
-                        selectedOwnership = null,
-                        selectedProviders = DUMMY_SELECTED_PROVIDERS,
+                        selectedOwnership = Constraint.Any,
+                        selectedProviders = Constraint.Only(DUMMY_SELECTED_PROVIDERS),
                     )
             )
             onNodeWithText("Ownership").performClick()
@@ -80,8 +81,8 @@ class FilterScreenTest {
                 state =
                     RelayFilterUiState(
                         providerToOwnerships = DUMMY_RELAY_ALL_PROVIDERS,
-                        selectedOwnership = Ownership.MullvadOwned,
-                        selectedProviders = DUMMY_SELECTED_PROVIDERS,
+                        selectedOwnership = Constraint.Only(Ownership.MullvadOwned),
+                        selectedProviders = Constraint.Only(DUMMY_SELECTED_PROVIDERS),
                     )
             )
             onNodeWithText("Ownership").performClick()
@@ -95,8 +96,8 @@ class FilterScreenTest {
                 state =
                     RelayFilterUiState(
                         providerToOwnerships = DUMMY_RELAY_ALL_PROVIDERS,
-                        selectedOwnership = Ownership.Rented,
-                        selectedProviders = DUMMY_SELECTED_PROVIDERS,
+                        selectedOwnership = Constraint.Only(Ownership.Rented),
+                        selectedProviders = Constraint.Only(DUMMY_SELECTED_PROVIDERS),
                     )
             )
             onNodeWithText("Ownership").performClick()
@@ -110,8 +111,8 @@ class FilterScreenTest {
                 state =
                     RelayFilterUiState(
                         providerToOwnerships = DUMMY_RELAY_ALL_PROVIDERS,
-                        selectedOwnership = null,
-                        selectedProviders = DUMMY_SELECTED_PROVIDERS,
+                        selectedOwnership = Constraint.Any,
+                        selectedProviders = Constraint.Only(DUMMY_SELECTED_PROVIDERS),
                     )
             )
 
@@ -128,8 +129,8 @@ class FilterScreenTest {
                 state =
                     RelayFilterUiState(
                         providerToOwnerships = DUMMY_RELAY_ALL_PROVIDERS,
-                        selectedOwnership = null,
-                        selectedProviders = listOf(ProviderId("31173")),
+                        selectedOwnership = Constraint.Any,
+                        selectedProviders = Constraint.Only(listOf(ProviderId("31173"))),
                     ),
                 onApplyClick = mockClickListener,
             )
@@ -178,6 +179,6 @@ class FilterScreenTest {
                 ProviderId("xtom") to setOf(Ownership.Rented),
             )
 
-        private val DUMMY_SELECTED_PROVIDERS = DUMMY_RELAY_ALL_PROVIDERS.keys.toList()
+        private val DUMMY_SELECTED_PROVIDERS = DUMMY_RELAY_ALL_PROVIDERS.keys.toList().dropLast(3)
     }
 }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/FilterUiStatePreviewParameterProvider.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/FilterUiStatePreviewParameterProvider.kt
index 0aa3be8158..2a5188e8ee 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/FilterUiStatePreviewParameterProvider.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/preview/FilterUiStatePreviewParameterProvider.kt
@@ -2,6 +2,7 @@ package net.mullvad.mullvadvpn.compose.preview
 
 import androidx.compose.ui.tooling.preview.PreviewParameterProvider
 import net.mullvad.mullvadvpn.compose.state.RelayFilterUiState
+import net.mullvad.mullvadvpn.lib.model.Constraint
 import net.mullvad.mullvadvpn.lib.model.Ownership
 import net.mullvad.mullvadvpn.lib.model.ProviderId
 
@@ -12,8 +13,8 @@ class FilterUiStatePreviewParameterProvider : PreviewParameterProvider<RelayFilt
         sequenceOf(
             RelayFilterUiState(
                 providerToOwnerships = PROVIDER_TO_OWNERSHIPS,
-                selectedOwnership = Ownership.MullvadOwned,
-                selectedProviders = PROVIDER_TO_OWNERSHIPS.keys.toList(),
+                selectedOwnership = Constraint.Only(Ownership.MullvadOwned),
+                selectedProviders = Constraint.Only(PROVIDER_TO_OWNERSHIPS.keys),
             )
         )
 }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/FilterScreen.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/FilterScreen.kt
index 3d3f937c73..00d1e3bb9a 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/FilterScreen.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/FilterScreen.kt
@@ -45,8 +45,10 @@ import net.mullvad.mullvadvpn.compose.preview.FilterUiStatePreviewParameterProvi
 import net.mullvad.mullvadvpn.compose.state.RelayFilterUiState
 import net.mullvad.mullvadvpn.compose.transitions.SlideInFromRightTransition
 import net.mullvad.mullvadvpn.compose.util.CollectSideEffectWithLifecycle
+import net.mullvad.mullvadvpn.lib.model.Constraint
 import net.mullvad.mullvadvpn.lib.model.Ownership
 import net.mullvad.mullvadvpn.lib.model.ProviderId
+import net.mullvad.mullvadvpn.lib.model.Providers
 import net.mullvad.mullvadvpn.lib.theme.AppTheme
 import net.mullvad.mullvadvpn.lib.theme.Dimens
 import net.mullvad.mullvadvpn.viewmodel.FilterScreenSideEffect
@@ -96,7 +98,7 @@ fun FilterScreen(
     state: RelayFilterUiState,
     onBackClick: () -> Unit,
     onApplyClick: () -> Unit,
-    onSelectedOwnership: (ownership: Ownership?) -> Unit,
+    onSelectedOwnership: (ownership: Constraint<Ownership>) -> Unit,
     onAllProviderCheckChange: (isChecked: Boolean) -> Unit,
     onSelectedProvider: (checked: Boolean, provider: ProviderId) -> Unit,
 ) {
@@ -121,14 +123,14 @@ fun FilterScreen(
             }
             if (ownershipExpanded) {
                 item(key = Keys.OWNERSHIP_ALL, contentType = ContentType.ITEM) {
-                    AnyOwnership(state, onSelectedOwnership)
+                    AnyOwnership(state, { onSelectedOwnership(Constraint.Any) })
                 }
                 itemsWithDivider(
                     key = { it.name },
                     contentType = { ContentType.ITEM },
                     items = state.selectableOwnerships,
                 ) { ownership ->
-                    Ownership(ownership, state, onSelectedOwnership)
+                    Ownership(ownership, state, { onSelectedOwnership(Constraint.Only(it)) })
                 }
             }
             itemWithDivider(key = Keys.PROVIDERS_TITLE, contentType = ContentType.HEADER) {
@@ -171,14 +173,11 @@ private fun LazyItemScope.OwnershipHeader(expanded: Boolean, onToggleExpanded: (
 }
 
 @Composable
-private fun LazyItemScope.AnyOwnership(
-    state: RelayFilterUiState,
-    onSelectedOwnership: (ownership: Ownership?) -> Unit,
-) {
+private fun LazyItemScope.AnyOwnership(state: RelayFilterUiState, onSelectedOwnership: () -> Unit) {
     SelectableCell(
         title = stringResource(id = R.string.any),
-        isSelected = state.selectedOwnership == null,
-        onCellClicked = { onSelectedOwnership(null) },
+        isSelected = state.selectedOwnership is Constraint.Any,
+        onCellClicked = { onSelectedOwnership() },
         modifier = Modifier.animateItem(),
     )
 }
@@ -187,11 +186,11 @@ private fun LazyItemScope.AnyOwnership(
 private fun LazyItemScope.Ownership(
     ownership: Ownership,
     state: RelayFilterUiState,
-    onSelectedOwnership: (ownership: Ownership?) -> Unit,
+    onSelectedOwnership: (ownership: Ownership) -> Unit,
 ) {
     SelectableCell(
         title = stringResource(id = ownership.stringResource()),
-        isSelected = ownership == state.selectedOwnership,
+        isSelected = ownership == state.selectedOwnership.getOrNull(),
         onCellClicked = { onSelectedOwnership(ownership) },
         modifier = Modifier.animateItem(),
     )
@@ -230,12 +229,18 @@ private fun LazyItemScope.Provider(
 ) {
     CheckboxCell(
         title = providerId.value,
-        checked = providerId in state.selectedProviders,
+        checked = providerId.isChecked(state.selectedProviders),
         onCheckedChange = { checked -> onSelectedProvider(checked, providerId) },
         modifier = Modifier.animateItem(),
     )
 }
 
+private fun ProviderId.isChecked(constraint: Constraint<Providers>) =
+    when (constraint) {
+        Constraint.Any -> true
+        is Constraint.Only -> this in constraint.value
+    }
+
 @Composable
 private fun LazyItemScope.RemovedProvider(
     providerId: ProviderId,
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/FilterConstrainExtensions.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/FilterConstrainExtensions.kt
index 2ddf35fad5..b7531f15ec 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/FilterConstrainExtensions.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/FilterConstrainExtensions.kt
@@ -2,7 +2,6 @@ package net.mullvad.mullvadvpn.compose.state
 
 import net.mullvad.mullvadvpn.lib.model.Constraint
 import net.mullvad.mullvadvpn.lib.model.Ownership
-import net.mullvad.mullvadvpn.lib.model.ProviderId
 import net.mullvad.mullvadvpn.lib.model.Providers
 
 fun Ownership?.toOwnershipConstraint(): Constraint<Ownership> =
@@ -11,15 +10,9 @@ fun Ownership?.toOwnershipConstraint(): Constraint<Ownership> =
         else -> Constraint.Only(this)
     }
 
-fun Constraint<Providers>.toSelectedProviders(allProviders: List<ProviderId>): List<ProviderId> =
-    when (this) {
-        Constraint.Any -> allProviders
-        is Constraint.Only -> value.providers.toList()
-    }
-
-fun List<ProviderId>.toConstraintProviders(allProviders: List<ProviderId>): Constraint<Providers> =
+fun Providers.toConstraintProviders(allProviders: Providers): Constraint<Providers> =
     if (size == allProviders.size) {
         Constraint.Any
     } else {
-        Constraint.Only(Providers(toHashSet()))
+        Constraint.Only(this)
     }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/RelayFilterUiState.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/RelayFilterUiState.kt
index 0c81ac0aa5..d3e8a2b685 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/RelayFilterUiState.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/RelayFilterUiState.kt
@@ -1,37 +1,41 @@
 package net.mullvad.mullvadvpn.compose.state
 
+import net.mullvad.mullvadvpn.lib.model.Constraint
 import net.mullvad.mullvadvpn.lib.model.Ownership
 import net.mullvad.mullvadvpn.lib.model.ProviderId
+import net.mullvad.mullvadvpn.lib.model.Providers
 
 data class RelayFilterUiState(
     private val providerToOwnerships: Map<ProviderId, Set<Ownership>> = emptyMap(),
-    val selectedOwnership: Ownership? = null,
-    val selectedProviders: List<ProviderId> = emptyList(),
+    val selectedOwnership: Constraint<Ownership> = Constraint.Any,
+    val selectedProviders: Constraint<Providers> = Constraint.Any,
 ) {
-    val allProviders: List<ProviderId> = providerToOwnerships.keys.toList().sorted()
+    val allProviders: Providers = providerToOwnerships.keys
 
     val selectableOwnerships: List<Ownership> =
-        if (selectedProviders.isEmpty()) {
-                Ownership.entries
-            } else {
-                providerToOwnerships
-                    .filterKeys { it in selectedProviders }
-                    .values
-                    .flatten()
-                    .distinct()
-            }
-            .sorted()
+        when (selectedProviders) {
+            Constraint.Any -> Ownership.entries
+            is Constraint.Only ->
+                if (selectedProviders.value.isEmpty()) {
+                    Ownership.entries
+                } else {
+                    providerToOwnerships
+                        .filterKeys { it in selectedProviders.value }
+                        .values
+                        .flatten()
+                        .distinct()
+                }
+        }.sorted()
 
     val selectableProviders: List<ProviderId> =
-        if (selectedOwnership != null) {
-            providerToOwnerships.filterValues { selectedOwnership in it }.keys.toList().sorted()
-        } else {
-            allProviders
-        }
+        when (selectedOwnership) {
+            Constraint.Any -> allProviders.toList()
+            is Constraint.Only ->
+                providerToOwnerships.filterValues { selectedOwnership.value in it }.keys
+        }.sorted()
 
+    val isApplyButtonEnabled = selectedProviders.getOrNull()?.isNotEmpty() != false
     val removedProviders: List<ProviderId> = selectedProviders - allProviders
 
-    val isApplyButtonEnabled = selectedProviders.isNotEmpty()
-
-    val isAllProvidersChecked = allProviders.size == selectedProviders.size
+    val isAllProvidersChecked = selectedProviders is Constraint.Any
 }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItemExtensions.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItemExtensions.kt
index 21d3294de4..d58da5bc9a 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItemExtensions.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItemExtensions.kt
@@ -46,7 +46,7 @@ private fun RelayItem.Location.hasProvider(providersConstraint: Constraint<Provi
         when (this) {
             is RelayItem.Location.Country -> cities.any { it.hasProvider(providersConstraint) }
             is RelayItem.Location.City -> relays.any { it.hasProvider(providersConstraint) }
-            is RelayItem.Location.Relay -> provider in providersConstraint.value.providers
+            is RelayItem.Location.Relay -> provider in providersConstraint.value
         }
     } else {
         true
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/FilterChipUseCase.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/FilterChipUseCase.kt
index 361e29944e..8c542202ff 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/FilterChipUseCase.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/FilterChipUseCase.kt
@@ -57,7 +57,7 @@ class FilterChipUseCase(
             when (selectedConstraintProviders) {
                 is Constraint.Any -> null
                 is Constraint.Only ->
-                    selectedConstraintProviders.value.providers
+                    selectedConstraintProviders.value
                         .filter { providerId ->
                             if (ownershipFilter == null) {
                                 true
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/FilterViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/FilterViewModel.kt
index c05515126e..41d98b40fb 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/FilterViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/FilterViewModel.kt
@@ -2,6 +2,7 @@ package net.mullvad.mullvadvpn.viewmodel
 
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
+import kotlin.collections.plus
 import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.SharingStarted
@@ -10,13 +11,13 @@ import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.receiveAsFlow
 import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
 import net.mullvad.mullvadvpn.compose.state.RelayFilterUiState
-import net.mullvad.mullvadvpn.compose.state.toConstraintProviders
-import net.mullvad.mullvadvpn.compose.state.toOwnershipConstraint
-import net.mullvad.mullvadvpn.compose.state.toSelectedProviders
+import net.mullvad.mullvadvpn.lib.model.Constraint
 import net.mullvad.mullvadvpn.lib.model.Ownership
 import net.mullvad.mullvadvpn.lib.model.ProviderId
+import net.mullvad.mullvadvpn.lib.model.Providers
 import net.mullvad.mullvadvpn.repository.RelayListFilterRepository
 import net.mullvad.mullvadvpn.usecase.ProviderToOwnershipsUseCase
 
@@ -27,24 +28,13 @@ class FilterViewModel(
     private val _uiSideEffect = Channel<FilterScreenSideEffect>()
     val uiSideEffect = _uiSideEffect.receiveAsFlow()
 
-    private val selectedOwnership = MutableStateFlow<Ownership?>(null)
-    private val selectedProviders = MutableStateFlow<List<ProviderId>>(emptyList())
+    private val selectedOwnership = MutableStateFlow<Constraint<Ownership>>(Constraint.Any)
+    private val selectedProviders = MutableStateFlow<Constraint<Providers>>(Constraint.Any)
 
     init {
         viewModelScope.launch {
-            selectedProviders.value =
-                combine(
-                        providerToOwnershipsUseCase(),
-                        relayListFilterRepository.selectedProviders,
-                    ) { providerToOwnerships, selectedConstraintProviders ->
-                        selectedConstraintProviders.toSelectedProviders(
-                            providerToOwnerships.keys.toList()
-                        )
-                    }
-                    .first()
-
-            val ownershipConstraint = relayListFilterRepository.selectedOwnership.first()
-            selectedOwnership.value = ownershipConstraint.getOrNull()
+            selectedProviders.value = relayListFilterRepository.selectedProviders.first()
+            selectedOwnership.value = relayListFilterRepository.selectedOwnership.first()
         }
     }
 
@@ -54,8 +44,8 @@ class FilterViewModel(
 
     private fun createState(
         providerToOwnerships: Map<ProviderId, Set<Ownership>>,
-        selectedOwnership: Ownership?,
-        selectedProviders: List<ProviderId>,
+        selectedOwnership: Constraint<Ownership>,
+        selectedProviders: Constraint<Providers>,
     ): RelayFilterUiState =
         RelayFilterUiState(
             providerToOwnerships = providerToOwnerships,
@@ -63,34 +53,61 @@ class FilterViewModel(
             selectedProviders = selectedProviders,
         )
 
-    fun setSelectedOwnership(ownership: Ownership?) {
+    fun setSelectedOwnership(ownership: Constraint<Ownership>) {
         selectedOwnership.value = ownership
     }
 
     fun setSelectedProvider(checked: Boolean, provider: ProviderId) {
-        selectedProviders.value =
+        selectedProviders.update {
             if (checked) {
-                selectedProviders.value + provider
+                it.check(provider, uiState.value.allProviders)
             } else {
-                selectedProviders.value - provider
+                it.uncheck(provider, uiState.value.allProviders)
+            }
+        }
+    }
+
+    private fun Constraint<Providers>.check(
+        provider: ProviderId,
+        allProviders: Providers,
+    ): Constraint<Providers> {
+        return when (this) {
+            is Constraint.Any -> Constraint.Any
+            is Constraint.Only -> {
+                val newProviderList = value + provider
+                if (allProviders.size == newProviderList.size) {
+                    Constraint.Any
+                } else {
+                    Constraint.Only(newProviderList)
+                }
             }
+        }
+    }
+
+    private fun Constraint<Providers>.uncheck(
+        provider: ProviderId,
+        allProviders: Providers,
+    ): Constraint<Providers> {
+        return when (this) {
+            is Constraint.Any -> Constraint.Only(allProviders - provider)
+            is Constraint.Only -> Constraint.Only(value - provider)
+        }
     }
 
     fun setAllProviders(isChecked: Boolean) {
         viewModelScope.launch {
             selectedProviders.value =
                 if (isChecked) {
-                    providerToOwnershipsUseCase().first().keys.toList()
+                    Constraint.Any
                 } else {
-                    emptyList()
+                    Constraint.Only(emptySet())
                 }
         }
     }
 
     fun onApplyButtonClicked() {
-        val newSelectedOwnership = selectedOwnership.value.toOwnershipConstraint()
-        val newSelectedProviders =
-            selectedProviders.value.toConstraintProviders(uiState.value.allProviders)
+        val newSelectedOwnership = selectedOwnership.value
+        val newSelectedProviders = selectedProviders.value
 
         viewModelScope.launch {
             relayListFilterRepository.updateSelectedOwnershipAndProviderFilter(
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/repository/RelayListFilterRepositoryTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/repository/RelayListFilterRepositoryTest.kt
index 6ec1ae1fea..4575d8c895 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/repository/RelayListFilterRepositoryTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/repository/RelayListFilterRepositoryTest.kt
@@ -59,8 +59,7 @@ class RelayListFilterRepositoryTest {
     fun `when settings is updated selected providers should update`() = runTest {
         // Arrange
         val mockSettings: Settings = mockk()
-        val selectedProviders: Constraint<Providers> =
-            Constraint.Only(Providers(setOf(ProviderId("Prove"))))
+        val selectedProviders: Constraint<Providers> = Constraint.Only(setOf(ProviderId("Prove")))
         every { mockSettings.relaySettings.relayConstraints.providers } returns selectedProviders
 
         // Act, Assert
@@ -147,7 +146,7 @@ class RelayListFilterRepositoryTest {
     @Test
     fun `when successfully updating selected providers should return successful`() = runTest {
         // Arrange
-        val providers = Constraint.Only(Providers(setOf(ProviderId("Mopp"))))
+        val providers = Constraint.Only(setOf(ProviderId("Mopp")))
         coEvery { mockManagementService.setProviders(providers) } returns Unit.right()
 
         // Act
@@ -162,7 +161,7 @@ class RelayListFilterRepositoryTest {
     fun `when failing to update selected providers should return SetWireguardConstraintsError`() =
         runTest {
             // Arrange
-            val providers = Constraint.Only(Providers(setOf(ProviderId("Mopp"))))
+            val providers = Constraint.Only(setOf(ProviderId("Mopp")))
             val error = SetWireguardConstraintsError.Unknown(mockk())
             coEvery { mockManagementService.setProviders(providers) } returns error.left()
 
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/FilterChipUseCaseTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/FilterChipUseCaseTest.kt
index bd1a759b77..32921b4ab0 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/FilterChipUseCaseTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/FilterChipUseCaseTest.kt
@@ -71,7 +71,7 @@ class FilterChipUseCaseTest {
     @Test
     fun `when provider filter is applied should return correct number of providers`() = runTest {
         // Arrange
-        val expectedProviders = Providers(providers = setOf(ProviderId("1"), ProviderId("2")))
+        val expectedProviders = setOf(ProviderId("1"), ProviderId("2"))
         selectedProviders.value = Constraint.Only(expectedProviders)
         providerToOwnerships.value =
             mapOf(
@@ -88,7 +88,7 @@ class FilterChipUseCaseTest {
     fun `when provider and ownership filter is applied should return correct filter chips`() =
         runTest {
             // Arrange
-            val expectedProviders = Providers(providers = setOf(ProviderId("1")))
+            val expectedProviders = setOf(ProviderId("1"))
             val expectedOwnership = Ownership.MullvadOwned
             selectedProviders.value = Constraint.Only(expectedProviders)
             selectedOwnership.value = Constraint.Only(expectedOwnership)
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/FilterViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/FilterViewModelTest.kt
index 49b9b4baff..7c905c0119 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/FilterViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/FilterViewModelTest.kt
@@ -16,7 +16,7 @@ import kotlinx.coroutines.test.runTest
 import net.mullvad.mullvadvpn.compose.state.toConstraintProviders
 import net.mullvad.mullvadvpn.compose.state.toOwnershipConstraint
 import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
-import net.mullvad.mullvadvpn.lib.common.test.assertLists
+import net.mullvad.mullvadvpn.lib.common.test.assertSet
 import net.mullvad.mullvadvpn.lib.model.Constraint
 import net.mullvad.mullvadvpn.lib.model.Ownership
 import net.mullvad.mullvadvpn.lib.model.ProviderId
@@ -55,15 +55,15 @@ class FilterViewModelTest {
             ProviderId("Tzulo") to setOf(Ownership.Rented),
             ProviderId("xtom") to setOf(Ownership.Rented),
         )
-    private val mockSelectedProviders: List<ProviderId> =
-        listOf(ProviderId("31173"), ProviderId("Blix"), ProviderId("Creanova"))
+    private val mockSelectedProviders: Providers =
+        setOf(ProviderId("31173"), ProviderId("Blix"), ProviderId("Creanova"))
 
     @BeforeEach
     fun setup() {
         every { mockRelayListFilterRepository.selectedOwnership } returns selectedOwnership
         every { mockProvidersOwnershipUseCase() } returns flowOf(dummyListOfAllProviders)
         every { mockRelayListFilterRepository.selectedProviders } returns
-            MutableStateFlow(Constraint.Only(Providers(mockSelectedProviders.toHashSet())))
+            MutableStateFlow(Constraint.Only(mockSelectedProviders))
         viewModel =
             FilterViewModel(
                 providerToOwnershipsUseCase = mockProvidersOwnershipUseCase,
@@ -84,9 +84,9 @@ class FilterViewModelTest {
             val mockOwnership = Ownership.Rented
             // Assert
             viewModel.uiState.test {
-                assertEquals(awaitItem().selectedOwnership, Ownership.MullvadOwned)
-                viewModel.setSelectedOwnership(mockOwnership)
-                assertEquals(mockOwnership, awaitItem().selectedOwnership)
+                assertEquals(Constraint.Only(Ownership.MullvadOwned), awaitItem().selectedOwnership)
+                viewModel.setSelectedOwnership(Constraint.Only(mockOwnership))
+                assertEquals(Constraint.Only(mockOwnership), awaitItem().selectedOwnership)
             }
         }
 
@@ -97,11 +97,11 @@ class FilterViewModelTest {
             val mockSelectedProvidersList = ProviderId("ptisp")
             // Assert
             viewModel.uiState.test {
-                assertLists(awaitItem().selectedProviders, mockSelectedProviders)
+                assertSet(mockSelectedProviders, awaitItem().selectedProviders.getOrNull()!!)
                 viewModel.setSelectedProvider(true, mockSelectedProvidersList)
-                assertLists(
-                    listOf(mockSelectedProvidersList) + mockSelectedProviders,
-                    awaitItem().selectedProviders,
+                assertSet(
+                    setOf(mockSelectedProvidersList) + mockSelectedProviders,
+                    awaitItem().selectedProviders.getOrNull()!!,
                 )
             }
         }
@@ -109,14 +109,12 @@ class FilterViewModelTest {
     @Test
     fun `setAllProvider with true should emit uiState with selectedProviders includes all providers`() =
         runTest {
-            // Arrange
-            val mockProvidersList = dummyListOfAllProviders.keys.toList()
             // Act
             viewModel.setAllProviders(true)
             // Assert
             viewModel.uiState.test {
                 val state = awaitItem()
-                assertEquals(mockProvidersList, state.selectedProviders)
+                assertEquals(Constraint.Any, state.selectedProviders)
             }
         }
 
@@ -126,7 +124,7 @@ class FilterViewModelTest {
             // Arrange
             val mockOwnership = Ownership.MullvadOwned.toOwnershipConstraint()
             val mockSelectedProviders =
-                mockSelectedProviders.toConstraintProviders(dummyListOfAllProviders.keys.toList())
+                mockSelectedProviders.toConstraintProviders(dummyListOfAllProviders.keys)
             coEvery {
                 mockRelayListFilterRepository.updateSelectedOwnershipAndProviderFilter(
                     mockOwnership,
@@ -149,12 +147,12 @@ class FilterViewModelTest {
     @Test
     fun `ensure that providers with multiple ownership are only returned once`() = runTest {
         // Arrange
-        val expectedProviderList = dummyListOfAllProviders.keys.toList()
+        val expectedProviderList = dummyListOfAllProviders.keys
 
         // Assert
         viewModel.uiState.test {
             val state = awaitItem()
-            assertLists(expectedProviderList, state.allProviders)
+            assertSet(expectedProviderList, state.allProviders)
         }
     }
 
@@ -163,7 +161,7 @@ class FilterViewModelTest {
         // Assert
         viewModel.uiState.test {
             val state = awaitItem()
-            assertEquals(state.allProviders.sorted(), state.allProviders)
+            assertSet(state.allProviders, state.allProviders)
             assertEquals(state.selectableProviders.sorted(), state.selectableProviders)
         }
     }
diff --git a/android/lib/common-test/src/main/java/net/mullvad/mullvadvpn/lib/common/test/TestUtils.kt b/android/lib/common-test/src/main/java/net/mullvad/mullvadvpn/lib/common/test/TestUtils.kt
index 60f991fce6..d535019e57 100644
--- a/android/lib/common-test/src/main/java/net/mullvad/mullvadvpn/lib/common/test/TestUtils.kt
+++ b/android/lib/common-test/src/main/java/net/mullvad/mullvadvpn/lib/common/test/TestUtils.kt
@@ -14,3 +14,16 @@ fun <T> assertLists(expected: List<T>, actual: List<T>, message: String? = null)
     """
                 .trimMargin(),
     )
+
+fun <T> assertSet(expected: Set<T>, actual: Set<T>, message: String? = null) =
+    assertTrue(
+        expected.size == actual.size &&
+            expected.containsAll(actual) &&
+            actual.containsAll(expected),
+        message
+            ?: """Expected list should have same size and contains same items.
+        | Expected(${expected.size}): $expected
+        | Actual(${actual.size})  : $actual
+    """
+                .trimMargin(),
+    )
diff --git a/android/lib/daemon-grpc/src/main/kotlin/net/mullvad/mullvadvpn/lib/daemon/grpc/mapper/FromDomain.kt b/android/lib/daemon-grpc/src/main/kotlin/net/mullvad/mullvadvpn/lib/daemon/grpc/mapper/FromDomain.kt
index f62124a171..e4d3e84a4f 100644
--- a/android/lib/daemon-grpc/src/main/kotlin/net/mullvad/mullvadvpn/lib/daemon/grpc/mapper/FromDomain.kt
+++ b/android/lib/daemon-grpc/src/main/kotlin/net/mullvad/mullvadvpn/lib/daemon/grpc/mapper/FromDomain.kt
@@ -46,7 +46,7 @@ internal fun Constraint<RelayItemId>.fromDomain(): ManagementInterface.LocationC
 internal fun Constraint<Providers>.fromDomain(): List<String> =
     when (this) {
         is Constraint.Any -> emptyList()
-        is Constraint.Only -> value.providers.map { it.value }
+        is Constraint.Only -> value.map { it.value }
     }
 
 internal fun DnsOptions.fromDomain(): ManagementInterface.DnsOptions =
diff --git a/android/lib/daemon-grpc/src/main/kotlin/net/mullvad/mullvadvpn/lib/daemon/grpc/mapper/ToDomain.kt b/android/lib/daemon-grpc/src/main/kotlin/net/mullvad/mullvadvpn/lib/daemon/grpc/mapper/ToDomain.kt
index 91399b48b6..daa04fc8d9 100644
--- a/android/lib/daemon-grpc/src/main/kotlin/net/mullvad/mullvadvpn/lib/daemon/grpc/mapper/ToDomain.kt
+++ b/android/lib/daemon-grpc/src/main/kotlin/net/mullvad/mullvadvpn/lib/daemon/grpc/mapper/ToDomain.kt
@@ -365,7 +365,7 @@ internal fun ManagementInterface.GeographicLocationConstraint.toDomain(): GeoLoc
 }
 
 internal fun List<String>.toDomain(): Constraint<Providers> =
-    if (isEmpty()) Constraint.Any else Constraint.Only(Providers(map { ProviderId(it) }.toSet()))
+    if (isEmpty()) Constraint.Any else Constraint.Only(map { ProviderId(it) }.toSet())
 
 internal fun ManagementInterface.WireguardConstraints.toDomain(): WireguardConstraints =
     WireguardConstraints(
diff --git a/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/lib/model/Providers.kt b/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/lib/model/Providers.kt
index 73cf9facdb..0e511c9407 100644
--- a/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/lib/model/Providers.kt
+++ b/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/lib/model/Providers.kt
@@ -1,8 +1,3 @@
 package net.mullvad.mullvadvpn.lib.model
 
-import arrow.optics.optics
-
-@optics
-data class Providers(val providers: Set<ProviderId>) {
-    companion object
-}
+typealias Providers = Set<ProviderId>