Clear OOT notification on login/revoke screens

There are cases when a user may get to the login or device revoked
screens without their previously scheduled account expiry notification
being cancelled.

This can happen if for some reason the dameon has not sent the logged
out/revoked event, for example if the account is logged out via the
website while the app is not running.

This PR cancels the notification when the user gets to the login or
device revoked screens.

10 files changed, 77 insertions, 27 deletions

diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/MullvadApplication.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/MullvadApplication.kt
index 43686b723b..655ed045e6 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/MullvadApplication.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/MullvadApplication.kt
@@ -59,17 +59,11 @@ class MullvadApplication : Application() {
                 when (action) {
                     NotificationAction.CancelExisting -> {
                         accountExpiryNotificationProvider.cancelNotification()
-                        scheduleNotificationAlarmUseCase(
-                            context = this@MullvadApplication,
-                            accountExpiry = null,
-                        )
+                        scheduleNotificationAlarmUseCase(accountExpiry = null)
                     }
 
                     is NotificationAction.ScheduleAlarm ->
-                        scheduleNotificationAlarmUseCase(
-                            context = this@MullvadApplication,
-                            accountExpiry = action.alarmTime,
-                        )
+                        scheduleNotificationAlarmUseCase(accountExpiry = action.alarmTime)
                 }
             }
         }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/AppModule.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/AppModule.kt
