Fix selection of custom list location not being retained

Author: Jon Petersson

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)
         }
     }
 }

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,

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)