2 files changed, 20 insertions, 33 deletions

diff --git a/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/lib/model/RelayItem.kt b/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/lib/model/RelayItem.kt
index a31a6f67df..17bc563a8d 100644
--- a/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/lib/model/RelayItem.kt
+++ b/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/lib/model/RelayItem.kt
@@ -2,28 +2,26 @@ package net.mullvad.mullvadvpn.lib.model
 
 import arrow.optics.optics
 
+typealias DomainCustomList = CustomList
+
 @optics
 sealed interface RelayItem {
     val id: RelayItemId
     val name: String
+
     val active: Boolean
     val hasChildren: Boolean
-    val expanded: Boolean
 
     @optics
     data class CustomList(
-        override val id: CustomListId,
-        val customListName: CustomListName,
+        val customList: DomainCustomList,
         val locations: List<Location>,
-        override val expanded: Boolean,
     ) : RelayItem {
-        override val name: String = customListName.value
+        override val name: String = customList.name.value
+        override val id = customList.id
 
-        override val active
-            get() = locations.any { location -> location.active }
-
-        override val hasChildren
-            get() = locations.isNotEmpty()
+        override val active = locations.any { it.active }
+        override val hasChildren: Boolean = locations.isNotEmpty()
 
         companion object
     }
@@ -36,16 +34,11 @@ sealed interface RelayItem {
         data class Country(
             override val id: GeoLocationId.Country,
             override val name: String,
-            override val expanded: Boolean,
             val cities: List<City>
         ) : Location {
             val relays = cities.flatMap { city -> city.relays }
-
-            override val active
-                get() = cities.any { city -> city.active }
-
-            override val hasChildren
-                get() = cities.isNotEmpty()
+            override val active = cities.any { it.active }
+            override val hasChildren: Boolean = cities.isNotEmpty()
 
             companion object
         }
@@ -54,15 +47,10 @@ sealed interface RelayItem {
         data class City(
             override val id: GeoLocationId.City,
             override val name: String,
-            override val expanded: Boolean,
             val relays: List<Relay>
         ) : Location {
-
-            override val active
-                get() = relays.any { relay -> relay.active }
-
-            override val hasChildren
-                get() = relays.isNotEmpty()
+            override val active = relays.any { it.active }
+            override val hasChildren: Boolean = relays.isNotEmpty()
 
             companion object
         }
@@ -73,10 +61,8 @@ sealed interface RelayItem {
             val provider: Provider,
             override val active: Boolean,
         ) : Location {
-            override val name: String = id.hostname
-
-            override val hasChildren = false
-            override val expanded = false
+            override val name: String = id.code
+            override val hasChildren: Boolean = false
 
             companion object
         }
diff --git a/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/lib/model/RelayItemId.kt b/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/lib/model/RelayItemId.kt
index c560fab49c..85326d8d2a 100644
--- a/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/lib/model/RelayItemId.kt
+++ b/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/lib/model/RelayItemId.kt
@@ -21,28 +21,29 @@ value class CustomListId(val value: String) : RelayItemId, Parcelable {
 sealed interface GeoLocationId : RelayItemId, Parcelable {
     @optics
     @Parcelize
-    data class Country(val countryCode: String) : GeoLocationId {
+    data class Country(override val code: String) : GeoLocationId {
         companion object
     }
 
     @optics
     @Parcelize
-    data class City(val countryCode: Country, val cityCode: String) : GeoLocationId {
+    data class City(override val country: Country, override val code: String) : GeoLocationId {
         companion object
     }
 
     @optics
     @Parcelize
-    data class Hostname(val city: City, val hostname: String) : GeoLocationId {
+    data class Hostname(val city: City, override val code: String) : GeoLocationId {
         companion object
     }
 
+    val code: String
     val country: Country
         get() =
             when (this) {
                 is Country -> this
-                is City -> this.countryCode
-                is Hostname -> this.city.countryCode
+                is City -> this.country
+                is Hostname -> this.city.country
             }
 
     companion object