index 58b1a6e478..05c5e6d92a 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/AppModule.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/AppModule.kt
@@ -64,7 +64,7 @@ val appModule = module {
     single { ConnectionProxy(get(), get(), get()) }
     single { LocaleRepository(get()) }
     single { RelayLocationTranslationRepository(get(), get(), MainScope()) }
-    single { ScheduleNotificationAlarmUseCase(get()) }
+    single { ScheduleNotificationAlarmUseCase(androidContext(), get()) }
     single { AccountExpiryNotificationActionUseCase(get(), get()) }
 
     single { NotificationChannel.TunnelUpdates } bind NotificationChannel::class
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt
index 7b2ae9986e..ffe1d078ed 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt
@@ -255,11 +255,11 @@ val uiModule = module {
     }
     viewModel { DeviceListViewModel(get(), get()) }
     viewModel { ManageDevicesViewModel(get(), get()) }
-    viewModel { DeviceRevokedViewModel(get(), get()) }
+    viewModel { DeviceRevokedViewModel(get(), get(), get(), get()) }
     viewModel { MtuDialogViewModel(get(), get()) }
     viewModel { DnsDialogViewModel(get(), get(), get(), get()) }
     viewModel { WireguardCustomPortDialogViewModel(get()) }
-    viewModel { LoginViewModel(get(), get(), get()) }
+    viewModel { LoginViewModel(get(), get(), get(), get(), get()) }
     viewModel { PrivacyDisclaimerViewModel(get(), IS_PLAY_BUILD) }
     viewModel {
         SelectLocationViewModel(get(), get(), get(), get(), get(), get(), get(), get(), get())
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/receiver/NotificationAlarmReceiver.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/receiver/NotificationAlarmReceiver.kt
index 49bd290230..a638152530 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/receiver/NotificationAlarmReceiver.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/receiver/NotificationAlarmReceiver.kt
@@ -33,7 +33,7 @@ class NotificationAlarmReceiver : BroadcastReceiver(), KoinComponent {
         goAsync {
             // Only schedule the next alarm if we still have time left on the account.
             if (context != null && expiry > ZonedDateTime.now()) {
-                scheduleNotificationAlarmUseCase(context = context, accountExpiry = expiry)
+                scheduleNotificationAlarmUseCase(accountExpiry = expiry, customContext = context)
             }
         }
     }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/receiver/ScheduleNotificationBootCompletedReceiver.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/receiver/ScheduleNotificationBootCompletedReceiver.kt
index d7ad2660e4..e972a55b93 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/receiver/ScheduleNotificationBootCompletedReceiver.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/receiver/ScheduleNotificationBootCompletedReceiver.kt
@@ -27,6 +27,6 @@ class ScheduleNotificationBootCompletedReceiver : BroadcastReceiver(), KoinCompo
 
     private suspend fun scheduleAccountExpiryNotification(context: Context) {
         val expiry = userPreferencesRepository.accountExpiry() ?: return
-        scheduleNotificationAlarmUseCase(context, expiry)
+        scheduleNotificationAlarmUseCase(expiry, customContext = context)
     }
 }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/ScheduleNotificationAlarmUseCase.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/ScheduleNotificationAlarmUseCase.kt
index 1c2bc63eb6..41045beb36 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/ScheduleNotificationAlarmUseCase.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/ScheduleNotificationAlarmUseCase.kt
@@ -12,13 +12,14 @@ import net.mullvad.mullvadvpn.repository.UserPreferencesRepository
 import net.mullvad.mullvadvpn.service.notifications.accountexpiry.accountExpiryNotificationTriggerAt
 
 class ScheduleNotificationAlarmUseCase(
-    private val userPreferencesRepository: UserPreferencesRepository
+    private val applicationContext: Context,
+    private val userPreferencesRepository: UserPreferencesRepository,
 ) {
-    suspend operator fun invoke(context: Context, accountExpiry: ZonedDateTime?) {
-        val appContext = context.applicationContext
-        val alarmManager = appContext.getSystemService(AlarmManager::class.java) ?: return
+    suspend operator fun invoke(accountExpiry: ZonedDateTime?, customContext: Context? = null) {
+        val context = customContext ?: applicationContext
+        val alarmManager = context.getSystemService(AlarmManager::class.java) ?: return
 
-        cancelExisting(appContext, alarmManager)
+        cancelExisting(context, alarmManager)
 
         if (accountExpiry == null) {
             userPreferencesRepository.clearAccountExpiry()
@@ -29,7 +30,7 @@ class ScheduleNotificationAlarmUseCase(
             accountExpiryNotificationTriggerAt(now = ZonedDateTime.now(), expiry = accountExpiry)
         val triggerAtMillis = triggerAt.toInstant().toEpochMilli()
 
-        val intent = alarmIntent(appContext, accountExpiry)
+        val intent = alarmIntent(context, accountExpiry)
         alarmManager.set(AlarmManager.RTC, triggerAtMillis, intent)
 
         // Change to UTC to avoid leaking the user's time zone in the logs
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceRevokedViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceRevokedViewModel.kt
index 020a384c6c..60a1101f43 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceRevokedViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceRevokedViewModel.kt
@@ -2,13 +2,11 @@ package net.mullvad.mullvadvpn.viewmodel
 
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
-import kotlinx.coroutines.CoroutineDispatcher
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.WhileSubscribed
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.flow.receiveAsFlow
 import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
@@ -16,15 +14,22 @@ import net.mullvad.mullvadvpn.compose.state.DeviceRevokedUiState
 import net.mullvad.mullvadvpn.constant.VIEW_MODEL_STOP_TIMEOUT
 import net.mullvad.mullvadvpn.lib.shared.AccountRepository
 import net.mullvad.mullvadvpn.lib.shared.ConnectionProxy
+import net.mullvad.mullvadvpn.service.notifications.accountexpiry.AccountExpiryNotificationProvider
+import net.mullvad.mullvadvpn.usecase.ScheduleNotificationAlarmUseCase
 
 class DeviceRevokedViewModel(
     private val accountRepository: AccountRepository,
     private val connectionProxy: ConnectionProxy,
-    dispatcher: CoroutineDispatcher = Dispatchers.IO,
+    private val scheduleNotificationAlarmUseCase: ScheduleNotificationAlarmUseCase,
+    private val accountExpiryNotificationProvider: AccountExpiryNotificationProvider,
 ) : ViewModel() {
 
     val uiState =
         connectionProxy.tunnelState
+            .onStart {
+                accountExpiryNotificationProvider.cancelNotification()
+                scheduleNotificationAlarmUseCase(accountExpiry = null)
+            }
             .map {
                 if (it.isSecured()) {
                     DeviceRevokedUiState.SECURED
@@ -33,7 +38,7 @@ class DeviceRevokedViewModel(
                 }
             }
             .stateIn(
-                scope = CoroutineScope(dispatcher),
+                scope = viewModelScope,
                 started = SharingStarted.WhileSubscribed(VIEW_MODEL_STOP_TIMEOUT),
                 initialValue = DeviceRevokedUiState.UNKNOWN,
             )
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModel.kt
index e114ccde7c..6092c30bf1 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModel.kt
@@ -33,7 +33,9 @@ import net.mullvad.mullvadvpn.lib.model.CreateAccountError
 import net.mullvad.mullvadvpn.lib.model.LoginAccountError
 import net.mullvad.mullvadvpn.lib.shared.AccountRepository
 import net.mullvad.mullvadvpn.repository.NewDeviceRepository
+import net.mullvad.mullvadvpn.service.notifications.accountexpiry.AccountExpiryNotificationProvider
 import net.mullvad.mullvadvpn.usecase.InternetAvailableUseCase
+import net.mullvad.mullvadvpn.usecase.ScheduleNotificationAlarmUseCase
 import net.mullvad.mullvadvpn.util.delayAtLeast
 import net.mullvad.mullvadvpn.util.getOrDefault
 import net.mullvad.mullvadvpn.viewmodel.LoginUiSideEffect.NavigateToWelcome
@@ -59,6 +61,8 @@ class LoginViewModel(
     private val accountRepository: AccountRepository,
     private val newDeviceRepository: NewDeviceRepository,
     private val internetAvailableUseCase: InternetAvailableUseCase,
+    private val scheduleNotificationAlarmUseCase: ScheduleNotificationAlarmUseCase,
+    private val accountExpiryNotificationProvider: AccountExpiryNotificationProvider,
     private val dispatcher: CoroutineDispatcher = Dispatchers.IO,
 ) : ViewModel() {
     private val _loginState = MutableStateFlow(LoginUiState.INITIAL.loginState)
@@ -77,7 +81,13 @@ class LoginViewModel(
 
     val uiState: StateFlow<LoginUiState> =
         _uiState
-            .onStart { viewModelScope.launch { accountRepository.fetchAccountHistory() } }
+            .onStart {
+                viewModelScope.launch {
+                    accountRepository.fetchAccountHistory()
+                    accountExpiryNotificationProvider.cancelNotification()
+                    scheduleNotificationAlarmUseCase(accountExpiry = null)
+                }
+            }
             .stateIn(
                 viewModelScope,
                 SharingStarted.WhileSubscribed(VIEW_MODEL_STOP_TIMEOUT),
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 d0af241314..f3be9af276 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
@@ -12,13 +12,14 @@ import io.mockk.mockk
 import io.mockk.unmockkAll
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.MutableSharedFlow
-import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.runTest
 import net.mullvad.mullvadvpn.compose.state.DeviceRevokedUiState
 import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
 import net.mullvad.mullvadvpn.lib.model.TunnelState
 import net.mullvad.mullvadvpn.lib.shared.AccountRepository
 import net.mullvad.mullvadvpn.lib.shared.ConnectionProxy
+import net.mullvad.mullvadvpn.service.notifications.accountexpiry.AccountExpiryNotificationProvider
+import net.mullvad.mullvadvpn.usecase.ScheduleNotificationAlarmUseCase
 import org.junit.jupiter.api.AfterEach
 import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.BeforeEach
@@ -37,6 +38,12 @@ class DeviceRevokedViewModelTest {
 
     private val tunnelStateFlow = MutableSharedFlow<TunnelState>()
 
+    private val mockScheduleNotificationAlarmUseCase =
+        mockk<ScheduleNotificationAlarmUseCase>(relaxed = true)
+
+    private val mockAccountExpiryNotificationProvider =
+        mockk<AccountExpiryNotificationProvider>(relaxed = true)
+
     @BeforeEach
     fun setup() {
         MockKAnnotations.init(this)
@@ -45,7 +52,8 @@ class DeviceRevokedViewModelTest {
             DeviceRevokedViewModel(
                 accountRepository = mockedAccountRepository,
                 connectionProxy = mockConnectionProxy,
-                dispatcher = UnconfinedTestDispatcher(),
+                scheduleNotificationAlarmUseCase = mockScheduleNotificationAlarmUseCase,
+                accountExpiryNotificationProvider = mockAccountExpiryNotificationProvider,
             )
     }
 
@@ -69,6 +77,17 @@ class DeviceRevokedViewModelTest {
     }
 
     @Test
+    fun `when subscription starts the user account expiry notification should be cancelled`() =
+        runTest {
+            // Act, Assert
+            viewModel.uiState.test {
+                assertEquals(DeviceRevokedUiState.UNKNOWN, awaitItem())
+                coVerify { mockScheduleNotificationAlarmUseCase(null, null) }
+                coVerify { mockAccountExpiryNotificationProvider.cancelNotification() }
+            }
+        }
+
+    @Test
     fun `onGoToLoginClicked should invoke logout on AccountRepository`() {
         // Arrange
         coEvery { mockConnectionProxy.disconnect() } returns true.right()
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModelTest.kt
index 81e0afefbe..90ec8f1c8c 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModelTest.kt
@@ -29,7 +29,9 @@ import net.mullvad.mullvadvpn.lib.model.AccountNumber
 import net.mullvad.mullvadvpn.lib.model.CreateAccountError
 import net.mullvad.mullvadvpn.lib.model.LoginAccountError
 import net.mullvad.mullvadvpn.lib.shared.AccountRepository
+import net.mullvad.mullvadvpn.service.notifications.accountexpiry.AccountExpiryNotificationProvider
 import net.mullvad.mullvadvpn.usecase.InternetAvailableUseCase
+import net.mullvad.mullvadvpn.usecase.ScheduleNotificationAlarmUseCase
 import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
@@ -45,6 +47,12 @@ class LoginViewModelTest {
     private lateinit var loginViewModel: LoginViewModel
     private lateinit var accountHistoryFlow: MutableStateFlow<AccountNumber?>
 
+    private val mockScheduleNotificationAlarmUseCase =
+        mockk<ScheduleNotificationAlarmUseCase>(relaxed = true)
+
+    private val mockAccountExpiryNotificationProvider =
+        mockk<AccountExpiryNotificationProvider>(relaxed = true)
+
     @BeforeEach
     fun setup() {
         MockKAnnotations.init(this, relaxUnitFun = true)
@@ -58,6 +66,8 @@ class LoginViewModelTest {
                 accountRepository = mockedAccountRepository,
                 newDeviceRepository = mockk(relaxUnitFun = true),
                 internetAvailableUseCase = connectivityUseCase,
+                scheduleNotificationAlarmUseCase = mockScheduleNotificationAlarmUseCase,
+                accountExpiryNotificationProvider = mockAccountExpiryNotificationProvider,
                 UnconfinedTestDispatcher(),
             )
     }
@@ -94,6 +104,17 @@ class LoginViewModelTest {
     }
 
     @Test
+    fun `when subscription starts the user account expiry notification should be cancelled`() =
+        runTest {
+            // Act, Assert
+            loginViewModel.uiState.test {
+                assertEquals(LoginUiState.INITIAL, awaitItem())
+                coVerify { mockScheduleNotificationAlarmUseCase(null, null) }
+                coVerify { mockAccountExpiryNotificationProvider.cancelNotification() }
+            }
+        }
+
+    @Test
     fun `createAccount call should result in NavigateToWelcome side effect`() = runTest {
         turbineScope {
             // Arrange