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

2019-08-20 13:45:41 -04:00 · 2019-08-20 13:45:41 -04:00 · ac6c1ec2ee
parent 9ee8c00928
commit ac6c1ec2ee
10 changed files with 305 additions and 179 deletions
--- a/app/src/main/java/org/mozilla/fenix/collections/CollectionCreationComponent.kt
+++ b/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 {
-    object SelectTabs : SaveCollectionStep()
-    object SelectCollection : SaveCollectionStep()
-    object NameCollection : SaveCollectionStep()
-    object RenameCollection : SaveCollectionStep()
+enum class SaveCollectionStep {
+    SelectTabs,
+    SelectCollection,
+    NameCollection,
+    RenameCollection
 }

 data class CollectionCreationState(
-    val tabs: List<Tab> = listOf(),
-    val selectedTabs: Set<Tab> = setOf(),
+    val tabs: List<Tab> = emptyList(),
+    val selectedTabs: Set<Tab> = emptySet(),
    val saveCollectionStep: SaveCollectionStep = SaveCollectionStep.SelectTabs,
-    val tabCollections: List<TabCollection> = listOf(),
-    val selectedTabCollection: TabCollection?
+    val tabCollections: List<TabCollection> = emptyList(),
+    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> =
-            { state, change ->
-                when (change) {
-                    is CollectionCreationChange.AddAllTabs -> state.copy(selectedTabs = state.tabs.toSet())
-                    is CollectionCreationChange.RemoveAllTabs -> state.copy(selectedTabs = setOf())
-                    is CollectionCreationChange.TabListChange -> state.copy(tabs = change.tabs)
-                    is CollectionCreationChange.TabAdded -> {
-                        val selectedTabs = state.selectedTabs + setOf(change.tab)
-                        state.copy(selectedTabs = selectedTabs)
-                    }
-                    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)
-                    }
-                }
+        val reducer: Reducer<CollectionCreationState, CollectionCreationChange> = { state, change ->
+            when (change) {
+                is CollectionCreationChange.AddAllTabs -> state.copy(selectedTabs = state.tabs.toSet())
+                is CollectionCreationChange.RemoveAllTabs -> state.copy(selectedTabs = emptySet())
+                is CollectionCreationChange.TabListChange -> state.copy(tabs = change.tabs)
+                is CollectionCreationChange.TabAdded -> state.copy(selectedTabs = state.selectedTabs + change.tab)
+                is CollectionCreationChange.TabRemoved -> state.copy(selectedTabs = state.selectedTabs - change.tab)
+                is CollectionCreationChange.StepChanged -> state.copy(saveCollectionStep = change.saveCollectionStep)
+                is CollectionCreationChange.CollectionSelected -> state.copy(selectedTabCollection = change.collection)
            }
+        }
    }
 }
--- a/app/src/main/java/org/mozilla/fenix/collections/CollectionCreationTabListAdapter.kt
+++ b/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 =
-                        if (checkChanged.shouldHideCheckBox) View.GONE else View.VISIBLE
+                    holder.itemView.tab_selected_checkbox.isGone = checkChanged.shouldHideCheckBox
                }
            }
        }
