Switch to new feature-tab-collections API.

app/src/main/java/org/mozilla/fenix/browser/BrowserFragment.kt

@@ -19,6 +19,7 @@ import kotlinx.android.synthetic.main.fragment_browser.view.*
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import mozilla.components.browser.session.Session
 import mozilla.components.browser.state.selector.findTab
+import mozilla.components.browser.state.state.TabSessionState
 import mozilla.components.browser.thumbnails.BrowserThumbnails
 import mozilla.components.browser.toolbar.BrowserToolbar
 import mozilla.components.feature.app.links.AppLinksUseCases
@@ -257,11 +258,11 @@ class BrowserFragment : BaseBrowserFragment(), UserInteractionHandler {
     }
 
     private val collectionStorageObserver = object : TabCollectionStorage.Observer {
-        override fun onCollectionCreated(title: String, sessions: List<Session>, id: Long?) {
+        override fun onCollectionCreated(title: String, sessions: List<TabSessionState>, id: Long?) {
             showTabSavedToCollectionSnackbar(sessions.size, true)
         }
 
-        override fun onTabsAdded(tabCollection: TabCollection, sessions: List<Session>) {
+        override fun onTabsAdded(tabCollection: TabCollection, sessions: List<TabSessionState>) {
             showTabSavedToCollectionSnackbar(sessions.size)
         }
 

app/src/main/java/org/mozilla/fenix/collections/CollectionCreationController.kt

@@ -9,14 +9,15 @@ package org.mozilla.fenix.collections
 import androidx.annotation.VisibleForTesting
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.launch
-import mozilla.components.browser.session.Session
+import mozilla.components.browser.state.selector.findTab
-import mozilla.components.browser.session.SessionManager
+import mozilla.components.browser.state.selector.normalTabs
+import mozilla.components.browser.state.state.TabSessionState
+import mozilla.components.browser.state.store.BrowserStore
 import mozilla.components.feature.tab.collections.TabCollection
 import org.mozilla.fenix.components.TabCollectionStorage
 import org.mozilla.fenix.components.metrics.Event
 import org.mozilla.fenix.components.metrics.MetricController
 import org.mozilla.fenix.ext.getDefaultCollectionNumber
-import org.mozilla.fenix.ext.normalSessionSize
 import org.mozilla.fenix.home.Tab
 
 interface CollectionCreationController {
@@ -59,24 +60,24 @@ interface CollectionCreationController {
     fun removeTabFromSelection(tab: Tab)
 }
 
-fun List<Tab>.toSessionBundle(sessionManager: SessionManager): List<Session> {
+fun List<Tab>.toTabSessionStateList(store: BrowserStore): List<TabSessionState> {
-    return this.mapNotNull { sessionManager.findSessionById(it.sessionId) }
+    return this.mapNotNull { store.state.findTab(it.sessionId) }
 }
 
 /**
  * @param store Store used to hold in-memory collection state.
+ * @param browserStore The global `BrowserStore` instance.
  * @param dismiss Callback to dismiss the collection creation dialog.
  * @param metrics Controller that handles telemetry events.
  * @param tabCollectionStorage Storage used to save tab collections to disk.
- * @param sessionManager Used to query and serialize tabs.
  * @param scope Coroutine scope to launch coroutines.
  */
 class DefaultCollectionCreationController(
     private val store: CollectionCreationStore,
+    private val browserStore: BrowserStore,
     private val dismiss: () -> Unit,
     private val metrics: MetricController,
     private val tabCollectionStorage: TabCollectionStorage,
-    private val sessionManager: SessionManager,
     private val scope: CoroutineScope
 ) : CollectionCreationController {
 
@@ -88,13 +89,13 @@ class DefaultCollectionCreationController(
     override fun saveCollectionName(tabs: List<Tab>, name: String) {
         dismiss()
 
-        val sessionBundle = tabs.toSessionBundle(sessionManager)
+        val sessionBundle = tabs.toTabSessionStateList(browserStore)
         scope.launch {
             tabCollectionStorage.createCollection(name, sessionBundle)
         }
 
         metrics.track(
-            Event.CollectionSaved(sessionManager.normalSessionSize(), sessionBundle.size)
+            Event.CollectionSaved(browserStore.state.normalTabs.size, sessionBundle.size)
         )
     }
 
@@ -129,14 +130,14 @@ class DefaultCollectionCreationController(
 
     override fun selectCollection(collection: TabCollection, tabs: List<Tab>) {
         dismiss()
-        val sessionBundle = tabs.toList().toSessionBundle(sessionManager)
+        val sessionBundle = tabs.toList().toTabSessionStateList(browserStore)
         scope.launch {
             tabCollectionStorage
                 .addTabsToCollection(collection, sessionBundle)
         }
 
         metrics.track(
-            Event.CollectionTabsAdded(sessionManager.normalSessionSize(), sessionBundle.size)
+            Event.CollectionTabsAdded(browserStore.state.normalTabs.size, sessionBundle.size)
         )
     }
 

app/src/main/java/org/mozilla/fenix/collections/CollectionCreationFragment.kt

@@ -55,10 +55,10 @@ class CollectionCreationFragment : DialogFragment() {
         collectionCreationInteractor = DefaultCollectionCreationInteractor(
             DefaultCollectionCreationController(
                 collectionCreationStore,
+                requireComponents.core.store,
                 ::dismiss,
                 requireComponents.analytics.metrics,
                 requireComponents.core.tabCollectionStorage,
-                requireComponents.core.sessionManager,
                 scope = lifecycleScope
             )
         )

app/src/main/java/org/mozilla/fenix/components/Core.kt

@@ -346,7 +346,6 @@ class Core(
     val tabCollectionStorage by lazyMonitored {
         TabCollectionStorage(
             context,
-            sessionManager,
             strictMode
         )
     }

app/src/main/java/org/mozilla/fenix/components/TabCollectionStorage.kt

@@ -8,12 +8,11 @@ import android.content.Context
 import android.os.StrictMode
 import androidx.lifecycle.LiveData
 import androidx.lifecycle.asLiveData
-import androidx.paging.DataSource
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
-import mozilla.components.browser.session.Session
+import mozilla.components.browser.session.storage.BrowserStateSerializer
-import mozilla.components.browser.session.SessionManager
+import mozilla.components.browser.state.state.TabSessionState
 import mozilla.components.feature.tab.collections.Tab
 import mozilla.components.feature.tab.collections.TabCollection
 import mozilla.components.feature.tab.collections.TabCollectionStorage
@@ -28,7 +27,6 @@ import org.mozilla.fenix.utils.Mockable
 @Mockable
 class TabCollectionStorage(
     private val context: Context,
-    private val sessionManager: SessionManager,
     strictMode: StrictModeManager,
     private val delegate: Observable<Observer> = ObserverRegistry()
 ) : Observable<org.mozilla.fenix.components.TabCollectionStorage.Observer> by delegate {
@@ -40,12 +38,12 @@ class TabCollectionStorage(
         /**
          * A collection has been created
          */
-        fun onCollectionCreated(title: String, sessions: List<Session>, id: Long?) = Unit
+        fun onCollectionCreated(title: String, sessions: List<TabSessionState>, id: Long?) = Unit
 
         /**
          *  Tab(s) have been added to collection
          */
-        fun onTabsAdded(tabCollection: TabCollection, sessions: List<Session>) = Unit
+        fun onTabsAdded(tabCollection: TabCollection, sessions: List<TabSessionState>) = Unit
 
         /**
          *  Collection has been renamed
@@ -58,32 +56,24 @@ class TabCollectionStorage(
 
     private val collectionStorage by lazy {
         strictMode.resetAfter(StrictMode.allowThreadDiskReads()) {
-            TabCollectionStorage(context, sessionManager)
+            TabCollectionStorage(context, BrowserStateSerializer())
         }
     }
 
-    suspend fun createCollection(title: String, sessions: List<Session>) = ioScope.launch {
+    suspend fun createCollection(title: String, sessions: List<TabSessionState>) = ioScope.launch {
         val id = collectionStorage.createCollection(title, sessions)
         notifyObservers { onCollectionCreated(title, sessions, id) }
     }.join()
 
-    suspend fun addTabsToCollection(tabCollection: TabCollection, sessions: List<Session>) = ioScope.launch {
+    suspend fun addTabsToCollection(tabCollection: TabCollection, sessions: List<TabSessionState>) = ioScope.launch {
         collectionStorage.addTabsToCollection(tabCollection, sessions)
         notifyObservers { onTabsAdded(tabCollection, sessions) }
     }.join()
 
-    fun getTabCollectionsCount(): Int {
-        return collectionStorage.getTabCollectionsCount()
-    }
-
     fun getCollections(): LiveData<List<TabCollection>> {
         return collectionStorage.getCollections().asLiveData()
     }
 
-    fun getCollectionsPaged(): DataSource.Factory<Int, TabCollection> {
-        return collectionStorage.getCollectionsPaged()
-    }
-
     suspend fun removeCollection(tabCollection: TabCollection) = ioScope.launch {
         collectionStorage.removeCollection(tabCollection)
     }.join()

app/src/main/java/org/mozilla/fenix/home/HomeFragment.kt

@@ -239,6 +239,7 @@ class HomeFragment : Fragment() {
                 sessionManager = sessionManager,
                 tabCollectionStorage = components.core.tabCollectionStorage,
                 addTabUseCase = components.useCases.tabsUseCases.addTab,
+                restoreUseCase = components.useCases.tabsUseCases.restore,
                 reloadUrlUseCase = components.useCases.sessionUseCases.reload,
                 fragmentStore = homeFragmentStore,
                 navController = findNavController(),

app/src/main/java/org/mozilla/fenix/home/sessioncontrol/SessionControlController.kt

@@ -18,7 +18,7 @@ import mozilla.components.concept.engine.Engine
 import mozilla.components.concept.engine.prompt.ShareData
 import mozilla.components.feature.session.SessionUseCases
 import mozilla.components.feature.tab.collections.TabCollection
-import mozilla.components.feature.tab.collections.ext.restore
+import mozilla.components.feature.tab.collections.ext.invoke
 import mozilla.components.feature.tabs.TabsUseCases
 import mozilla.components.feature.top.sites.TopSite
 import mozilla.components.support.ktx.android.view.showKeyboard
@@ -177,6 +177,7 @@ class DefaultSessionControlController(
     private val store: BrowserStore,
     private val tabCollectionStorage: TabCollectionStorage,
     private val addTabUseCase: TabsUseCases.AddNewTabUseCase,
+    private val restoreUseCase: TabsUseCases.RestoreUseCase,
     private val reloadUrlUseCase: SessionUseCases.ReloadUrlUseCase,
     private val fragmentStore: HomeFragmentStore,
     private val navController: NavController,
@@ -208,7 +209,8 @@ class DefaultSessionControlController(
 
     override fun handleCollectionOpenTabClicked(tab: ComponentTab) {
         dismissSearchDialogIfDisplayed()
-        sessionManager.restore(
+
+        restoreUseCase.invoke(
             activity,
             engine,
             tab,
@@ -229,7 +231,7 @@ class DefaultSessionControlController(
     }
 
     override fun handleCollectionOpenTabsTapped(collection: TabCollection) {
-        sessionManager.restore(
+        restoreUseCase.invoke(
             activity,
             engine,
             collection,

app/src/main/java/org/mozilla/fenix/tabtray/TabTrayController.kt

@@ -9,11 +9,12 @@ import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.launch
 import mozilla.appservices.places.BookmarkRoot
-import mozilla.components.browser.session.Session
 import mozilla.components.browser.session.SessionManager
+import mozilla.components.browser.state.selector.findTab
 import mozilla.components.browser.state.selector.getNormalOrPrivateTabs
 import mozilla.components.browser.state.selector.normalTabs
 import mozilla.components.browser.state.state.BrowserState
+import mozilla.components.browser.state.state.TabSessionState
 import mozilla.components.browser.state.store.BrowserStore
 import mozilla.components.concept.base.profiler.Profiler
 import mozilla.components.concept.engine.prompt.ShareData
@@ -100,8 +101,8 @@ class DefaultTabTrayController(
     private val registerCollectionStorageObserver: () -> Unit,
     private val tabTrayDialogFragmentStore: TabTrayDialogFragmentStore,
     private val selectTabUseCase: TabsUseCases.SelectTabUseCase,
-    private val showChooseCollectionDialog: (List<Session>) -> Unit,
+    private val showChooseCollectionDialog: (List<TabSessionState>) -> Unit,
-    private val showAddNewCollectionDialog: (List<Session>) -> Unit,
+    private val showAddNewCollectionDialog: (List<TabSessionState>) -> Unit,
     private val showUndoSnackbarForTabs: () -> Unit,
     private val showBookmarksSnackbar: () -> Unit
 ) : TabTrayController {
@@ -129,7 +130,7 @@ class DefaultTabTrayController(
         metrics.track(Event.TabsTraySaveToCollectionPressed)
 
         val sessionList = selectedTabs.map {
-            sessionManager.findSessionById(it.id) ?: return
+            browserStore.state.findTab(it.id) ?: return
         }
 
         // Only register the observer right before moving to collection creation

app/src/main/java/org/mozilla/fenix/tabtray/TabTrayDialogFragment.kt

@@ -86,11 +86,11 @@ class TabTrayDialogFragment : AppCompatDialogFragment(), UserInteractionHandler
             else null
 
     private val collectionStorageObserver = object : TabCollectionStorage.Observer {
-        override fun onCollectionCreated(title: String, sessions: List<Session>, id: Long?) {
+        override fun onCollectionCreated(title: String, sessions: List<TabSessionState>, id: Long?) {
             showCollectionSnackbar(sessions.size, true, collectionToSelect = id)
         }
 
-        override fun onTabsAdded(tabCollection: TabCollection, sessions: List<Session>) {
+        override fun onTabsAdded(tabCollection: TabCollection, sessions: List<TabSessionState>) {
             showCollectionSnackbar(
                 sessions.size,
                 collectionToSelect = tabCollection.id
@@ -424,7 +424,7 @@ class TabTrayDialogFragment : AppCompatDialogFragment(), UserInteractionHandler
         return true
     }
 
-    private fun showChooseCollectionDialog(sessionList: List<Session>) {
+    private fun showChooseCollectionDialog(sessionList: List<TabSessionState>) {
         context?.let {
             val tabCollectionStorage = it.components.core.tabCollectionStorage
             val collections =
@@ -467,7 +467,7 @@ class TabTrayDialogFragment : AppCompatDialogFragment(), UserInteractionHandler
         }
     }
 
-    private fun showAddNewCollectionDialog(sessionList: List<Session>) {
+    private fun showAddNewCollectionDialog(sessionList: List<TabSessionState>) {
         context?.let {
             val tabCollectionStorage = it.components.core.tabCollectionStorage
             val customLayout =

app/src/test/java/org/mozilla/fenix/collections/DefaultCollectionCreationControllerTest.kt

@@ -1,3 +1,7 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
 package org.mozilla.fenix.collections
 
 import io.mockk.MockKAnnotations
@@ -10,10 +14,12 @@ import io.mockk.verifyAll
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.TestCoroutineScope
 import kotlinx.coroutines.test.runBlockingTest
-import mozilla.components.browser.session.Session
+import mozilla.components.browser.state.action.TabListAction
-import mozilla.components.browser.session.SessionManager
 import mozilla.components.browser.state.state.MediaState
+import mozilla.components.browser.state.state.createTab
+import mozilla.components.browser.state.store.BrowserStore
 import mozilla.components.feature.tab.collections.TabCollection
+import mozilla.components.support.test.ext.joinBlocking
 import org.junit.After
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertNull
@@ -35,7 +41,7 @@ class DefaultCollectionCreationControllerTest {
     @MockK(relaxed = true) private lateinit var dismiss: () -> Unit
     @MockK(relaxUnitFun = true) private lateinit var metrics: MetricController
     @MockK(relaxUnitFun = true) private lateinit var tabCollectionStorage: TabCollectionStorage
-    @MockK private lateinit var sessionManager: SessionManager
+    private lateinit var browserStore: BrowserStore
 
     @Before
     fun before() {
@@ -47,9 +53,15 @@ class DefaultCollectionCreationControllerTest {
         )
         every { store.state } answers { state }
 
+        browserStore = BrowserStore()
+
         controller = DefaultCollectionCreationController(
-            store, dismiss, metrics,
+            store,
-            tabCollectionStorage, sessionManager, testCoroutineScope
+            browserStore,
+            dismiss,
+            metrics,
+            tabCollectionStorage,
+            testCoroutineScope
         )
     }
 
@@ -60,14 +72,13 @@ class DefaultCollectionCreationControllerTest {
 
     @Test
     fun `GIVEN tab list WHEN saveCollectionName is called THEN collection should be created`() {
-        val session = mockSession(sessionId = "session-1")
+        val tab1 = createTab("https://www.mozilla.org", id = "session-1")
-        val sessions = listOf(
+        val tab2 = createTab("https://www.mozilla.org", id = "session-2")
-            session,
+
-            mockSession(sessionId = "session-2")
+        browserStore.dispatch(
-        )
+            TabListAction.AddMultipleTabsAction(listOf(tab1, tab2))
-        every { sessionManager.findSessionById("session-1") } returns session
+        ).joinBlocking()
-        every { sessionManager.findSessionById("null-session") } returns null
+
-        every { sessionManager.sessions } returns sessions
         val tabs = listOf(
             Tab("session-1", "", "", "", mediaState = MediaState.State.NONE),
             Tab("null-session", "", "", "", mediaState = MediaState.State.NONE)
@@ -76,7 +87,7 @@ class DefaultCollectionCreationControllerTest {
         controller.saveCollectionName(tabs, "name")
 
         verify { dismiss() }
-        coVerify { tabCollectionStorage.createCollection("name", listOf(session)) }
+        coVerify { tabCollectionStorage.createCollection("name", listOf(tab1)) }
         verify { metrics.track(Event.CollectionSaved(2, 1)) }
     }
 
@@ -154,13 +165,13 @@ class DefaultCollectionCreationControllerTest {
 
     @Test
     fun `WHEN selectCollection is called THEN add tabs should be added to collection`() {
-        val session = mockSession(sessionId = "session-1")
+        val tab1 = createTab("https://www.mozilla.org", id = "session-1")
-        val sessions = listOf(
+        val tab2 = createTab("https://www.mozilla.org", id = "session-2")
-            session,
+
-            mockSession(sessionId = "session-2")
+        browserStore.dispatch(
-        )
+            TabListAction.AddMultipleTabsAction(listOf(tab1, tab2))
-        every { sessionManager.findSessionById("session-1") } returns session
+        ).joinBlocking()
-        every { sessionManager.sessions } returns sessions
+
         val tabs = listOf(
             Tab("session-1", "", "", "", mediaState = MediaState.State.NONE)
         )
@@ -169,7 +180,7 @@ class DefaultCollectionCreationControllerTest {
         controller.selectCollection(collection, tabs)
 
         verify { dismiss() }
-        coVerify { tabCollectionStorage.addTabsToCollection(collection, listOf(session)) }
+        coVerify { tabCollectionStorage.addTabsToCollection(collection, listOf(tab1)) }
         verify { metrics.track(Event.CollectionTabsAdded(2, 1)) }
     }
 
@@ -246,14 +257,4 @@ class DefaultCollectionCreationControllerTest {
 
         verify { store.dispatch(CollectionCreationAction.StepChanged(SaveCollectionStep.SelectCollection, 2)) }
     }
-
-    private fun mockSession(
-        sessionId: String? = null,
-        isPrivate: Boolean = false,
-        isCustom: Boolean = false
-    ) = mockk<Session> {
-        sessionId?.let { every { id } returns it }
-        every { private } returns isPrivate
-        every { isCustomTabSession() } returns isCustom
-    }
 }

app/src/test/java/org/mozilla/fenix/home/DefaultSessionControlControllerTest.kt

@@ -15,7 +15,9 @@ import kotlinx.coroutines.test.TestCoroutineScope
 import mozilla.components.browser.session.SessionManager
 import mozilla.components.browser.state.search.SearchEngine
 import mozilla.components.browser.state.state.BrowserState
+import mozilla.components.browser.state.state.ReaderState
 import mozilla.components.browser.state.state.SearchState
+import mozilla.components.browser.state.state.recover.RecoverableTab
 import mozilla.components.browser.state.store.BrowserStore
 import mozilla.components.concept.engine.Engine
 import mozilla.components.feature.session.SessionUseCases
@@ -108,6 +110,8 @@ class DefaultSessionControlControllerTest {
         every { activity.components.analytics } returns analytics
         every { analytics.metrics } returns metrics
 
+        val restoreUseCase: TabsUseCases.RestoreUseCase = mockk(relaxed = true)
+
         controller = DefaultSessionControlController(
             activity = activity,
             store = store,
@@ -118,6 +122,7 @@ class DefaultSessionControlControllerTest {
             tabCollectionStorage = tabCollectionStorage,
             addTabUseCase = tabsUseCases.addTab,
             reloadUrlUseCase = reloadUrlUseCase.reload,
+            restoreUseCase = restoreUseCase,
             fragmentStore = fragmentStore,
             navController = navController,
             viewLifecycleScope = scope,
@@ -173,12 +178,22 @@ class DefaultSessionControlControllerTest {
 
     @Test
     fun `handleCollectionOpenTabClicked onTabRestored`() {
+        val restoredTab = RecoverableTab(
+            id = "test",
+            parentId = null,
+            url = "https://www.mozilla.org",
+            title = "Mozilla",
+            state = null,
+            contextId = null,
+            readerState = ReaderState(),
+            lastAccess = 0,
+            private = false
+        )
+
         val tab = mockk<ComponentTab> {
-            every { restore(activity, engine, restoreSessionId = false) } returns mockk {
+            every { restore(activity, engine, restoreSessionId = false) } returns restoredTab
-                every { session } returns mockk()
-                every { engineSessionState } returns mockk()
-            }
         }
+
         controller.handleCollectionOpenTabClicked(tab)
 
         verify { metrics.track(Event.CollectionTabRestored) }

app/src/test/java/org/mozilla/fenix/tabtray/DefaultTabTrayControllerTest.kt

@@ -21,6 +21,7 @@ import kotlinx.coroutines.test.TestCoroutineScope
 import mozilla.components.browser.session.Session
 import mozilla.components.browser.session.SessionManager
 import mozilla.components.browser.state.state.BrowserState
+import mozilla.components.browser.state.state.TabSessionState
 import mozilla.components.browser.state.state.createTab
 import mozilla.components.browser.state.store.BrowserStore
 import mozilla.components.concept.base.profiler.Profiler
@@ -48,20 +49,12 @@ class DefaultTabTrayControllerTest {
     private val profiler: Profiler? = mockk(relaxed = true)
     private val navController: NavController = mockk()
     private val sessionManager: SessionManager = mockk(relaxed = true)
-    var store = BrowserStore(
-        BrowserState(
-            tabs = listOf(
-                createTab(url = "http://firefox.com", id = "5678"),
-                createTab(url = "http://mozilla.org", id = "1234")
-            ), selectedTabId = "1234"
-        )
-    )
     private val browsingModeManager: BrowsingModeManager = mockk(relaxed = true)
     private val dismissTabTray: (() -> Unit) = mockk(relaxed = true)
     private val dismissTabTrayAndNavigateHome: ((String) -> Unit) = mockk(relaxed = true)
     private val registerCollectionStorageObserver: (() -> Unit) = mockk(relaxed = true)
-    private val showChooseCollectionDialog: ((List<Session>) -> Unit) = mockk(relaxed = true)
+    private val showChooseCollectionDialog: ((List<TabSessionState>) -> Unit) = mockk(relaxed = true)
-    private val showAddNewCollectionDialog: ((List<Session>) -> Unit) = mockk(relaxed = true)
+    private val showAddNewCollectionDialog: ((List<TabSessionState>) -> Unit) = mockk(relaxed = true)
     private val tabCollectionStorage: TabCollectionStorage = mockk(relaxed = true)
     private val bookmarksStorage: BookmarksStorage = mockk(relaxed = true)
     private val tabCollection: TabCollection = mockk()
@@ -76,6 +69,9 @@ class DefaultTabTrayControllerTest {
 
     private lateinit var controller: DefaultTabTrayController
 
+    private val tab1 = createTab(url = "http://firefox.com", id = "5678")
+    private val tab2 = createTab(url = "http://mozilla.org", id = "1234")
+
     private val session = Session(
         "mozilla.org",
         true
@@ -90,6 +86,12 @@ class DefaultTabTrayControllerTest {
     fun setUp() {
         mockkStatic("org.mozilla.fenix.ext.SessionManagerKt")
 
+        val store = BrowserStore(
+            BrowserState(
+                tabs = listOf(tab1, tab2), selectedTabId = tab2.id
+            )
+        )
+
         every { sessionManager.sessionsOfType(private = true) } returns listOf(session).asSequence()
         every { sessionManager.sessionsOfType(private = false) } returns listOf(nonPrivateSession).asSequence()
         every { sessionManager.createSessionSnapshot(any()) } returns SessionManager.Snapshot.Item(
@@ -268,13 +270,14 @@ class DefaultTabTrayControllerTest {
 
     @Test
     fun onSaveToCollectionClicked() {
-        val tab = Tab("1234", "mozilla.org")
+        val tab = Tab(tab2.id, tab2.content.url)
 
         controller.handleSaveToCollectionClicked(setOf(tab))
+
         verify {
             metrics.track(Event.TabsTraySaveToCollectionPressed)
             registerCollectionStorageObserver()
-            showChooseCollectionDialog(listOf(session))
+            showChooseCollectionDialog(listOf(tab2))
         }
     }