Merge branch 'pj/fix/tv-city-connection-issues' into 'develop'

pawel-jurczyk · pawel-jurczyk · commit e8dd0dd2e25f · 2025-12-17T10:40:52.000+01:00
Merge-Request: apple/vpn/protonvpn!2836
Approved-by: Chris Janusiewicz &lt;krzysztof.janusiewicz@proton.ch&gt;
Approved-by: John Biggs &lt;john.biggs@proton.ch&gt;
diff --git a/apps/tvos/TestPlans/ProtonVPN-tvOS-Unit-All.xctestplan b/apps/tvos/TestPlans/ProtonVPN-tvOS-Unit-All.xctestplan
@@ -9,20 +9,6 @@
     }
   ],
   "defaultOptions" : {
-    "codeCoverage" : {
-      "targets" : [
-        {
-          "containerPath" : "container:..\/..\/libraries\/ConnectionDetails",
-          "identifier" : "ConnectionDetails",
-          "name" : "ConnectionDetails"
-        },
-        {
-          "containerPath" : "container:..\/..\/libraries\/tvOS",
-          "identifier" : "tvOS",
-          "name" : "tvOS"
-        }
-      ]
-    },
     "language" : "en",
     "region" : "US",
     "targetForVariableExpansion" : {
diff --git a/libraries/Features/tvos_app/Sources/tvos_app/Features/Home/CountryList/CountryListView.swift b/libraries/Features/tvos_app/Sources/tvos_app/Features/Home/CountryList/CountryListView.swift
@@ -17,6 +17,7 @@
 //  along with ProtonVPN.  If not, see <https://www.gnu.org/licenses/>.
 
 import ComposableArchitecture
+import Connection
 import Foundation
 import ProtonCoreUIFoundations
 import SwiftUI
@@ -29,6 +30,8 @@ struct CountryListView: View {
     // Watch which item is focused to highlight selected row
     @FocusState private var focusedIndex: ItemCoordinate?
 
+    @SharedReader(.connectionState) var connectionState: ConnectionState
+
     static let columnCount = 6
     private static let gridItemWidth: Double = 210
     private static let unfocusedOpacity: Double = 0.5 // "Unfocused" items are half transparent
@@ -86,7 +89,6 @@ struct CountryListView: View {
         } primaryAction: {
             store.send(.selectItem(.country(code: item.code)))
         }
-        .menuStyle(.borderlessButton)
         .buttonStyle(CountryListButtonStyle())
         .padding(.top, .themeSpacing8)
         .padding(.bottom, .themeSpacing32)
@@ -108,7 +110,13 @@ struct CountryListView: View {
             Button {
                 store.send(.selectItem(.city(name: city, code: item.code)))
             } label: {
-                Text(city)
+                if case let .connected(intent, _, _, _) = connectionState,
+                   case let .city(name, code) = intent.spec.location,
+                   item.code == code, city == name {
+                    Label(city, systemImage: "lock.fill")
+                } else {
+                    Text(city)
+                }
             }
         }
     }
diff --git a/libraries/Features/tvos_app/Sources/tvos_app/Features/Main/MainFeature.swift b/libraries/Features/tvos_app/Sources/tvos_app/Features/Main/MainFeature.swift
@@ -53,6 +53,8 @@ struct MainFeature {
         case homeLoading(HomeLoadingFeature.Action)
         case settings(SettingsFeature.Action)
 
+        case userSelectedItem(ServerGroupInfo.Kind)
+
         case onAppear
         case onLogout
         case updateUserLocation
@@ -113,8 +115,11 @@ struct MainFeature {
                 return .none
 
             case let .homeLoading(.loaded(.countryList(.selectItem(kind)))):
+                return .send(.userSelectedItem(kind))
+
+            case let .userSelectedItem(kind):
                 let (code, cityName) = kind.selectedItem
-                switch handleConnectionIntent(to: code, currentConnectionState: state.connectionState) {
+                switch handleConnectionIntent(to: kind, currentConnectionState: state.connectionState) {
                 case .connect:
                     return .send(.connectDisconnectingIfNecessary(countryCode: code, cityName: cityName))
                 case .disconnect:
@@ -187,31 +192,45 @@ struct MainFeature {
     }
 
     private func handleConnectionIntent(
-        to targetCountryCode: String,
+        to target: ServerGroupInfo.Kind,
         currentConnectionState: ConnectionState
     ) -> ConnectionStrategy {
-        guard let activeCountryCode = activeCountryCode(from: currentConnectionState) else {
+        guard let activeKind = activeKind(from: currentConnectionState) else {
             // If we're not already connecting/connected, we can just connect to the selected country
             return .connect
         }
-        if targetCountryCode == activeCountryCode {
-            // If the selected country is the same as the connecting/connected one, disconnect
+        if target == activeKind {
+            // If the selected kind is the same as the connecting/connected one, disconnect
             return .disconnect
-        } else {
-            // Otherwise, proceed with connection to the selected country
-            return .connect
         }
+        if case let .country(targetCode) = target,
+           case let .city(_, activeCode) = activeKind {
+            // if target is country and active kind is city in that country, also disconnect
+            return targetCode == activeCode ? .disconnect : .connect
+        }
+
+        // Otherwise, proceed with connection to the selected country
+        return .connect
     }
 
-    private func activeCountryCode(from connectionState: ConnectionState) -> String? {
-        if case let .connected(_, server, _, _) = connectionState {
-            return server.logical.exitCountryCode
+    private func activeKind(from connectionState: ConnectionState) -> ServerGroupInfo.Kind? {
+        if case let .connected(intent, server, _, _) = connectionState {
+            switch intent.spec.location {
+            case let .city(name, code):
+                return .city(name: name, code: code)
+            default:
+                return .country(code: server.logical.exitCountryCode)
+            }
         }
         if case let .connecting(.unresolved(intent)) = connectionState {
-            return intent.spec.countryCode
+            if let code = intent.spec.countryCode {
+                return .country(code: code)
+            } else {
+                return nil
+            }
         }
         if case let .connecting(.resolved(_, server)) = connectionState {
-            return server.logical.exitCountryCode
+            return .country(code: server.logical.exitCountryCode)
         }
         return nil
     }
@@ -227,6 +246,19 @@ struct MainFeature {
     }
 }
 
+private extension ServerGroupInfo.Kind {
+    var code: String? {
+        switch self {
+        case let .city(_, code):
+            code
+        case let .country(code):
+            code
+        case let .gateway(name):
+            name
+        }
+    }
+}
+
 enum ServerResolutionError: Error {
     case noActiveServers(String)
     case serverHasNoEndpoints
diff --git a/libraries/Features/tvos_app/Tests/tvos_appTests/MainFeatureTests.swift b/libraries/Features/tvos_app/Tests/tvos_appTests/MainFeatureTests.swift
@@ -74,6 +74,112 @@ final class MainFeatureTests: XCTestCase {
         await clock.advance(by: .seconds(1))
     }
 
+    @MainActor
+    func testUserClickedCountryItemShouldConnect() async {
+        let store = TestStore(initialState: MainFeature.State()) {
+            MainFeature()
+        } withDependencies: {
+            $0.serverRepository = .empty()
+        }
+        await store.send(.userSelectedItem(.country(code: "PL")))
+        await store.receive {
+            if case .connectDisconnectingIfNecessary("PL", nil) = $0 {
+                true
+            } else {
+                false
+            }
+        }
+        await store.receive(\.connection.input.connect)
+        await store.receive(\.connection.delegate.intentResolutionFailed)
+    }
+
+    @MainActor
+    func testUserClickedCityItemShouldConnect() async {
+        let store = TestStore(initialState: MainFeature.State()) {
+            MainFeature()
+        } withDependencies: {
+            $0.serverRepository = .empty()
+        }
+        await store.send(.userSelectedItem(.city(name: "Warsaw", code: "PL")))
+        await store.receive {
+            if case .connectDisconnectingIfNecessary("PL", "Warsaw") = $0 {
+                true
+            } else {
+                false
+            }
+        }
+        await store.receive(\.connection.input.connect)
+        await store.receive(\.connection.delegate.intentResolutionFailed)
+    }
+
+    @MainActor
+    func testUserClickedCityItemWhileConnectedToThatCityShouldDisconnect() async {
+        let connectionState = ConnectionState.connected(.mock(withSpecLocation: .city(name: "irrelevant", code: "CA")), .ca, .now, nil)
+        let store = TestStore(initialState: MainFeature.State(connectionState: connectionState)) {
+            MainFeature()
+        } withDependencies: {
+            $0.serverRepository = .empty()
+        }
+        await store.send(.userSelectedItem(.city(name: "irrelevant", code: "CA")))
+        await store.receive(\.connection.input.disconnect) {
+            $0.connection.core.shouldDisconnectWhenAllowed = true
+        }
+    }
+
+    @MainActor
+    func testUserClickedCityItemWhileConnectedToAnotherCityInTheSameCountryShouldConnect() async {
+        let connectionState = ConnectionState.connected(.mock(withSpecLocation: .city(name: "irrelevant", code: "CA")), .ca, .now, nil)
+        let store = TestStore(initialState: MainFeature.State(connectionState: connectionState)) {
+            MainFeature()
+        } withDependencies: {
+            $0.serverRepository = .empty()
+        }
+        await store.send(.userSelectedItem(.city(name: "not irrelevant", code: "CA")))
+        await store.receive {
+            if case .connectDisconnectingIfNecessary("CA", "not irrelevant") = $0 {
+                true
+            } else {
+                false
+            }
+        }
+        await store.receive(\.connection.input.connect)
+        await store.receive(\.connection.delegate.intentResolutionFailed)
+    }
+
+    @MainActor
+    func testUserClickedCountryItemWhileConnectedToCityInAnotherCountryShouldConnect() async {
+        let connectionState = ConnectionState.connected(.mock(withSpecLocation: .city(name: "irrelevant", code: "CA")), .ca, .now, nil)
+        let store = TestStore(initialState: MainFeature.State(connectionState: connectionState)) {
+            MainFeature()
+        } withDependencies: {
+            $0.serverRepository = .empty()
+        }
+        await store.send(.userSelectedItem(.country(code: "US")))
+        await store.receive {
+            if case .connectDisconnectingIfNecessary("US", nil) = $0 {
+                true
+            } else {
+                false
+            }
+        }
+        await store.receive(\.connection.input.connect)
+        await store.receive(\.connection.delegate.intentResolutionFailed)
+    }
+
+    @MainActor
+    func testUserClickedCountryItemWhileConnectedToCityInThatCountryShouldDisconnect() async {
+        let connectionState = ConnectionState.connected(.mock(withSpecLocation: .city(name: "irrelevant", code: "CA")), .ca, .now, nil)
+        let store = TestStore(initialState: MainFeature.State(connectionState: connectionState)) {
+            MainFeature()
+        } withDependencies: {
+            $0.serverRepository = .empty()
+        }
+        await store.send(.userSelectedItem(.country(code: "CA")))
+        await store.receive(\.connection.input.disconnect) {
+            $0.connection.core.shouldDisconnectWhenAllowed = true
+        }
+    }
+
     @MainActor
     func testUserClickedConnect() async {
         let clock = TestClock()