@ -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
--- a/app/src/main/java/org/mozilla/fenix/collections/CollectionCreationUIView.kt
+++ b/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, _ ->
-            if (actionId == EditorInfo.IME_ACTION_DONE && v.text.toString().isNotBlank()) {
+        view.name_collection_edittext.setOnEditorActionListener { view, actionId, _ ->
+            val text = view.text.toString()
+            if (actionId == EditorInfo.IME_ACTION_DONE && text.isNotBlank()) {
                when (step) {
-                    is SaveCollectionStep.NameCollection -> {
-                        actionEmitter.onNext(
-                            CollectionCreationAction.SaveCollectionName(
-                                selectedTabs.toList(),
-                                v.text.toString()
-                            )
-                        )
-                    }
-                    is SaveCollectionStep.RenameCollection -> {
-                        selectedCollection?.let {
-                            actionEmitter.onNext(CollectionCreationAction.RenameCollection(it, v.text.toString()))
-                        }
-                    }
+                    SaveCollectionStep.NameCollection ->
+                        CollectionCreationAction.SaveCollectionName(selectedTabs.toList(), text)
+                    SaveCollectionStep.RenameCollection ->
+                        selectedCollection?.let { CollectionCreationAction.RenameCollection(it, text) }
+                    else -> null
+                }?.let { action ->
+                    actionEmitter.onNext(action)
                }
            }
            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) {
-            when (step) {
-                SaveCollectionStep.SelectTabs -> {
-                    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
+        return if (event?.action == KeyEvent.ACTION_UP && keyCode == KeyEvent.KEYCODE_BACK) {
+            actionEmitter.onNext(CollectionCreationAction.BackPressed(step))
+            true
        } else {
-            return false
+            false
        }
    }

--- a/app/src/main/java/org/mozilla/fenix/collections/CreateCollectionFragment.kt
+++ b/app/src/main/java/org/mozilla/fenix/collections/CreateCollectionFragment.kt
@ -50,15 +50,7 @@ class CreateCollectionFragment : DialogFragment() {
                this,
                CollectionCreationViewModel::class.java
            ) {
-                CollectionCreationViewModel(
-                    CollectionCreationState(
-                        viewModel.tabs,
-                        viewModel.selectedTabs,
-                        viewModel.saveCollectionStep,
-                        viewModel.tabCollections,
-                        viewModel.selectedTabCollection
-                    )
-                )
+                CollectionCreationViewModel(viewModel.state)
            }
        )
        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) {
-            SaveCollectionStep.SelectTabs -> dismiss()
-            SaveCollectionStep.SelectCollection -> {
-                if (viewModel.tabs.size <= 1) dismiss() else {
-                    getManagedEmitter<CollectionCreationChange>().onNext(
-                        CollectionCreationChange.StepChanged(SaveCollectionStep.SelectTabs)
-                    )
-                }
+        val newStep = stepBack(backPressFrom)
+        if (newStep != null) {
+            getManagedEmitter<CollectionCreationChange>().onNext(CollectionCreationChange.StepChanged(newStep))
+        } else {
+            dismiss()
+        }
+    }
+
+    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 -> {
-                if (viewModel.tabCollections.isEmpty() && viewModel.tabs.size == 1) {
-                    dismiss()
-                } else {
-                    getManagedEmitter<CollectionCreationChange>()
-                        .onNext(
-                            CollectionCreationChange.StepChanged(
-                                viewModel.tabCollections.getBackStepForCollectionsSize()
-                            )
-                        )
-                }
-            }
-            SaveCollectionStep.RenameCollection -> {
-                dismiss()
+            SaveCollectionStep.NameCollection -> if (state.tabCollections.isEmpty()) {
+                stepBack(SaveCollectionStep.SelectCollection)
+            } else {
+                SaveCollectionStep.SelectCollection
            }
        }
    }

    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 }
-
-        tabs.forEach {
-            requireComponents.core.sessionManager.findSessionById(it.sessionId)?.let { session ->
-                requireComponents.useCases.tabsUseCases.removeTab.invoke(session)
-            }
+        if (viewModel.previousFragmentId == R.id.browserFragment) {
+            val components = requireComponents
+            tabs.asSequence()
+                .mapNotNull { tab -> components.core.sessionManager.findSessionById(tab.sessionId) }
+                .forEach { session -> components.useCases.tabsUseCases.removeTab(session) }
        }
    }
 }
--- a/app/src/main/java/org/mozilla/fenix/collections/CreateCollectionViewModel.kt
+++ b/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 tabs = listOf<Tab>()
-    var saveCollectionStep: SaveCollectionStep = SaveCollectionStep.SelectTabs
-    var tabCollections = listOf<TabCollection>()
-    var selectedTabCollection: TabCollection? = null
-    var snackbarAnchorView: View? = null
+    var state = CollectionCreationState()
+        private set
+
    var previousFragmentId: Int? = null

