Support new tunnel state API in the Android frontend.

Remove `get_current_location` from jni.

18 files changed, 64 insertions, 171 deletions

diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ConnectScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ConnectScreenTest.kt
index dae3bf2ed3..3838bdc7a0 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ConnectScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/ConnectScreenTest.kt
@@ -79,7 +79,6 @@ class ConnectScreenTest {
                         inAddress = null,
                         outAddress = "",
                         showLocation = false,
-                        isTunnelInfoExpanded = false,
                         deviceName = "",
                         daysLeftUntilExpiry = null,
                         inAppNotification = InAppNotification.TunnelStateBlocked,
@@ -115,7 +114,6 @@ class ConnectScreenTest {
                         inAddress = null,
                         outAddress = "",
                         showLocation = false,
-                        isTunnelInfoExpanded = false,
                         deviceName = "",
                         daysLeftUntilExpiry = null,
                         inAppNotification = InAppNotification.TunnelStateBlocked,
@@ -149,7 +147,6 @@ class ConnectScreenTest {
                         inAddress = null,
                         outAddress = "",
                         showLocation = false,
-                        isTunnelInfoExpanded = false,
                         deviceName = "",
                         daysLeftUntilExpiry = null,
                         inAppNotification = null,
@@ -182,7 +179,6 @@ class ConnectScreenTest {
                         inAddress = null,
                         outAddress = "",
                         showLocation = false,
-                        isTunnelInfoExpanded = false,
                         deviceName = "",
                         daysLeftUntilExpiry = null,
                         inAppNotification = null,
@@ -216,7 +212,6 @@ class ConnectScreenTest {
                         inAddress = null,
                         outAddress = "",
                         showLocation = true,
-                        isTunnelInfoExpanded = false,
                         deviceName = "",
                         daysLeftUntilExpiry = null,
                         inAppNotification = null,
@@ -245,12 +240,11 @@ class ConnectScreenTest {
                     ConnectUiState(
                         location = null,
                         relayLocation = mockRelayLocation,
-                        tunnelUiState = TunnelState.Disconnected,
-                        tunnelRealState = TunnelState.Disconnected,
+                        tunnelUiState = TunnelState.Disconnected(),
+                        tunnelRealState = TunnelState.Disconnected(),
                         inAddress = null,
                         outAddress = "",
                         showLocation = true,
-                        isTunnelInfoExpanded = false,
                         deviceName = "",
                         daysLeftUntilExpiry = null,
                         inAppNotification = null,
@@ -286,7 +280,6 @@ class ConnectScreenTest {
                         inAddress = null,
                         outAddress = "",
                         showLocation = true,
-                        isTunnelInfoExpanded = false,
                         deviceName = "",
                         daysLeftUntilExpiry = null,
                         inAppNotification =
@@ -326,7 +319,6 @@ class ConnectScreenTest {
                         inAddress = null,
                         outAddress = "",
                         showLocation = true,
-                        isTunnelInfoExpanded = false,
                         deviceName = "",
                         daysLeftUntilExpiry = null,
                         inAppNotification =
@@ -363,7 +355,6 @@ class ConnectScreenTest {
                         inAddress = null,
                         outAddress = "",
                         showLocation = false,
-                        isTunnelInfoExpanded = false,
                         deviceName = "",
                         daysLeftUntilExpiry = null,
                         inAppNotification = InAppNotification.TunnelStateBlocked,
@@ -399,7 +390,6 @@ class ConnectScreenTest {
                         inAddress = null,
                         outAddress = "",
                         showLocation = true,
-                        isTunnelInfoExpanded = false,
                         deviceName = "",
                         daysLeftUntilExpiry = null,
                         inAppNotification = InAppNotification.TunnelStateBlocked,
@@ -430,12 +420,11 @@ class ConnectScreenTest {
                     ConnectUiState(
                         location = null,
                         relayLocation = mockRelayLocation,
-                        tunnelUiState = TunnelState.Disconnected,
-                        tunnelRealState = TunnelState.Disconnected,
+                        tunnelUiState = TunnelState.Disconnected(),
+                        tunnelRealState = TunnelState.Disconnected(),
                         inAddress = null,
                         outAddress = "",
                         showLocation = false,
-                        isTunnelInfoExpanded = false,
                         deviceName = "",
                         daysLeftUntilExpiry = null,
                         inAppNotification = null,
@@ -468,7 +457,6 @@ class ConnectScreenTest {
                         inAddress = null,
                         outAddress = "",
                         showLocation = false,
-                        isTunnelInfoExpanded = false,
                         deviceName = "",
                         daysLeftUntilExpiry = null,
                         inAppNotification = null,
@@ -501,7 +489,6 @@ class ConnectScreenTest {
                         inAddress = null,
                         outAddress = "",
                         showLocation = false,
-                        isTunnelInfoExpanded = false,
                         deviceName = "",
                         daysLeftUntilExpiry = null,
                         inAppNotification = null,
@@ -528,12 +515,11 @@ class ConnectScreenTest {
                     ConnectUiState(
                         location = null,
                         relayLocation = null,
-                        tunnelUiState = TunnelState.Disconnected,
-                        tunnelRealState = TunnelState.Disconnected,
+                        tunnelUiState = TunnelState.Disconnected(),
+                        tunnelRealState = TunnelState.Disconnected(),
                         inAddress = null,
                         outAddress = "",
                         showLocation = false,
-                        isTunnelInfoExpanded = false,
                         deviceName = "",
                         daysLeftUntilExpiry = null,
                         inAppNotification = null,
@@ -565,7 +551,6 @@ class ConnectScreenTest {
                         inAddress = null,
                         outAddress = "",
                         showLocation = false,
-                        isTunnelInfoExpanded = false,
                         deviceName = "",
                         daysLeftUntilExpiry = null,
                         inAppNotification = null,
@@ -583,39 +568,6 @@ class ConnectScreenTest {
     }
 
     @Test
-    fun testToggleTunnelInfo() {
-        // Arrange
-        val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-        val dummyLocation = GeoIpLocation(null, null, "dummy country", null, "dummy hostname")
-        composeTestRule.setContentWithTheme {
-            ConnectScreen(
-                uiState =
-                    ConnectUiState(
-                        location = dummyLocation,
-                        relayLocation = null,
-                        tunnelUiState = TunnelState.Connecting(null, null),
-                        tunnelRealState = TunnelState.Connecting(null, null),
-                        inAddress = null,
-                        outAddress = "",
-                        showLocation = false,
-                        isTunnelInfoExpanded = false,
-                        deviceName = "",
-                        daysLeftUntilExpiry = null,
-                        inAppNotification = null,
-                        isPlayBuild = false
-                    ),
-                onToggleTunnelInfo = mockedClickHandler
-            )
-        }
-
-        // Act
-        composeTestRule.onNodeWithTag(LOCATION_INFO_TEST_TAG).performClick()
-
-        // Assert
-        verify { mockedClickHandler.invoke() }
-    }
-
-    @Test
     fun showLocationInfo() {
         // Arrange
         val mockLocation: GeoIpLocation = mockk(relaxed = true)
@@ -638,7 +590,6 @@ class ConnectScreenTest {
                         inAddress = mockInAddress,
                         outAddress = mockOutAddress,
                         showLocation = false,
-                        isTunnelInfoExpanded = true,
                         deviceName = "",
                         daysLeftUntilExpiry = null,
                         inAppNotification = null,
@@ -647,6 +598,8 @@ class ConnectScreenTest {
             )
         }
 
+        composeTestRule.onNodeWithTag(LOCATION_INFO_TEST_TAG).performClick()
+
         // Assert
         composeTestRule.apply {
             onNodeWithText(mockHostName).assertExists()
@@ -677,7 +630,6 @@ class ConnectScreenTest {
                         inAddress = null,
                         outAddress = "",
                         showLocation = false,
-                        isTunnelInfoExpanded = false,
                         deviceName = "",
                         daysLeftUntilExpiry = null,
                         inAppNotification = InAppNotification.UpdateAvailable(versionInfo),
@@ -714,7 +666,6 @@ class ConnectScreenTest {
                         inAddress = null,
                         outAddress = "",
                         showLocation = false,
-                        isTunnelInfoExpanded = false,
                         deviceName = "",
                         daysLeftUntilExpiry = null,
                         inAppNotification = InAppNotification.UnsupportedVersion(versionInfo),
@@ -748,7 +699,6 @@ class ConnectScreenTest {
                         inAddress = null,
                         outAddress = "",
                         showLocation = false,
-                        isTunnelInfoExpanded = false,
                         deviceName = "",
                         daysLeftUntilExpiry = null,
                         inAppNotification = InAppNotification.AccountExpiry(expiryDate),
@@ -787,7 +737,6 @@ class ConnectScreenTest {
                         inAddress = null,
                         outAddress = "",
                         showLocation = false,
-                        isTunnelInfoExpanded = false,
                         deviceName = "",
                         daysLeftUntilExpiry = null,
                         inAppNotification = InAppNotification.UnsupportedVersion(versionInfo),
@@ -820,7 +769,6 @@ class ConnectScreenTest {
                         inAddress = null,
                         outAddress = "",
                         showLocation = false,
-                        isTunnelInfoExpanded = false,
                         deviceName = "",
                         daysLeftUntilExpiry = null,
                         inAppNotification = InAppNotification.AccountExpiry(expiryDate),
@@ -838,10 +786,8 @@ class ConnectScreenTest {
 
     @Test
     fun testOpenAccountView() {
-
-        val onAccountClickMockk: () -> Unit = mockk(relaxed = true)
-
         // Arrange
+        val onAccountClickMockk: () -> Unit = mockk(relaxed = true)
         composeTestRule.setContentWithTheme {
             ConnectScreen(uiState = ConnectUiState.INITIAL, onAccountClick = onAccountClickMockk)
         }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/component/ConnectionStatusText.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/component/ConnectionStatusText.kt
index 742302ce91..903bb412dc 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/component/ConnectionStatusText.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/component/ConnectionStatusText.kt
@@ -18,7 +18,7 @@ import net.mullvad.talpid.tunnel.ErrorStateCause
 private fun PreviewConnectionStatusText() {
     AppTheme {
         SpacedColumn {
-            ConnectionStatusText(TunnelState.Disconnected)
+            ConnectionStatusText(TunnelState.Disconnected())
             ConnectionStatusText(TunnelState.Connecting(null, null))
             ConnectionStatusText(
                 state = TunnelState.Error(ErrorState(ErrorStateCause.Ipv6Unavailable, true))
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/ConnectScreen.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/ConnectScreen.kt
index b1de1ad809..646e89c987 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/ConnectScreen.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/ConnectScreen.kt
@@ -19,7 +19,9 @@ import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.collectAsState
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableLongStateOf
+import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
+import androidx.compose.runtime.saveable.rememberSaveable
 import androidx.compose.runtime.setValue
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
@@ -108,7 +110,6 @@ fun Connect(navigator: DestinationsNavigator) {
         onSwitchLocationClick = {
             navigator.navigate(SelectLocationDestination) { launchSingleTop = true }
         },
-        onToggleTunnelInfo = connectViewModel::toggleTunnelInfoExpansion,
         onUpdateVersionClick = {
             val intent =
                 Intent(
@@ -137,7 +138,6 @@ fun ConnectScreen(
     onConnectClick: () -> Unit = {},
     onCancelClick: () -> Unit = {},
     onSwitchLocationClick: () -> Unit = {},
-    onToggleTunnelInfo: () -> Unit = {},
     onUpdateVersionClick: () -> Unit = {},
     onManageAccountClick: () -> Unit = {},
     onSettingsClick: () -> Unit = {},
@@ -233,12 +233,13 @@ fun ConnectScreen(
                 color = MaterialTheme.colorScheme.onPrimary,
                 modifier = Modifier.padding(horizontal = Dimens.sideMargin)
             )
+            var expanded by rememberSaveable { mutableStateOf(false) }
             LocationInfo(
-                onToggleTunnelInfo = onToggleTunnelInfo,
+                onToggleTunnelInfo = { expanded = !expanded },
                 isVisible =
-                    uiState.tunnelRealState != TunnelState.Disconnected &&
+                    uiState.tunnelRealState !is TunnelState.Disconnected &&
                         uiState.location?.hostname != null,
-                isExpanded = uiState.isTunnelInfoExpanded,
+                isExpanded = expanded,
                 location = uiState.location,
                 inAddress = uiState.inAddress,
                 outAddress = uiState.outAddress,
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/OutOfTimeScreen.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/OutOfTimeScreen.kt
index d9071be7d8..d0d0c7460d 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/OutOfTimeScreen.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/OutOfTimeScreen.kt
@@ -62,7 +62,7 @@ private fun PreviewOutOfTimeScreenDisconnected() {
         OutOfTimeScreen(
             uiState =
                 OutOfTimeUiState(
-                    tunnelState = TunnelState.Disconnected,
+                    tunnelState = TunnelState.Disconnected(),
                     "Heroic Frog",
                     showSitePayment = true
                 ),
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/ConnectUiState.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/ConnectUiState.kt
index 54c1a0d7c0..dc26e24741 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/ConnectUiState.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/ConnectUiState.kt
@@ -15,7 +15,6 @@ data class ConnectUiState(
     val outAddress: String,
     val showLocation: Boolean,
     val inAppNotification: InAppNotification?,
-    val isTunnelInfoExpanded: Boolean,
     val deviceName: String?,
     val daysLeftUntilExpiry: Int?,
     val isPlayBuild: Boolean
@@ -25,12 +24,11 @@ data class ConnectUiState(
             ConnectUiState(
                 location = null,
                 relayLocation = null,
-                tunnelUiState = TunnelState.Disconnected,
-                tunnelRealState = TunnelState.Disconnected,
+                tunnelUiState = TunnelState.Disconnected(),
+                tunnelRealState = TunnelState.Disconnected(),
                 inAddress = null,
                 outAddress = "",
                 showLocation = false,
-                isTunnelInfoExpanded = false,
                 inAppNotification = null,
                 deviceName = null,
                 daysLeftUntilExpiry = null,
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/OutOfTimeUiState.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/OutOfTimeUiState.kt
index 54fd414f86..d72e015194 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/OutOfTimeUiState.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/OutOfTimeUiState.kt
@@ -3,7 +3,7 @@ package net.mullvad.mullvadvpn.compose.state
 import net.mullvad.mullvadvpn.model.TunnelState
 
 data class OutOfTimeUiState(
-    val tunnelState: TunnelState = TunnelState.Disconnected,
+    val tunnelState: TunnelState = TunnelState.Disconnected(),
     val deviceName: String = "",
     val showSitePayment: Boolean = false,
     val billingPaymentState: PaymentState? = null,
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/WelcomeUiState.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/WelcomeUiState.kt
index e2673a0ddf..e43cf6bb98 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/WelcomeUiState.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/WelcomeUiState.kt
@@ -3,7 +3,7 @@ package net.mullvad.mullvadvpn.compose.state
 import net.mullvad.mullvadvpn.model.TunnelState
 
 data class WelcomeUiState(
-    val tunnelState: TunnelState = TunnelState.Disconnected,
+    val tunnelState: TunnelState = TunnelState.Disconnected(),
     val accountNumber: String? = null,
     val deviceName: String? = null,
     val showSitePayment: Boolean = false,
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ConnectionProxy.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ConnectionProxy.kt
index d51bad461d..bbc267b2fa 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ConnectionProxy.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ConnectionProxy.kt
@@ -19,8 +19,8 @@ const val ANTICIPATED_STATE_TIMEOUT_MS = 1500L
 class ConnectionProxy(private val connection: Messenger, eventDispatcher: EventDispatcher) {
     private var resetAnticipatedStateJob: Job? = null
 
-    val onStateChange = EventNotifier<TunnelState>(TunnelState.Disconnected)
-    val onUiStateChange = EventNotifier<TunnelState>(TunnelState.Disconnected)
+    val onStateChange = EventNotifier<TunnelState>(TunnelState.Disconnected())
+    val onUiStateChange = EventNotifier<TunnelState>(TunnelState.Disconnected())
 
     var state by onStateChange.notifiable()
         private set
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/LocationInfoCache.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/LocationInfoCache.kt
deleted file mode 100644
index 48f77d397d..0000000000
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/LocationInfoCache.kt
+++ /dev/null
@@ -1,26 +0,0 @@
-package net.mullvad.mullvadvpn.ui.serviceconnection
-
-import kotlin.properties.Delegates.observable
-import net.mullvad.mullvadvpn.lib.ipc.Event
-import net.mullvad.mullvadvpn.lib.ipc.EventDispatcher
-import net.mullvad.mullvadvpn.model.GeoIpLocation
-
-class LocationInfoCache(eventDispatcher: EventDispatcher) {
-    private var location: GeoIpLocation? by
-        observable(null) { _, _, newLocation -> onNewLocation?.invoke(newLocation) }
-
-    var onNewLocation by
-        observable<((GeoIpLocation?) -> Unit)?>(null) { _, _, callback ->
-            callback?.invoke(location)
-        }
-
-    init {
-        eventDispatcher.registerHandler(Event.NewLocation::class) { event ->
-            location = event.location
-        }
-    }
-
-    fun onDestroy() {
-        onNewLocation = null
-    }
-}
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionContainer.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionContainer.kt
index ca156bed66..8aabe6c9f5 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionContainer.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionContainer.kt
@@ -28,7 +28,6 @@ class ServiceConnectionContainer(
     val authTokenCache = AuthTokenCache(connection, dispatcher)
     val connectionProxy = ConnectionProxy(connection, dispatcher)
     val deviceDataSource = ServiceConnectionDeviceDataSource(connection, dispatcher)
-    val locationInfoCache = LocationInfoCache(dispatcher)
     val settingsListener = SettingsListener(connection, dispatcher)
 
     val splitTunneling = SplitTunneling(connection, dispatcher)
@@ -62,7 +61,6 @@ class ServiceConnectionContainer(
 
         authTokenCache.onDestroy()
         connectionProxy.onDestroy()
-        locationInfoCache.onDestroy()
         settingsListener.onDestroy()
         voucherRedeemer.onDestroy()
 
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/TunnelStateNotificationUseCase.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/TunnelStateNotificationUseCase.kt
index f228bd7dbe..dec794c86c 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/TunnelStateNotificationUseCase.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/TunnelStateNotificationUseCase.kt
@@ -39,7 +39,7 @@ class TunnelStateNotificationUseCase(
             }
             is TunnelState.Error -> InAppNotification.TunnelStateError(tunnelUiState.errorState)
             is TunnelState.Connected,
-            TunnelState.Disconnected -> null
+            is TunnelState.Disconnected -> null
         }
 
     private fun ConnectionProxy.tunnelUiStateFlow(): Flow<TunnelState> =
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/ConnectViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/ConnectViewModel.kt
index 3066006083..d25f360b51 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/ConnectViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/ConnectViewModel.kt
@@ -5,30 +5,30 @@ import androidx.lifecycle.viewModelScope
 import kotlinx.coroutines.FlowPreview
 import kotlinx.coroutines.channels.BufferOverflow
 import kotlinx.coroutines.channels.Channel
-import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.SharedFlow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.callbackFlow
 import kotlinx.coroutines.flow.debounce
 import kotlinx.coroutines.flow.emptyFlow
+import kotlinx.coroutines.flow.filter
+import kotlinx.coroutines.flow.filterIsInstance
 import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.flow.receiveAsFlow
 import kotlinx.coroutines.flow.shareIn
 import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
 import net.mullvad.mullvadvpn.compose.state.ConnectUiState
+import net.mullvad.mullvadvpn.model.GeoIpLocation
 import net.mullvad.mullvadvpn.model.TunnelState
 import net.mullvad.mullvadvpn.repository.AccountRepository
 import net.mullvad.mullvadvpn.repository.DeviceRepository
 import net.mullvad.mullvadvpn.repository.InAppNotificationController
 import net.mullvad.mullvadvpn.ui.serviceconnection.ConnectionProxy
-import net.mullvad.mullvadvpn.ui.serviceconnection.LocationInfoCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionContainer
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManager
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionState
@@ -71,48 +71,46 @@ class ConnectViewModel(
             }
             .shareIn(viewModelScope, SharingStarted.WhileSubscribed())
 
-    private val _isTunnelInfoExpanded = MutableStateFlow(false)
-
     val uiState: StateFlow<ConnectUiState> =
         _shared
             .flatMapLatest { serviceConnection ->
                 combine(
-                    serviceConnection.locationInfoCache.locationCallbackFlow(),
                     relayListUseCase.selectedRelayItem(),
                     inAppNotificationController.notifications,
                     serviceConnection.connectionProxy.tunnelUiStateFlow(),
                     serviceConnection.connectionProxy.tunnelRealStateFlow(),
+                    serviceConnection.connectionProxy.lastKnownDisconnectedLocation(),
                     accountRepository.accountExpiryState,
-                    _isTunnelInfoExpanded,
                     deviceRepository.deviceState.map { it.deviceName() }
                 ) {
-                    location,
                     relayLocation,
                     notifications,
                     tunnelUiState,
                     tunnelRealState,
+                    lastKnownDisconnectedLocation,
                     accountExpiry,
-                    isTunnelInfoExpanded,
                     deviceName ->
                     ConnectUiState(
                         location =
                             when (tunnelRealState) {
-                                is TunnelState.Connected -> tunnelRealState.location
+                                is TunnelState.Disconnected -> tunnelRealState.location()
+                                        ?: lastKnownDisconnectedLocation
                                 is TunnelState.Connecting -> tunnelRealState.location
-                                else -> null
-                            }
-                                ?: location,
+                                        ?: relayLocation?.location?.location
+                                is TunnelState.Connected -> tunnelRealState.location
+                                is TunnelState.Disconnecting -> lastKnownDisconnectedLocation
+                                is TunnelState.Error -> null
+                            },
                         relayLocation = relayLocation,
                         tunnelUiState = tunnelUiState,
                         tunnelRealState = tunnelRealState,
-                        isTunnelInfoExpanded = isTunnelInfoExpanded,
                         inAddress =
                             when (tunnelRealState) {
                                 is TunnelState.Connected -> tunnelRealState.endpoint.toInAddress()
                                 is TunnelState.Connecting -> tunnelRealState.endpoint?.toInAddress()
                                 else -> null
                             },
-                        outAddress = location?.toOutAddress() ?: "",
+                        outAddress = tunnelRealState.location()?.toOutAddress() ?: "",
                         showLocation =
                             when (tunnelUiState) {
                                 is TunnelState.Disconnected -> true
@@ -149,20 +147,18 @@ class ConnectViewModel(
         }
     }
 
-    private fun LocationInfoCache.locationCallbackFlow() = callbackFlow {
-        onNewLocation = { this.trySend(it) }
-        awaitClose { onNewLocation = null }
-    }
-
     private fun ConnectionProxy.tunnelUiStateFlow(): Flow<TunnelState> =
         callbackFlowFromNotifier(this.onUiStateChange)
 
     private fun ConnectionProxy.tunnelRealStateFlow(): Flow<TunnelState> =
         callbackFlowFromNotifier(this.onStateChange)
 
-    fun toggleTunnelInfoExpansion() {
-        _isTunnelInfoExpanded.value = _isTunnelInfoExpanded.value.not()
-    }
+    private fun ConnectionProxy.lastKnownDisconnectedLocation(): Flow<GeoIpLocation?> =
+        tunnelRealStateFlow()
+            .filterIsInstance<TunnelState.Disconnected>()
+            .filter { it.location != null }
+            .map { it.location }
+            .onStart { emit(null) }
 
     fun onDisconnectClick() {
         serviceConnectionManager.connectionProxy()?.disconnect()
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/OutOfTimeUseCaseTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/OutOfTimeUseCaseTest.kt
index 74683813ae..bdc5ea49b8 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/OutOfTimeUseCaseTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/OutOfTimeUseCaseTest.kt
@@ -63,7 +63,7 @@ class OutOfTimeUseCaseTest {
         val expiredAccountExpiry = AccountExpiry.Available(DateTime.now().plusDays(1))
         val tunnelStateChanges =
             listOf(
-                    TunnelState.Disconnected,
+                    TunnelState.Disconnected(),
                     TunnelState.Connected(mockk(), null),
                     TunnelState.Connecting(null, null),
                     TunnelState.Disconnecting(mockk()),
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/TunnelStateNotificationUseCaseTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/TunnelStateNotificationUseCaseTest.kt
index 1b89c92be7..e67630642d 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/TunnelStateNotificationUseCaseTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/TunnelStateNotificationUseCaseTest.kt
@@ -35,7 +35,7 @@ class TunnelStateNotificationUseCaseTest {
         MutableStateFlow<ServiceConnectionState>(ServiceConnectionState.Disconnected)
     private lateinit var tunnelStateNotificationUseCase: TunnelStateNotificationUseCase
 
-    private val eventNotifierTunnelUiState = EventNotifier<TunnelState>(TunnelState.Disconnected)
+    private val eventNotifierTunnelUiState = EventNotifier<TunnelState>(TunnelState.Disconnected())
 
     @Before
     fun setup() {
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/ConnectViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/ConnectViewModelTest.kt
index 7271c07433..47364652ce 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/ConnectViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/ConnectViewModelTest.kt
@@ -6,12 +6,11 @@ import io.mockk.coEvery
 import io.mockk.every
 import io.mockk.mockk
 import io.mockk.mockkStatic
-import io.mockk.slot
 import io.mockk.unmockkAll
 import io.mockk.verify
 import kotlin.test.assertEquals
 import kotlin.test.assertIs
-import kotlin.test.assertTrue
+import kotlin.test.assertNull
 import kotlinx.coroutines.async
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -33,7 +32,6 @@ import net.mullvad.mullvadvpn.ui.VersionInfo
 import net.mullvad.mullvadvpn.ui.serviceconnection.AppVersionInfoCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.AuthTokenCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.ConnectionProxy
-import net.mullvad.mullvadvpn.ui.serviceconnection.LocationInfoCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionContainer
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManager
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionState
@@ -73,7 +71,6 @@ class ConnectViewModelTest {
 
     // Service connections
     private val mockServiceConnectionContainer: ServiceConnectionContainer = mockk()
-    private val mockLocationInfoCache: LocationInfoCache = mockk(relaxUnitFun = true)
     private lateinit var mockAppVersionInfoCache: AppVersionInfoCache
     private val mockConnectionProxy: ConnectionProxy = mockk()
     private val mockLocation: GeoIpLocation = mockk(relaxed = true)
@@ -93,12 +90,10 @@ class ConnectViewModelTest {
     // Payment use case
     private val mockPaymentUseCase: PaymentUseCase = mockk(relaxed = true)
 
-    // Captures
-    private val locationSlot = slot<((GeoIpLocation?) -> Unit)>()
-
     // Event notifiers
-    private val eventNotifierTunnelUiState = EventNotifier<TunnelState>(TunnelState.Disconnected)
-    private val eventNotifierTunnelRealState = EventNotifier<TunnelState>(TunnelState.Disconnected)
+    private val eventNotifierTunnelUiState = EventNotifier<TunnelState>(TunnelState.Disconnected())
+    private val eventNotifierTunnelRealState =
+        EventNotifier<TunnelState>(TunnelState.Disconnected())
 
     // Flows
     private val selectedRelayFlow = MutableStateFlow<RelayItem?>(null)
@@ -118,7 +113,6 @@ class ConnectViewModelTest {
             }
 
         every { mockServiceConnectionManager.connectionState } returns serviceConnectionState
-        every { mockServiceConnectionContainer.locationInfoCache } returns mockLocationInfoCache
         every { mockServiceConnectionContainer.appVersionInfoCache } returns mockAppVersionInfoCache
         every { mockServiceConnectionContainer.connectionProxy } returns mockConnectionProxy
 
@@ -134,7 +128,6 @@ class ConnectViewModelTest {
         every { mockLocation.country } returns "dummy country"
 
         // Listeners
-        every { mockLocationInfoCache.onNewLocation = capture(locationSlot) } answers {}
         every { mockAppVersionInfoCache.onUpdate = any() } answers {}
 
         // Flows
@@ -167,29 +160,14 @@ class ConnectViewModelTest {
     }
 
     @Test
-    fun testTunnelInfoExpandedUpdate() =
-        runTest(testCoroutineRule.testDispatcher) {
-            viewModel.uiState.test {
-                assertEquals(ConnectUiState.INITIAL, awaitItem())
-                serviceConnectionState.value =
-                    ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
-                locationSlot.captured.invoke(mockLocation)
-                viewModel.toggleTunnelInfoExpansion()
-                val result = awaitItem()
-                assertTrue(result.isTunnelInfoExpanded)
-            }
-        }
-
-    @Test
     fun testTunnelRealStateUpdate() =
         runTest(testCoroutineRule.testDispatcher) {
-            val tunnelRealStateTestItem = TunnelState.Connected(mockk(relaxed = true), mockk())
+            val tunnelRealStateTestItem = TunnelState.Connected(mockk(relaxed = true), null)
 
             viewModel.uiState.test {
                 assertEquals(ConnectUiState.INITIAL, awaitItem())
                 serviceConnectionState.value =
                     ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
-                locationSlot.captured.invoke(mockLocation)
                 eventNotifierTunnelRealState.notify(tunnelRealStateTestItem)
                 val result = awaitItem()
                 assertEquals(tunnelRealStateTestItem, result.tunnelRealState)
@@ -205,7 +183,6 @@ class ConnectViewModelTest {
                 assertEquals(ConnectUiState.INITIAL, awaitItem())
                 serviceConnectionState.value =
                     ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
-                locationSlot.captured.invoke(mockLocation)
                 eventNotifierTunnelUiState.notify(tunnelUiStateTestItem)
                 val result = awaitItem()
                 assertEquals(tunnelUiStateTestItem, result.tunnelUiState)
@@ -223,7 +200,6 @@ class ConnectViewModelTest {
                 assertEquals(ConnectUiState.INITIAL, awaitItem())
                 serviceConnectionState.value =
                     ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
-                locationSlot.captured.invoke(mockLocation)
                 val result = awaitItem()
                 assertEquals(relayTestItem, result.relayLocation)
             }
@@ -241,13 +217,19 @@ class ConnectViewModelTest {
                     hostname = "Host"
                 )
 
+            // Act, Assert
             viewModel.uiState.test {
                 assertEquals(ConnectUiState.INITIAL, awaitItem())
+                eventNotifierTunnelRealState.notify(TunnelState.Disconnected(null))
+
                 serviceConnectionState.value =
                     ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
-                locationSlot.captured.invoke(locationTestItem)
-                val result = awaitItem()
-                assertEquals(locationTestItem, result.location)
+                // Start of with no location
+                assertNull(awaitItem().location)
+
+                // After updated we show latest
+                eventNotifierTunnelRealState.notify(TunnelState.Disconnected(locationTestItem))
+                assertEquals(locationTestItem, awaitItem().location)
             }
         }
 
@@ -262,7 +244,6 @@ class ConnectViewModelTest {
                 assertEquals(ConnectUiState.INITIAL, awaitItem())
                 serviceConnectionState.value =
                     ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
-                locationSlot.captured.invoke(locationTestItem)
                 expectNoEvents()
                 val result = awaitItem()
                 assertEquals(locationTestItem, result.location)
@@ -320,7 +301,6 @@ class ConnectViewModelTest {
                 assertEquals(ConnectUiState.INITIAL, awaitItem())
                 serviceConnectionState.value =
                     ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
-                locationSlot.captured.invoke(mockLocation)
                 eventNotifierTunnelUiState.notify(tunnelUiState)
                 val result = awaitItem()
                 assertEquals(expectedConnectNotificationState, result.inAppNotification)
@@ -356,7 +336,6 @@ class ConnectViewModelTest {
                 awaitItem()
                 serviceConnectionState.value =
                     ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
-                locationSlot.captured.invoke(mockLocation)
                 outOfTimeViewFlow.value = true
                 awaitItem()
             }
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceRevokedViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceRevokedViewModelTest.kt
index 8959b1f9e3..a73ecfc4e7 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceRevokedViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceRevokedViewModelTest.kt
@@ -105,7 +105,7 @@ class DeviceRevokedViewModelTest {
         // Arrange
         val mockedContainer =
             mockk<ServiceConnectionContainer>().also {
-                every { it.connectionProxy.state } returns TunnelState.Disconnected
+                every { it.connectionProxy.state } returns TunnelState.Disconnected()
                 every { it.connectionProxy.disconnect() } just Runs
                 every { mockedAccountRepository.logout() } just Runs
             }
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/OutOfTimeViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/OutOfTimeViewModelTest.kt
index ab861b1e20..50a9d8fb98 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/OutOfTimeViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/OutOfTimeViewModelTest.kt
@@ -59,7 +59,8 @@ class OutOfTimeViewModelTest {
     private val mockConnectionProxy: ConnectionProxy = mockk()
 
     // Event notifiers
-    private val eventNotifierTunnelRealState = EventNotifier<TunnelState>(TunnelState.Disconnected)
+    private val eventNotifierTunnelRealState =
+        EventNotifier<TunnelState>(TunnelState.Disconnected())
 
     private val mockAccountRepository: AccountRepository = mockk(relaxed = true)
     private val mockDeviceRepository: DeviceRepository = mockk()
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/WelcomeViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/WelcomeViewModelTest.kt
index ac047db2ea..29b0dfde75 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/WelcomeViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/WelcomeViewModelTest.kt
@@ -58,7 +58,7 @@ class WelcomeViewModelTest {
     private val mockConnectionProxy: ConnectionProxy = mockk()
 
     // Event notifiers
-    private val eventNotifierTunnelUiState = EventNotifier<TunnelState>(TunnelState.Disconnected)
+    private val eventNotifierTunnelUiState = EventNotifier<TunnelState>(TunnelState.Disconnected())
 
     private val mockAccountRepository: AccountRepository = mockk(relaxed = true)
     private val mockDeviceRepository: DeviceRepository = mockk()