diff --git a/ios/MullvadVPN/View controllers/SelectLocation/CustomListsDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/CustomListsDataSource.swift index 7d5639e00d4e..99ad974f2086 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 8be5c295805a..ed639cc219b1 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 2b6fc5b5e829..4431498b7d7d 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)