Change `Relay` and `RelayCity` to reference parent

3 files changed, 21 insertions, 20 deletions

diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/Relay.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/Relay.kt
index 863dc69d8e..8620a64963 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/Relay.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/Relay.kt
@@ -2,20 +2,15 @@ package net.mullvad.mullvadvpn.relaylist
 
 import net.mullvad.mullvadvpn.model.LocationConstraint
 
-data class Relay(
-    val countryCode: String,
-    val cityCode: String,
-    override val name: String,
-    val cityName: String
-) : RelayItem {
+data class Relay(val city: RelayCity, override val name: String) : RelayItem {
     override val code = name
     override val type = RelayItemType.Relay
-    override val location = LocationConstraint.Hostname(countryCode, cityCode, name)
+    override val location = LocationConstraint.Hostname(city.country.code, city.code, name)
     override val hasChildren = false
 
     override val visibleChildCount = 0
 
-    override val locationName = "$cityName ($name)"
+    override val locationName = "${city.name} ($name)"
 
     override var expanded
         get() = false
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayCity.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayCity.kt
index 15c52fce43..9d4ba65ccf 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayCity.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayCity.kt
@@ -3,14 +3,14 @@ package net.mullvad.mullvadvpn.relaylist
 import net.mullvad.mullvadvpn.model.LocationConstraint
 
 class RelayCity(
+    val country: RelayCountry,
     override val name: String,
-    val countryCode: String,
     override val code: String,
     override var expanded: Boolean,
     val relays: List<Relay>
 ) : RelayItem {
     override val type = RelayItemType.City
-    override val location = LocationConstraint.City(countryCode, code)
+    override val location = LocationConstraint.City(country.code, code)
 
     override val hasChildren
         get() = relays.size > 1
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayList.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayList.kt
index 7c3870c4dd..26cd7ffa58 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayList.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayList.kt
@@ -9,19 +9,25 @@ class RelayList {
     constructor(model: net.mullvad.mullvadvpn.model.RelayList) {
         countries = model.countries
             .map { country ->
-                val cities = country.cities
-                    .map { city -> 
-                        val relays = city.relays
-                            .filter { relay -> relay.hasWireguardTunnels }
-                            .map { relay ->
-                                Relay(country.code, city.code, relay.hostname, city.name)
-                            }
+                val cities = mutableListOf<RelayCity>()
+                val relayCountry = RelayCountry(country.name, country.code, false, cities)
 
-                        RelayCity(city.name, country.code, city.code, false, relays)
+                for (city in country.cities) {
+                    val relays = mutableListOf<Relay>()
+                    val relayCity = RelayCity(relayCountry, city.name, city.code, false, relays)
+
+                    val validCityRelays = city.relays.filter { relay -> relay.hasWireguardTunnels }
+
+                    for (relay in validCityRelays) {
+                        relays.add(Relay(relayCity, relay.hostname))
                     }
-                    .filter { city -> city.relays.isNotEmpty() }
 
-                RelayCountry(country.name, country.code, false, cities)
+                    if (relays.isNotEmpty()) {
+                        cities.add(relayCity)
+                    }
+                }
+
+                relayCountry
             }
             .filter { country -> country.cities.isNotEmpty() }
     }