Closes #2690 - Clean up CreateCollection viewmodel (#4731)

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

@@ -5,8 +5,6 @@
 package org.mozilla.fenix.collections
 
 import android.view.ViewGroup
-import androidx.lifecycle.ViewModel
-import androidx.lifecycle.ViewModelProvider
 import org.mozilla.fenix.home.sessioncontrol.Tab
 import org.mozilla.fenix.home.sessioncontrol.TabCollection
 import org.mozilla.fenix.mvi.Action
@@ -18,19 +16,19 @@ import org.mozilla.fenix.mvi.UIComponentViewModelBase
 import org.mozilla.fenix.mvi.UIComponentViewModelProvider
 import org.mozilla.fenix.mvi.ViewState
 
-sealed class SaveCollectionStep {
+enum class SaveCollectionStep {
-    object SelectTabs : SaveCollectionStep()
+    SelectTabs,
-    object SelectCollection : SaveCollectionStep()
+    SelectCollection,
-    object NameCollection : SaveCollectionStep()
+    NameCollection,
-    object RenameCollection : SaveCollectionStep()
+    RenameCollection
 }
 
 data class CollectionCreationState(
-    val tabs: List<Tab> = listOf(),
+    val tabs: List<Tab> = emptyList(),
-    val selectedTabs: Set<Tab> = setOf(),
+    val selectedTabs: Set<Tab> = emptySet(),
     val saveCollectionStep: SaveCollectionStep = SaveCollectionStep.SelectTabs,
-    val tabCollections: List<TabCollection> = listOf(),
+    val tabCollections: List<TabCollection> = emptyList(),
-    val selectedTabCollection: TabCollection?
+    val selectedTabCollection: TabCollection? = null
 ) : ViewState
 
 sealed class CollectionCreationChange : Change {
@@ -84,36 +82,17 @@ class CollectionCreationViewModel(
         reducer
     ) {
 
-    class Factory(
-        private val initialState: CollectionCreationState
-    ) : ViewModelProvider.Factory {
-        @Suppress("UNCHECKED_CAST")
-        override fun <T : ViewModel?> create(modelClass: Class<T>): T =
-            CollectionCreationViewModel(initialState) as T
-    }
-
     companion object {
-        val reducer: Reducer<CollectionCreationState, CollectionCreationChange> =
+        val reducer: Reducer<CollectionCreationState, CollectionCreationChange> = { state, change ->
-            { state, change ->
+            when (change) {
-                when (change) {
+                is CollectionCreationChange.AddAllTabs -> state.copy(selectedTabs = state.tabs.toSet())
-                    is CollectionCreationChange.AddAllTabs -> state.copy(selectedTabs = state.tabs.toSet())
+                is CollectionCreationChange.RemoveAllTabs -> state.copy(selectedTabs = emptySet())
-                    is CollectionCreationChange.RemoveAllTabs -> state.copy(selectedTabs = setOf())
+                is CollectionCreationChange.TabListChange -> state.copy(tabs = change.tabs)
-                    is CollectionCreationChange.TabListChange -> state.copy(tabs = change.tabs)
+                is CollectionCreationChange.TabAdded -> state.copy(selectedTabs = state.selectedTabs + change.tab)
-                    is CollectionCreationChange.TabAdded -> {
+                is CollectionCreationChange.TabRemoved -> state.copy(selectedTabs = state.selectedTabs - change.tab)
-                        val selectedTabs = state.selectedTabs + setOf(change.tab)
+                is CollectionCreationChange.StepChanged -> state.copy(saveCollectionStep = change.saveCollectionStep)
-                        state.copy(selectedTabs = selectedTabs)
+                is CollectionCreationChange.CollectionSelected -> state.copy(selectedTabCollection = change.collection)
-                    }
-                    is CollectionCreationChange.TabRemoved -> {
-                        val selectedTabs = state.selectedTabs - setOf(change.tab)
-                        state.copy(selectedTabs = selectedTabs)
-                    }
-                    is CollectionCreationChange.StepChanged -> {
-                        state.copy(saveCollectionStep = change.saveCollectionStep)
-                    }
-                    is CollectionCreationChange.CollectionSelected -> {
-                        state.copy(selectedTabCollection = change.collection)
-                    }
-                }
             }
+        }
     }
 }

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

@@ -7,6 +7,8 @@ package org.mozilla.fenix.collections
 import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
+import androidx.core.view.isGone
+import androidx.core.view.isInvisible
 import androidx.recyclerview.widget.DiffUtil
 import androidx.recyclerview.widget.RecyclerView
 import io.reactivex.Observer
@@ -42,8 +44,7 @@ class CollectionCreationTabListAdapter(
                     } else if (checkChanged.shouldBeUnchecked) {
                         holder.itemView.tab_selected_checkbox.isChecked = false
                     }
-                    holder.itemView.tab_selected_checkbox.visibility =
+                    holder.itemView.tab_selected_checkbox.isGone = checkChanged.shouldHideCheckBox
-                        if (checkChanged.shouldHideCheckBox) View.GONE else View.VISIBLE
                 }
             }
         }