-    fun getStepForTabsAndCollectionSize(): SaveCollectionStep =
-        if (tabs.size > 1) SaveCollectionStep.SelectTabs else tabCollections.getStepForCollectionsSize()
+    fun updateCollection(
+        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
--- a/app/src/main/java/org/mozilla/fenix/collections/SaveCollectionListAdapter.kt
+++ b/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(this, selectedTabs.toList())
-                actionEmitter.onNext(action)
-            }
+            val action = CollectionCreationAction.SelectCollection(collection, selectedTabs.toList())
+            actionEmitter.onNext(action)
        }
    }

--- a/app/src/main/java/org/mozilla/fenix/components/toolbar/BrowserToolbarController.kt
+++ b/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 {
-                    viewModel.tabs = listOf(it)
-                    val selectedSet = mutableSetOf(it)
-                    viewModel.selectedTabs = selectedSet
-                    viewModel.tabCollections =
-                        context.components.core.tabCollectionStorage.cachedTabCollections.reversed()
-                    viewModel.saveCollectionStep = viewModel.tabCollections.getStepForCollectionsSize()
-                    viewModel.snackbarAnchorView = nestedScrollQuickActionView
+                currentSession?.toTab(context)?.let { currentSessionAsTab ->
+                    viewModel.saveTabToCollection(
+                        tabs = listOf(currentSessionAsTab),
+                        selectedTab = currentSessionAsTab,
+                        cachedTabCollections = context.components.core.tabCollectionStorage.cachedTabCollections
+                    )
                    viewModel.previousFragmentId = R.id.browserFragment

                    val directions = BrowserFragmentDirections.actionBrowserFragmentToCreateCollectionFragment()
--- a/app/src/main/java/org/mozilla/fenix/home/HomeFragment.kt
+++ b/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(
-                    selectedTabCollection = action.collection,
-                    step = SaveCollectionStep.SelectTabs
-                )
+                updateCollection(action.collection, SaveCollectionStep.SelectTabs)
            }
            is CollectionAction.Rename -> {
-                showCollectionCreationFragment(
-                    selectedTabCollection = action.collection,
-                    step = SaveCollectionStep.RenameCollection
-                )
+                updateCollection(action.collection, 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,
-        selectedTabCollection: TabCollection? = null,
-        step: SaveCollectionStep? = null
+        setupViewModel: (CreateCollectionViewModel, tabs: List<Tab>, cachedTabCollections: List<TabCollection>) -> Unit
    ) {
-        if (findNavController(this).currentDestination?.id == R.id.createCollectionFragment) return
-
-        val tabs = getListOfSessions().toTabs()
+        if (findNavController().currentDestination?.id == R.id.createCollectionFragment) return

        val viewModel: CreateCollectionViewModel by activityViewModels {
            ViewModelProvider.NewInstanceFactory() // this is a workaround for #4652
        }
-        viewModel.tabs = tabs
-        val selectedTabs =
-            tabs.find { tab -> tab.sessionId == selectedTabId }
-                ?: if (tabs.size == 1) tabs[0] else null
-        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()
+
+        val tabs = getListOfSessions().toTabs()
+        val cachedTabCollections = requireComponents.core.tabCollectionStorage.cachedTabCollections
+        setupViewModel(viewModel, tabs, cachedTabCollections)
+
        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(
--- a/app/src/test/java/org/mozilla/fenix/collections/CreateCollectionViewModelTest.kt
+++ b/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)
+    }
+}
--- a/app/src/test/java/org/mozilla/fenix/components/toolbar/DefaultBrowserToolbarControllerTest.kt
+++ b/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 { viewModel.selectedTabs = mutableSetOf(currentSessionAsTab) }
-        verify { viewModel.tabCollections = cachedTabCollections.reversed() }
-        verify { viewModel.saveCollectionStep = SaveCollectionStep.SelectCollection }
-        verify { viewModel.snackbarAnchorView = nestedScrollQuickActionView }
+        verify {
+            viewModel.saveTabToCollection(
+                listOf(currentSessionAsTab),
+                currentSessionAsTab,
+                cachedTabCollections
+            )
+        }
        verify { viewModel.previousFragmentId = R.id.browserFragment }
        verify {
            val directions = BrowserFragmentDirections