From 54e3a308888e012d4f8769755092086bb489b25c Mon Sep 17 00:00:00 2001
From: Jon Petersson <jon.petersson@kvadrat.se>
Date: Tue, 14 May 2024 09:57:51 +0200
Subject: [PATCH] Remove some unnecessary LocationNode subclasses

---
 .../AllLocationDataSource.swift               |  6 ++---
 .../SelectLocation/LocationCell.swift         | 25 +++++++++++--------
 .../SelectLocation/LocationNode.swift         |  6 -----
 .../SelectLocation/LocationNodeTests.swift    |  6 ++---
 4 files changed, 21 insertions(+), 22 deletions(-)

diff --git a/ios/MullvadVPN/View controllers/SelectLocation/AllLocationDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/AllLocationDataSource.swift
index f57e980deadc..578deeddd2ca 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/AllLocationDataSource.swift	
+++ b/ios/MullvadVPN/View controllers/SelectLocation/AllLocationDataSource.swift	
@@ -60,7 +60,7 @@ class AllLocationDataSource: LocationDataSourceProtocol {
     ) {
         switch location {
         case let .country(countryCode):
-            let countryNode = CountryLocationNode(
+            let countryNode = LocationNode(
                 name: serverLocation.country,
                 code: LocationNode.combineNodeCodes([countryCode]),
                 locations: [location],
@@ -73,7 +73,7 @@ class AllLocationDataSource: LocationDataSourceProtocol {
             }
 
         case let .city(countryCode, cityCode):
-            let cityNode = CityLocationNode(
+            let cityNode = LocationNode(
                 name: serverLocation.city,
                 code: LocationNode.combineNodeCodes([countryCode, cityCode]),
                 locations: [location],
@@ -88,7 +88,7 @@ class AllLocationDataSource: LocationDataSourceProtocol {
             }
 
         case let .hostname(countryCode, cityCode, hostCode):
-            let hostNode = HostLocationNode(
+            let hostNode = LocationNode(
                 name: relay.hostname,
                 code: LocationNode.combineNodeCodes([hostCode]),
                 locations: [location],
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift
index e777ca5e21d4..a51cebf59c86 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift	
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift	
@@ -316,18 +316,23 @@ extension LocationCell {
         checkboxButton.accessibilityIdentifier = .customListLocationCheckmarkButton
         checkboxButton.isSelected = item.isSelected
         checkboxButton.tintColor = item.isSelected ? .successColor : .white
+        accessibilityValue = item.node.code
 
-        if item.node is CountryLocationNode {
-            accessibilityIdentifier = .countryLocationCell
-            accessibilityValue = item.node.code
-        } else if item.node is CityLocationNode {
-            accessibilityIdentifier = .cityLocationCell
-            accessibilityValue = item.node.code
-        } else if item.node is HostLocationNode {
-            accessibilityIdentifier = .relayLocationCell
-            accessibilityValue = item.node.code
-        } else if item.node is CustomListLocationNode {
+        if item.node is CustomListLocationNode {
             accessibilityIdentifier = .customListLocationCell
+        } else {
+            // Only custom list nodes have more than one location. Therefore checking first
+            // location here is fine.
+            switch item.node.locations.first {
+            case .country:
+                accessibilityIdentifier = .countryLocationCell
+            case .city:
+                accessibilityIdentifier = .cityLocationCell
+            case .hostname:
+                accessibilityIdentifier = .relayLocationCell
+            case nil:
+                break
+            }
         }
 
         setBehavior(behavior)
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationNode.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationNode.swift
index 4534ac4d8e01..8a43493f0b8d 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationNode.swift	
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationNode.swift	
@@ -168,9 +168,3 @@ class CustomListLocationNode: LocationNode {
         )
     }
 }
-
-class CountryLocationNode: LocationNode {}
-
-class CityLocationNode: LocationNode {}
-
-class HostLocationNode: LocationNode {}
diff --git a/ios/MullvadVPNTests/MullvadVPN/View controllers/SelectLocation/LocationNodeTests.swift b/ios/MullvadVPNTests/MullvadVPN/View controllers/SelectLocation/LocationNodeTests.swift
index 9495129d0088..c1ba7c57be41 100644
--- a/ios/MullvadVPNTests/MullvadVPN/View controllers/SelectLocation/LocationNodeTests.swift	
+++ b/ios/MullvadVPNTests/MullvadVPN/View controllers/SelectLocation/LocationNodeTests.swift	
@@ -16,9 +16,9 @@ class LocationNodeTests: XCTestCase {
         showsChildren: false,
         customList: CustomList(name: "List", locations: [])
     )
-    let countryNode = CountryLocationNode(name: "Country", code: "country", showsChildren: false)
-    let cityNode = CityLocationNode(name: "City", code: "city", showsChildren: false)
-    let hostNode = HostLocationNode(name: "Host", code: "host", showsChildren: false)
+    let countryNode = LocationNode(name: "Country", code: "country", showsChildren: false)
+    let cityNode = LocationNode(name: "City", code: "city", showsChildren: false)
+    let hostNode = LocationNode(name: "Host", code: "host", showsChildren: false)
 
     override func setUp() async throws {
         createNodeTree()