From 08e8569e09db4a7c479bb2b3048872084211cde1 Mon Sep 17 00:00:00 2001
From: Jonatan Rhodin <jonatan.rhodin@mullvad.net>
Date: Tue, 19 Nov 2024 16:58:17 +0100
Subject: [PATCH] Use last known location instead of null when tunnel state is
 error

---
 .../mullvadvpn/viewmodel/ConnectViewModel.kt  |  4 +--
 .../viewmodel/ConnectViewModelTest.kt         | 31 ++++++++++++++-----
 2 files changed, 26 insertions(+), 9 deletions(-)

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 d9ca922f1f27..5572f9396170 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
@@ -78,11 +78,11 @@ class ConnectViewModel(
                     location =
                         when (tunnelState) {
                             is TunnelState.Disconnected ->
-                                tunnelState.location() ?: lastKnownDisconnectedLocation
+                                tunnelState.location ?: lastKnownDisconnectedLocation
                             is TunnelState.Connecting -> tunnelState.location
                             is TunnelState.Connected -> tunnelState.location
                             is TunnelState.Disconnecting -> lastKnownDisconnectedLocation
-                            is TunnelState.Error -> null
+                            is TunnelState.Error -> lastKnownDisconnectedLocation
                         },
                     selectedRelayItemTitle = selectedRelayItemTitle,
                     tunnelState = tunnelState,
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 33e836acd863..3dada2a4332b 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
@@ -17,7 +17,6 @@ import kotlinx.coroutines.async
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.first
-import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.test.runTest
 import net.mullvad.mullvadvpn.compose.state.ConnectUiState
 import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
@@ -79,6 +78,7 @@ class ConnectViewModelTest {
     // Flows
     private val tunnelState = MutableStateFlow<TunnelState>(TunnelState.Disconnected())
     private val selectedRelayItemFlow = MutableStateFlow<String?>(null)
+    private val lastKnownLocationFlow = MutableStateFlow<GeoIpLocation?>(null)
 
     // Out Of Time Use Case
     private val outOfTimeUseCase: OutOfTimeUseCase = mockk()
@@ -104,7 +104,8 @@ class ConnectViewModelTest {
 
         every { mockConnectionProxy.tunnelState } returns tunnelState
 
-        every { mockLastKnownLocationUseCase.lastKnownDisconnectedLocation } returns flowOf(null)
+        every { mockLastKnownLocationUseCase.lastKnownDisconnectedLocation } returns
+            lastKnownLocationFlow
 
         every { mockLocation.country } returns "dummy country"
 
@@ -142,15 +143,14 @@ class ConnectViewModelTest {
     }
 
     @Test
-    fun `given change in tunnelRealState uiState should emit new tunnelRealState`() = runTest {
-        val tunnelRealStateTestItem =
-            TunnelState.Connected(mockk(relaxed = true), null, emptyList())
+    fun `given change in tunnel state uiState should emit new tunnel state`() = runTest {
+        val tunnelStateTestItem = TunnelState.Connected(mockk(relaxed = true), null, emptyList())
 
         viewModel.uiState.test {
             assertEquals(ConnectUiState.INITIAL, awaitItem())
-            tunnelState.emit(tunnelRealStateTestItem)
+            tunnelState.emit(tunnelStateTestItem)
             val result = awaitItem()
-            assertEquals(tunnelRealStateTestItem, result.tunnelState)
+            assertEquals(tunnelStateTestItem, result.tunnelState)
         }
     }
 
@@ -326,4 +326,21 @@ class ConnectViewModelTest {
         // Assert
         assertIs<ConnectViewModel.UiSideEffect.OutOfTime>(deferred.await())
     }
+
+    @Test
+    fun `given tunnel state error should emit last known disconnected location as location`() =
+        runTest {
+            // Arrange
+            val tunnel = TunnelState.Error(mockk(relaxed = true))
+            val lastKnownLocation: GeoIpLocation = mockk(relaxed = true)
+            lastKnownLocationFlow.emit(lastKnownLocation)
+            tunnelState.emit(tunnel)
+
+            // Act, Assert
+            viewModel.uiState.test {
+                assertEquals(ConnectUiState.INITIAL, awaitItem())
+                val result = awaitItem()
+                assertEquals(lastKnownLocation, result.location)
+            }
+        }
 }