Merge branch 'fix-selection-of-custom-list-location-not-being-retained-ios-569'

3 files changed, 22 insertions, 5 deletions

diff --git a/ios/MullvadVPN/View controllers/SelectLocation/CustomListsDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/CustomListsDataSource.swift
index 7d5639e00d..99ad974f20 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/CustomListsDataSource.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/CustomListsDataSource.swift
@@ -85,18 +85,21 @@ class CustomListsDataSource: LocationDataSourceProtocol {
         return switch location {
         case let .country(countryCode):
             rootNode
-                .countryFor(code: countryCode)?.copy(withParent: parentNode)
+                .countryFor(code: countryCode)?
+                .copy(withParent: parentNode)
 
         case let .city(countryCode, cityCode):
             rootNode
-                .countryFor(code: countryCode)?.copy(withParent: parentNode)
-                .cityFor(codes: [countryCode, cityCode])
+                .countryFor(code: countryCode)?
+                .cityFor(codes: [countryCode, cityCode])?
+                .copy(withParent: parentNode)
 
         case let .hostname(countryCode, cityCode, hostCode):
             rootNode
-                .countryFor(code: countryCode)?.copy(withParent: parentNode)
+                .countryFor(code: countryCode)?
                 .cityFor(codes: [countryCode, cityCode])?
-                .hostFor(code: hostCode)
+                .hostFor(code: hostCode)?
+                .copy(withParent: parentNode)
         }
     }
 }
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationNode.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationNode.swift
index 8be5c29580..ed639cc219 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationNode.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationNode.swift
@@ -80,6 +80,8 @@ extension LocationNode {
 }
 
 extension LocationNode {
+    /// Recursively copies a node, its parent and its descendants from another
+    /// node (tree), with an optional custom root parent.
     func copy(withParent parent: LocationNode? = nil) -> LocationNode {
         let node = LocationNode(
             name: name,
diff --git a/ios/MullvadVPNTests/Location/CustomListsDataSourceTests.swift b/ios/MullvadVPNTests/Location/CustomListsDataSourceTests.swift
index 2b6fc5b5e8..4431498b7d 100644
--- a/ios/MullvadVPNTests/Location/CustomListsDataSourceTests.swift
+++ b/ios/MullvadVPNTests/Location/CustomListsDataSourceTests.swift
@@ -32,6 +32,18 @@ class CustomListsDataSourceTests: XCTestCase {
         XCTAssertNotNil(youtubeNode.descendantNodeFor(codes: ["youtube", "us", "dal"]))
     }
 
+    func testParents() throws {
+        let listNode = try XCTUnwrap(dataSource.nodes.first(where: { $0.name == "Netflix" }))
+        let countryNode = try XCTUnwrap(listNode.descendantNodeFor(codes: ["netflix-se"]))
+        let cityNode = try XCTUnwrap(listNode.descendantNodeFor(codes: ["netflix-se-got"]))
+        let hostNode = try XCTUnwrap(listNode.descendantNodeFor(codes: ["netflix-se10-wireguard"]))
+
+        XCTAssertNil(listNode.parent)
+        XCTAssertEqual(countryNode.parent, listNode)
+        XCTAssertEqual(cityNode.parent, countryNode)
+        XCTAssertEqual(hostNode.parent, cityNode)
+    }
+
     func testSearch() throws {
         let nodes = dataSource.search(by: "got")
         let rootNode = RootLocationNode(children: nodes)