Closes #16053: Fix crash when closing private tabs via notification

app/src/main/java/org/mozilla/fenix/HomeActivity.kt

@@ -180,7 +180,6 @@ open class HomeActivity : LocaleAwareAppCompatActivity(), NavHostActivity {
                 .attachViewToRunVisualCompletenessQueueLater(WeakReference(rootContainer))
         }
 
-        checkPrivateShortcutEntryPoint(intent)
         privateNotificationObserver = PrivateNotificationFeature(
             applicationContext,
             components.core.store,
@@ -603,17 +602,6 @@ open class HomeActivity : LocaleAwareAppCompatActivity(), NavHostActivity {
         return false
     }
 
-    private fun checkPrivateShortcutEntryPoint(intent: Intent) {
-        if (intent.hasExtra(OPEN_TO_SEARCH) &&
-            (intent.getStringExtra(OPEN_TO_SEARCH) ==
-                    StartSearchIntentProcessor.STATIC_SHORTCUT_NEW_PRIVATE_TAB ||
-                    intent.getStringExtra(OPEN_TO_SEARCH) ==
-                    StartSearchIntentProcessor.PRIVATE_BROWSING_PINNED_SHORTCUT)
-        ) {
-            PrivateNotificationService.isStartedFromPrivateShortcut = true
-        }
-    }
-
     private fun setupThemeAndBrowsingMode(mode: BrowsingMode) {
         settings().lastKnownMode = mode
         browsingModeManager = createBrowsingModeManager(mode)

app/src/main/java/org/mozilla/fenix/components/toolbar/ToolbarIntegration.kt

@@ -141,7 +141,6 @@ class DefaultToolbarIntegration(
 
         val tabsAction = TabCounterToolbarButton(
             lifecycleOwner,
-            isPrivate,
             onItemTapped = {
                 interactor.onTabCounterMenuItemTapped(it)
             },

app/src/main/java/org/mozilla/fenix/session/PrivateNotificationService.kt

@@ -4,9 +4,11 @@
 
 package org.mozilla.fenix.session
 
+import android.annotation.SuppressLint
 import android.content.Intent
 import androidx.core.app.NotificationCompat
 import androidx.core.content.ContextCompat
+import mozilla.components.browser.state.selector.selectedTab
 import mozilla.components.browser.state.store.BrowserStore
 import mozilla.components.feature.privatemode.notification.AbstractPrivateNotificationService
 import org.mozilla.fenix.HomeActivity
@@ -36,31 +38,35 @@ class PrivateNotificationService : AbstractPrivateNotificationService() {
         color = ContextCompat.getColor(this@PrivateNotificationService, R.color.pbm_notification_color)
     }
 
+    @SuppressLint("MissingSuperCall")
     override fun erasePrivateTabs() {
-        metrics.track(Event.PrivateBrowsingNotificationTapped)
+        val inPrivateMode = store.state.selectedTab?.content?.private ?: false
 
-        val homeScreenIntent = Intent(this, HomeActivity::class.java).apply {
-            flags = Intent.FLAG_ACTIVITY_NEW_TASK or Intent.FLAG_ACTIVITY_CLEAR_TASK
-            putExtra(HomeActivity.PRIVATE_BROWSING_MODE, isStartedFromPrivateShortcut)
-        }
+        // Trigger use case directly for now (instead of calling super.erasePrivateTabs)
+        // as otherwise SessionManager and the store will be out of sync.
+        components.useCases.tabsUseCases.removePrivateTabs()
 
-        if (VisibilityLifecycleCallback.finishAndRemoveTaskIfInBackground(this)) {
-            // Set start mode to be in background (recents screen)
-            homeScreenIntent.apply {
-                putExtra(HomeActivity.START_IN_RECENTS_SCREEN, true)
-            }
-        }
+        metrics.track(Event.PrivateBrowsingNotificationTapped)
 
-        startActivity(homeScreenIntent)
-        super.erasePrivateTabs()
-    }
+        // If the app is in private mode we launch to the private mode home screen as a
+        // confirmation that all private tabs have been deleted. If we don't do this the user
+        // will end up on a new selected tab in normal mode which isn't desired.
+        // If the app is in normal mode there's no reason to direct the user away to
+        // private mode as all private tabs have been deleted.
+        if (inPrivateMode) {
+            val homeScreenIntent = Intent(this, HomeActivity::class.java).apply {
+                flags = Intent.FLAG_ACTIVITY_NEW_TASK or Intent.FLAG_ACTIVITY_CLEAR_TASK
+                putExtra(HomeActivity.PRIVATE_BROWSING_MODE, true)
+            }
 
-    companion object {
+            if (VisibilityLifecycleCallback.finishAndRemoveTaskIfInBackground(this)) {
+                // Set start mode to be in background (recents screen)
+                homeScreenIntent.apply {
+                    putExtra(HomeActivity.START_IN_RECENTS_SCREEN, true)
+                }
+            }
 
-        /**
-         * Global used by [HomeActivity] to figure out if normal mode or private mode
-         * should be used after closing all private tabs.
-         */
-        var isStartedFromPrivateShortcut = false
+            startActivity(homeScreenIntent)
+        }
     }
 }

app/src/test/java/org/mozilla/fenix/session/PrivateNotificationServiceTest.kt

@@ -8,10 +8,20 @@ import android.content.ComponentName
 import android.content.Intent
 import io.mockk.every
 import io.mockk.mockk
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.TestCoroutineDispatcher
+import mozilla.components.browser.state.state.BrowserState
+import mozilla.components.browser.state.state.createTab
+import mozilla.components.browser.state.store.BrowserStore
 import mozilla.components.feature.privatemode.notification.AbstractPrivateNotificationService.Companion.ACTION_ERASE
 import mozilla.components.support.test.robolectric.testContext
+import mozilla.components.support.test.rule.MainCoroutineRule
+import org.junit.After
 import org.junit.Assert.assertEquals
+import org.junit.Assert.assertNotNull
+import org.junit.Assert.assertNull
 import org.junit.Before
+import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mozilla.fenix.HomeActivity
@@ -22,14 +32,21 @@ import org.robolectric.Robolectric
 import org.robolectric.Shadows.shadowOf
 import org.robolectric.android.controller.ServiceController
 
+@ExperimentalCoroutinesApi
 @RunWith(FenixRobolectricTestRunner::class)
 class PrivateNotificationServiceTest {
 
     private lateinit var controller: ServiceController<PrivateNotificationService>
+    private lateinit var store: BrowserStore
+
+    private val testDispatcher = TestCoroutineDispatcher()
+
+    @get:Rule
+    val coroutinesTestRule = MainCoroutineRule(testDispatcher)
 
     @Before
     fun setup() {
-        val store = testContext.components.core.store
+        store = testContext.components.core.store
         every { store.dispatch(any()) } returns mockk()
 
         controller = Robolectric.buildService(
@@ -38,27 +55,34 @@ class PrivateNotificationServiceTest {
         )
     }
 
+    @After
+    fun cleanUp() {
+        testDispatcher.cleanupTestCoroutines()
+    }
+
     @Test
-    fun `service opens home activity with PBM flag set to true`() {
-        PrivateNotificationService.isStartedFromPrivateShortcut = true
+    fun `service opens home activity in private mode if app is in private mode`() {
+        val selectedPrivateTab = createTab("https://mozilla.org", private = true)
+        every { store.state } returns BrowserState(tabs = listOf(selectedPrivateTab), selectedTabId = selectedPrivateTab.id)
+
         val service = shadowOf(controller.get())
         controller.startCommand(0, 0)
 
         val intent = service.nextStartedActivity
+        assertNotNull(intent)
         assertEquals(ComponentName(testContext, HomeActivity::class.java), intent.component)
         assertEquals(Intent.FLAG_ACTIVITY_NEW_TASK or Intent.FLAG_ACTIVITY_CLEAR_TASK, intent.flags)
         assertEquals(true, intent.extras?.getBoolean(PRIVATE_BROWSING_MODE))
     }
 
     @Test
-    fun `service opens home activity with PBM flag set to false`() {
-        PrivateNotificationService.isStartedFromPrivateShortcut = false
+    fun `service starts no activity if app is in normal mode`() {
+        val selectedPrivateTab = createTab("https://mozilla.org", private = false)
+        every { store.state } returns BrowserState(tabs = listOf(selectedPrivateTab), selectedTabId = selectedPrivateTab.id)
+
         val service = shadowOf(controller.get())
         controller.startCommand(0, 0)
 
-        val intent = service.nextStartedActivity
-        assertEquals(ComponentName(testContext, HomeActivity::class.java), intent.component)
-        assertEquals(Intent.FLAG_ACTIVITY_NEW_TASK or Intent.FLAG_ACTIVITY_CLEAR_TASK, intent.flags)
-        assertEquals(false, intent.extras?.getBoolean(PRIVATE_BROWSING_MODE))
+        assertNull(service.nextStartedActivity)
     }
 }