@@ -120,7 +121,6 @@ data class CheckChanged(val shouldBeChecked: Boolean, val shouldBeUnchecked: Boo
 
 class TabViewHolder(view: View) : RecyclerView.ViewHolder(view) {
 
-    private var tab: Tab? = null
     private val checkbox = view.tab_selected_checkbox!!
 
     init {
@@ -130,10 +130,9 @@ class TabViewHolder(view: View) : RecyclerView.ViewHolder(view) {
     }
 
     fun bind(tab: Tab, isSelected: Boolean, shouldHideCheckBox: Boolean) {
-        this.tab = tab
         itemView.hostname.text = tab.hostname
         itemView.tab_title.text = tab.title
-        checkbox.visibility = if (shouldHideCheckBox) View.INVISIBLE else View.VISIBLE
+        checkbox.isInvisible = shouldHideCheckBox
         itemView.isClickable = !shouldHideCheckBox
         if (checkbox.isChecked != isSelected) {
             checkbox.isChecked = isSelected

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

@@ -76,22 +76,17 @@ class CollectionCreationUIView(
         }
 
         view.name_collection_edittext.filters += InputFilter.LengthFilter(COLLECTION_NAME_MAX_LENGTH)
-        view.name_collection_edittext.setOnEditorActionListener { v, actionId, _ ->
+        view.name_collection_edittext.setOnEditorActionListener { view, actionId, _ ->
-            if (actionId == EditorInfo.IME_ACTION_DONE && v.text.toString().isNotBlank()) {
+            val text = view.text.toString()
+            if (actionId == EditorInfo.IME_ACTION_DONE && text.isNotBlank()) {
                 when (step) {
-                    is SaveCollectionStep.NameCollection -> {
+                    SaveCollectionStep.NameCollection ->
-                        actionEmitter.onNext(
+                        CollectionCreationAction.SaveCollectionName(selectedTabs.toList(), text)
-                            CollectionCreationAction.SaveCollectionName(
+                    SaveCollectionStep.RenameCollection ->
-                                selectedTabs.toList(),
+                        selectedCollection?.let { CollectionCreationAction.RenameCollection(it, text) }
-                                v.text.toString()
+                    else -> null
-                            )
+                }?.let { action ->
-                        )
+                    actionEmitter.onNext(action)
-                    }
-                    is SaveCollectionStep.RenameCollection -> {
-                        selectedCollection?.let {
-                            actionEmitter.onNext(CollectionCreationAction.RenameCollection(it, v.text.toString()))
-                        }
-                    }
                 }
             }
             false
@@ -116,7 +111,7 @@ class CollectionCreationUIView(
         selectedCollection = it.selectedTabCollection
 
         when (it.saveCollectionStep) {
-            is SaveCollectionStep.SelectTabs -> {
+            SaveCollectionStep.SelectTabs -> {
                 view.context.components.analytics.metrics.track(Event.CollectionTabSelectOpened)
 
                 view.tab_list.isClickable = true
@@ -194,7 +189,7 @@ class CollectionCreationUIView(
                     View.VISIBLE
                 }
             }
-            is SaveCollectionStep.SelectCollection -> {
+            SaveCollectionStep.SelectCollection -> {
                 view.tab_list.isClickable = false
 
                 save_button.visibility = View.GONE
@@ -224,7 +219,7 @@ class CollectionCreationUIView(
                 back_button.text =
                     view.context.getString(R.string.create_collection_select_collection)
             }
-            is SaveCollectionStep.NameCollection -> {
+            SaveCollectionStep.NameCollection -> {
                 view.tab_list.isClickable = false
 
                 collectionCreationTabListAdapter.updateData(it.selectedTabs.toList(), it.selectedTabs, true)
@@ -264,7 +259,7 @@ class CollectionCreationUIView(
                 back_button.text =
                     view.context.getString(R.string.create_collection_name_collection)
             }
-            is SaveCollectionStep.RenameCollection -> {
+            SaveCollectionStep.RenameCollection -> {
                 view.tab_list.isClickable = false
 
                 it.selectedTabCollection?.let { tabCollection ->
@@ -322,24 +317,11 @@ class CollectionCreationUIView(
     }
 
     fun onKey(keyCode: Int, event: KeyEvent?): Boolean {
-        if (event?.action == KeyEvent.ACTION_UP && keyCode == KeyEvent.KEYCODE_BACK) {
+        return if (event?.action == KeyEvent.ACTION_UP && keyCode == KeyEvent.KEYCODE_BACK) {
-            when (step) {
+            actionEmitter.onNext(CollectionCreationAction.BackPressed(step))
-                SaveCollectionStep.SelectTabs -> {
+            true
-                    actionEmitter.onNext(CollectionCreationAction.BackPressed(SaveCollectionStep.SelectTabs))
-                }
-                SaveCollectionStep.SelectCollection -> {
-                    actionEmitter.onNext(CollectionCreationAction.BackPressed(SaveCollectionStep.SelectCollection))
-                }
-                SaveCollectionStep.NameCollection -> {
-                    actionEmitter.onNext(CollectionCreationAction.BackPressed(SaveCollectionStep.NameCollection))
-                }
-                SaveCollectionStep.RenameCollection -> {
-                    actionEmitter.onNext(CollectionCreationAction.BackPressed(SaveCollectionStep.RenameCollection))
-                }
-            }
-            return true
         } else {
-            return false
+            false
         }
     }
 

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

@@ -50,15 +50,7 @@ class CreateCollectionFragment : DialogFragment() {
                 this,
                 CollectionCreationViewModel::class.java
             ) {
-                CollectionCreationViewModel(
+                CollectionCreationViewModel(viewModel.state)
-                    CollectionCreationState(
-                        viewModel.tabs,
-                        viewModel.selectedTabs,
-                        viewModel.saveCollectionStep,
-                        viewModel.tabCollections,
-                        viewModel.selectedTabCollection
-                    )
-                )
             }
         )
         return view
@@ -87,7 +79,11 @@ class CreateCollectionFragment : DialogFragment() {
                     getManagedEmitter<CollectionCreationChange>()
                         .onNext(
                             CollectionCreationChange.StepChanged(
-                                viewModel.tabCollections.getStepForCollectionsSize()
+                                if (viewModel.state.tabCollections.isEmpty()) {
+                                    SaveCollectionStep.NameCollection
+                                } else {
+                                    SaveCollectionStep.SelectCollection
+                                }
                             )
                         )
                 }
@@ -161,41 +157,38 @@ class CreateCollectionFragment : DialogFragment() {
     }
 
     private fun handleBackPress(backPressFrom: SaveCollectionStep) {
-        when (backPressFrom) {
+        val newStep = stepBack(backPressFrom)
-            SaveCollectionStep.SelectTabs -> dismiss()
+        if (newStep != null) {
-            SaveCollectionStep.SelectCollection -> {
+            getManagedEmitter<CollectionCreationChange>().onNext(CollectionCreationChange.StepChanged(newStep))
-                if (viewModel.tabs.size <= 1) dismiss() else {
+        } else {
-                    getManagedEmitter<CollectionCreationChange>().onNext(
+            dismiss()
-                        CollectionCreationChange.StepChanged(SaveCollectionStep.SelectTabs)
+        }
-                    )
+    }
-                }
+
+    private fun stepBack(backFromStep: SaveCollectionStep): SaveCollectionStep? {
+        val state = viewModel.state
+        return when (backFromStep) {
+            SaveCollectionStep.SelectTabs, SaveCollectionStep.RenameCollection -> null
+            SaveCollectionStep.SelectCollection -> if (state.tabs.size <= 1) {
+                stepBack(SaveCollectionStep.SelectTabs)
+            } else {
+                SaveCollectionStep.SelectTabs
             }
-            SaveCollectionStep.NameCollection -> {
+            SaveCollectionStep.NameCollection -> if (state.tabCollections.isEmpty()) {
-                if (viewModel.tabCollections.isEmpty() && viewModel.tabs.size == 1) {
+                stepBack(SaveCollectionStep.SelectCollection)
-                    dismiss()
+            } else {
-                } else {
+                SaveCollectionStep.SelectCollection
-                    getManagedEmitter<CollectionCreationChange>()
-                        .onNext(
-                            CollectionCreationChange.StepChanged(
-                                viewModel.tabCollections.getBackStepForCollectionsSize()
-                            )
-                        )
-                }
-            }
-            SaveCollectionStep.RenameCollection -> {
-                dismiss()
             }
         }
     }
 
     private fun closeTabsIfNecessary(tabs: List<Tab>) {
         // Only close the tabs if the user is not on the BrowserFragment
-        if (viewModel.previousFragmentId == R.id.browserFragment) { return }
+        if (viewModel.previousFragmentId == R.id.browserFragment) {
-
+            val components = requireComponents
-        tabs.forEach {
+            tabs.asSequence()
-            requireComponents.core.sessionManager.findSessionById(it.sessionId)?.let { session ->
+                .mapNotNull { tab -> components.core.sessionManager.findSessionById(tab.sessionId) }
-                requireComponents.useCases.tabsUseCases.removeTab.invoke(session)
+                .forEach { session -> components.useCases.tabsUseCases.removeTab(session) }
-            }
         }
     }
 }

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

@@ -4,26 +4,47 @@
 
 package org.mozilla.fenix.collections
 
-import android.view.View
 import androidx.lifecycle.ViewModel
 import org.mozilla.fenix.home.sessioncontrol.Tab
 import org.mozilla.fenix.home.sessioncontrol.TabCollection
 
 class CreateCollectionViewModel : ViewModel() {
-    var selectedTabs = mutableSetOf<Tab>()
+    var state = CollectionCreationState()
-    var tabs = listOf<Tab>()
+        private set
-    var saveCollectionStep: SaveCollectionStep = SaveCollectionStep.SelectTabs
+
-    var tabCollections = listOf<TabCollection>()
-    var selectedTabCollection: TabCollection? = null
-    var snackbarAnchorView: View? = null
     var previousFragmentId: Int? = null
 
-    fun getStepForTabsAndCollectionSize(): SaveCollectionStep =
+    fun updateCollection(
-        if (tabs.size > 1) SaveCollectionStep.SelectTabs else tabCollections.getStepForCollectionsSize()
+        tabs: List<Tab>,
+        saveCollectionStep: SaveCollectionStep,
+        selectedTabCollection: TabCollection,
+        cachedTabCollections: List<TabCollection>
+    ) {
+        state = CollectionCreationState(
+            tabs = tabs,
+            selectedTabs = if (tabs.size == 1) setOf(tabs.first()) else emptySet(),
+            tabCollections = cachedTabCollections.reversed(),
+            selectedTabCollection = selectedTabCollection,
+            saveCollectionStep = saveCollectionStep
+        )
+    }
+
+    fun saveTabToCollection(
+        tabs: List<Tab>,
+        selectedTab: Tab?,
+        cachedTabCollections: List<TabCollection>
+    ) {
+        val tabCollections = cachedTabCollections.reversed()
+        state = CollectionCreationState(
+            tabs = tabs,
+            selectedTabs = selectedTab?.let { setOf(it) } ?: emptySet(),
+            tabCollections = tabCollections,
+            selectedTabCollection = null,
+            saveCollectionStep = when {
+                tabs.size > 1 -> SaveCollectionStep.SelectTabs
+                tabCollections.isNotEmpty() -> SaveCollectionStep.SelectCollection
+                else -> SaveCollectionStep.NameCollection
+            }
+        )
+    }
 }
-
-fun List<TabCollection>.getStepForCollectionsSize(): SaveCollectionStep =
-    if (isEmpty()) SaveCollectionStep.NameCollection else SaveCollectionStep.SelectCollection
-
-fun List<TabCollection>.getBackStepForCollectionsSize(): SaveCollectionStep =
-    if (isEmpty()) SaveCollectionStep.SelectTabs else SaveCollectionStep.SelectCollection

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

@@ -35,10 +35,8 @@ class SaveCollectionListAdapter(
         val collection = tabCollections[position]
         holder.bind(collection)
         holder.itemView.setOnClickListener {
-            collection.apply {
+            val action = CollectionCreationAction.SelectCollection(collection, selectedTabs.toList())
-                val action = CollectionCreationAction.SelectCollection(this, selectedTabs.toList())
+            actionEmitter.onNext(action)
-                actionEmitter.onNext(action)
-            }
         }
     }
 

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

@@ -20,7 +20,6 @@ import org.mozilla.fenix.R
 import org.mozilla.fenix.browser.BrowserFragment
 import org.mozilla.fenix.browser.BrowserFragmentDirections
 import org.mozilla.fenix.collections.CreateCollectionViewModel
-import org.mozilla.fenix.collections.getStepForCollectionsSize
 import org.mozilla.fenix.components.metrics.Event
 import org.mozilla.fenix.ext.components
 import org.mozilla.fenix.ext.nav
@@ -123,14 +122,12 @@ class DefaultBrowserToolbarController(
                 context.components.analytics.metrics
                     .track(Event.CollectionSaveButtonPressed(TELEMETRY_BROWSER_IDENTIFIER))
 
-                currentSession?.toTab(context)?.let {
+                currentSession?.toTab(context)?.let { currentSessionAsTab ->
-                    viewModel.tabs = listOf(it)
+                    viewModel.saveTabToCollection(
-                    val selectedSet = mutableSetOf(it)
+                        tabs = listOf(currentSessionAsTab),
-                    viewModel.selectedTabs = selectedSet
+                        selectedTab = currentSessionAsTab,
-                    viewModel.tabCollections =
+                        cachedTabCollections = context.components.core.tabCollectionStorage.cachedTabCollections
-                        context.components.core.tabCollectionStorage.cachedTabCollections.reversed()
+                    )
-                    viewModel.saveCollectionStep = viewModel.tabCollections.getStepForCollectionsSize()
-                    viewModel.snackbarAnchorView = nestedScrollQuickActionView
                     viewModel.previousFragmentId = R.id.browserFragment
 
                     val directions = BrowserFragmentDirections.actionBrowserFragmentToCreateCollectionFragment()

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

@@ -25,7 +25,7 @@ import androidx.lifecycle.OnLifecycleEvent
 import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.lifecycleScope
 import androidx.navigation.fragment.FragmentNavigator
-import androidx.navigation.fragment.NavHostFragment.findNavController
+import androidx.navigation.fragment.findNavController
 import androidx.recyclerview.widget.LinearLayoutManager
 import androidx.recyclerview.widget.RecyclerView
 import androidx.recyclerview.widget.RecyclerView.SCROLL_STATE_IDLE
@@ -309,7 +309,7 @@ class HomeFragment : Fragment(), AccountObserver {
             is TabAction.SaveTabGroup -> {
                 if ((activity as HomeActivity).browsingModeManager.mode.isPrivate) return
                 invokePendingDeleteJobs()
-                showCollectionCreationFragment(action.selectedTabSessionId)
+                saveTabToCollection(action.selectedTabSessionId)
             }
             is TabAction.Select -> {
                 invokePendingDeleteJobs()
@@ -437,16 +437,10 @@ class HomeFragment : Fragment(), AccountObserver {
             }
             is CollectionAction.AddTab -> {
                 requireComponents.analytics.metrics.track(Event.CollectionAddTabPressed)
-                showCollectionCreationFragment(
+                updateCollection(action.collection, SaveCollectionStep.SelectTabs)
-                    selectedTabCollection = action.collection,
-                    step = SaveCollectionStep.SelectTabs
-                )
             }
             is CollectionAction.Rename -> {
-                showCollectionCreationFragment(
+                updateCollection(action.collection, SaveCollectionStep.RenameCollection)
-                    selectedTabCollection = action.collection,
-                    step = SaveCollectionStep.RenameCollection
-                )
                 requireComponents.analytics.metrics.track(Event.CollectionRenamePressed)
             }
             is CollectionAction.OpenTab -> {
@@ -647,27 +641,18 @@ class HomeFragment : Fragment(), AccountObserver {
     }
 
     private fun showCollectionCreationFragment(
-        selectedTabId: String? = null,
+        setupViewModel: (CreateCollectionViewModel, tabs: List<Tab>, cachedTabCollections: List<TabCollection>) -> Unit
-        selectedTabCollection: TabCollection? = null,
-        step: SaveCollectionStep? = null
     ) {
-        if (findNavController(this).currentDestination?.id == R.id.createCollectionFragment) return
+        if (findNavController().currentDestination?.id == R.id.createCollectionFragment) return
-
-        val tabs = getListOfSessions().toTabs()
 
         val viewModel: CreateCollectionViewModel by activityViewModels {
             ViewModelProvider.NewInstanceFactory() // this is a workaround for #4652
         }
-        viewModel.tabs = tabs
+
-        val selectedTabs =
+        val tabs = getListOfSessions().toTabs()
-            tabs.find { tab -> tab.sessionId == selectedTabId }
+        val cachedTabCollections = requireComponents.core.tabCollectionStorage.cachedTabCollections
-                ?: if (tabs.size == 1) tabs[0] else null
+        setupViewModel(viewModel, tabs, cachedTabCollections)
-        val selectedSet = if (selectedTabs == null) mutableSetOf() else mutableSetOf(selectedTabs)
+
-        viewModel.selectedTabs = selectedSet
-        viewModel.tabCollections = requireComponents.core.tabCollectionStorage.cachedTabCollections.reversed()
-        viewModel.selectedTabCollection = selectedTabCollection
-        viewModel.saveCollectionStep =
-            step ?: viewModel.getStepForTabsAndCollectionSize()
         viewModel.previousFragmentId = R.id.homeFragment
 
         // Only register the observer right before moving to collection creation
@@ -679,6 +664,27 @@ class HomeFragment : Fragment(), AccountObserver {
         }
     }
 
+    private fun saveTabToCollection(selectedTabId: String?) {
+        showCollectionCreationFragment { viewModel, tabs, cachedTabCollections ->
+            viewModel.saveTabToCollection(
+                tabs = tabs,
+                selectedTab = tabs.find { it.sessionId == selectedTabId } ?: if (tabs.size == 1) tabs[0] else null,
+                cachedTabCollections = cachedTabCollections
+            )
+        }
+    }
+
+    private fun updateCollection(selectedTabCollection: TabCollection, step: SaveCollectionStep) {
+        showCollectionCreationFragment { viewModel, tabs, cachedTabCollections ->
+            viewModel.updateCollection(
+                tabs = tabs,
+                saveCollectionStep = step,
+                selectedTabCollection = selectedTabCollection,
+                cachedTabCollections = cachedTabCollections
+            )
+        }
+    }
+
     private fun share(url: String? = null, tabs: List<ShareTab>? = null) {
         val directions =
             HomeFragmentDirections.actionHomeFragmentToShareFragment(

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

@@ -0,0 +1,150 @@
+/* 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
+import io.mockk.mockk
+import mozilla.components.feature.tab.collections.TabCollection
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertNull
+import org.junit.Before
+import org.junit.Test
+import org.mozilla.fenix.home.sessioncontrol.Tab
+
+class CreateCollectionViewModelTest {
+
+    private lateinit var viewModel: CreateCollectionViewModel
+
+    @Before
+    fun setup() {
+        MockKAnnotations.init(this)
+
+        viewModel = CreateCollectionViewModel()
+    }
+
+    @Test
+    fun `initial state defaults`() {
+        assertEquals(
+            CollectionCreationState(
+                tabs = emptyList(),
+                selectedTabs = emptySet(),
+                saveCollectionStep = SaveCollectionStep.SelectTabs,
+                tabCollections = emptyList(),
+                selectedTabCollection = null
+            ),
+            viewModel.state
+        )
+        assertNull(viewModel.previousFragmentId)
+    }
+
+    @Test
+    fun `updateCollection copies tabs to state`() {
+        val tabs = listOf<Tab>(mockk(), mockk())
+        val tabCollections = listOf<TabCollection>(mockk(), mockk())
+        val selectedCollection: TabCollection = mockk()
+        viewModel.updateCollection(
+            tabs = tabs,
+            saveCollectionStep = SaveCollectionStep.SelectCollection,
+            selectedTabCollection = selectedCollection,
+            cachedTabCollections = tabCollections
+        )
+        assertEquals(tabs, viewModel.state.tabs)
+        assertEquals(SaveCollectionStep.SelectCollection, viewModel.state.saveCollectionStep)
+        assertEquals(selectedCollection, viewModel.state.selectedTabCollection)
+        assertEquals(tabCollections.reversed(), viewModel.state.tabCollections)
+    }
+
+    @Test
+    fun `updateCollection selects the only tab`() {
+        val tab: Tab = mockk()
+        viewModel.updateCollection(
+            tabs = listOf(tab),
+            saveCollectionStep = mockk(),
+            selectedTabCollection = mockk(),
+            cachedTabCollections = emptyList()
+        )
+        assertEquals(setOf(tab), viewModel.state.selectedTabs)
+
+        viewModel.updateCollection(
+            tabs = listOf(tab, mockk()),
+            saveCollectionStep = mockk(),
+            selectedTabCollection = mockk(),
+            cachedTabCollections = emptyList()
+        )
+        assertEquals(emptySet<Tab>(), viewModel.state.selectedTabs)
+
+        viewModel.updateCollection(
+            tabs = emptyList(),
+            saveCollectionStep = mockk(),
+            selectedTabCollection = mockk(),
+            cachedTabCollections = emptyList()
+        )
+        assertEquals(emptySet<Tab>(), viewModel.state.selectedTabs)
+    }
+
+    @Test
+    fun `saveTabToCollection copies tabs to state`() {
+        val tabs = listOf<Tab>(mockk(), mockk())
+        val tabCollections = listOf<TabCollection>(mockk(), mockk())
+        viewModel.saveTabToCollection(
+            tabs = tabs,
+            selectedTab = null,
+            cachedTabCollections = tabCollections
+        )
+        assertEquals(tabs, viewModel.state.tabs)
+        assertEquals(SaveCollectionStep.SelectTabs, viewModel.state.saveCollectionStep)
+        assertNull(viewModel.state.selectedTabCollection)
+        assertEquals(tabCollections.reversed(), viewModel.state.tabCollections)
+    }
+
+    @Test
+    fun `saveTabToCollection selects selectedTab`() {
+        val tab: Tab = mockk()
+        viewModel.saveTabToCollection(
+            tabs = listOf(mockk()),
+            selectedTab = tab,
+            cachedTabCollections = emptyList()
+        )
+        assertEquals(setOf(tab), viewModel.state.selectedTabs)
+
+        viewModel.saveTabToCollection(
+            tabs = listOf(mockk()),
+            selectedTab = null,
+            cachedTabCollections = emptyList()
+        )
+        assertEquals(emptySet<Tab>(), viewModel.state.selectedTabs)
+    }
+
+    @Test
+    fun `saveTabToCollection sets saveCollectionStep`() {
+        viewModel.saveTabToCollection(
+            tabs = listOf(mockk(), mockk()),
+            selectedTab = null,
+            cachedTabCollections = listOf(mockk())
+        )
+        assertEquals(SaveCollectionStep.SelectTabs, viewModel.state.saveCollectionStep)
+
+        viewModel.saveTabToCollection(
+            tabs = listOf(mockk()),
+            selectedTab = null,
+            cachedTabCollections = listOf(mockk())
+        )
+        assertEquals(SaveCollectionStep.SelectCollection, viewModel.state.saveCollectionStep)
+
+        viewModel.saveTabToCollection(
+            tabs = emptyList(),
+            selectedTab = null,
+            cachedTabCollections = listOf(mockk())
+        )
+        assertEquals(SaveCollectionStep.SelectCollection, viewModel.state.saveCollectionStep)
+
+        viewModel.saveTabToCollection(
+            tabs = emptyList(),
+            selectedTab = null,
+            cachedTabCollections = emptyList()
+        )
+        assertEquals(SaveCollectionStep.NameCollection, viewModel.state.saveCollectionStep)
+    }
+}

app/src/test/java/org/mozilla/fenix/components/toolbar/DefaultBrowserToolbarControllerTest.kt

@@ -29,7 +29,6 @@ import org.mozilla.fenix.R
 import org.mozilla.fenix.browser.BrowserFragment
 import org.mozilla.fenix.browser.BrowserFragmentDirections
 import org.mozilla.fenix.collections.CreateCollectionViewModel
-import org.mozilla.fenix.collections.SaveCollectionStep
 import org.mozilla.fenix.components.Analytics
 import org.mozilla.fenix.components.metrics.Event
 import org.mozilla.fenix.components.metrics.MetricController
@@ -322,11 +321,13 @@ class DefaultBrowserToolbarControllerTest {
 
         verify { metrics.track(Event.BrowserMenuItemTapped(Event.BrowserMenuItemTapped.Item.SAVE_TO_COLLECTION)) }
         verify { metrics.track(Event.CollectionSaveButtonPressed(DefaultBrowserToolbarController.TELEMETRY_BROWSER_IDENTIFIER)) }
-        verify { viewModel.tabs = listOf(currentSessionAsTab) }
+        verify {
-        verify { viewModel.selectedTabs = mutableSetOf(currentSessionAsTab) }
+            viewModel.saveTabToCollection(
-        verify { viewModel.tabCollections = cachedTabCollections.reversed() }
+                listOf(currentSessionAsTab),
-        verify { viewModel.saveCollectionStep = SaveCollectionStep.SelectCollection }
+                currentSessionAsTab,
-        verify { viewModel.snackbarAnchorView = nestedScrollQuickActionView }
+                cachedTabCollections
+            )
+        }
         verify { viewModel.previousFragmentId = R.id.browserFragment }
         verify {
             val directions = BrowserFragmentDirections