For #3633 - Update StoreProvider to use a callback

Co-authored-by: Christian Sadilek <christian.sadilek@gmail.com>

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

@@ -9,25 +9,37 @@ import androidx.lifecycle.ViewModel
 import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.ViewModelProviders
 import androidx.lifecycle.get
-import mozilla.components.lib.state.Action
-import mozilla.components.lib.state.State
 import mozilla.components.lib.state.Store
 
 /**
- * Generic ViewModel to wrap a State object for state restoration
+ * Generic ViewModel to wrap a State object for state restoration.
+ *
+ * @property store [Store] instance attached to [ViewModel].
  */
-@Suppress("UNCHECKED_CAST")
-class StoreProvider<S : State, A : Action, T : Store<S, A>>(val store: T) : ViewModel() {
-    companion object {
-        fun <S : State, A : Action, T : Store<S, A>> get(fragment: Fragment, initialStore: T): T {
-            val factory = object : ViewModelProvider.Factory {
-                override fun <VM : ViewModel?> create(modelClass: Class<VM>): VM {
-                    return StoreProvider(initialStore) as VM
-                }
-            }
+class StoreProvider<T : Store<*, *>>(
+    val store: T
+) : ViewModel() {
 
-            val viewModel: StoreProvider<S, A, T> = ViewModelProviders.of(fragment, factory).get()
+    companion object {
+        fun <T : Store<*, *>> get(fragment: Fragment, createStore: () -> T): T {
+            val factory = StoreProviderFactory(createStore)
+            val viewModel: StoreProvider<T> = ViewModelProviders.of(fragment, factory).get()
             return viewModel.store
         }
     }
 }
+
+/**
+ * ViewModel factory to create [StoreProvider] instances.
+ *
+ * @param createStore Callback to create a new [Store], used when the [ViewModel] is first created.
+ */
+class StoreProviderFactory<T : Store<*, *>>(
+    private val createStore: () -> T
+) : ViewModelProvider.Factory {
+
+    @Suppress("UNCHECKED_CAST")
+    override fun <VM : ViewModel?> create(modelClass: Class<VM>): VM {
+        return StoreProvider(createStore()) as VM
+    }
+}

app/src/main/java/org/mozilla/fenix/library/history/HistoryFragment.kt

@@ -55,14 +55,13 @@ class HistoryFragment : Fragment(), BackHandler {
         savedInstanceState: Bundle?
     ): View? {
         val view = inflater.inflate(R.layout.fragment_history, container, false)
-        historyStore = StoreProvider.get(
-            this,
+        historyStore = StoreProvider.get(this) {
             HistoryStore(
                 HistoryState(
                     items = listOf(), mode = HistoryState.Mode.Normal
                 )
             )
-        )
+        }
         historyInteractor = HistoryInteractor(
             historyStore,
             ::openItem,

app/src/main/java/org/mozilla/fenix/search/SearchFragment.kt

@@ -72,20 +72,20 @@ class SearchFragment : Fragment(), BackHandler {
         val view = inflater.inflate(R.layout.fragment_search, container, false)
         val url = session?.url ?: ""
 
-        searchStore = StoreProvider.get(
-            this,
+        searchStore = StoreProvider.get(this) {
             SearchStore(
                 SearchState(
-                query = url,
-                showShortcutEnginePicker = false,
-                searchEngineSource = SearchEngineSource.Default(
-                    requireComponents.search.searchEngineManager.getDefaultSearchEngine(requireContext())
-                ),
-                showSuggestions = Settings.getInstance(requireContext()).showSearchSuggestions,
-                showVisitedSitesBookmarks = Settings.getInstance(requireContext()).shouldShowVisitedSitesBookmarks,
-                session = session)
+                    query = url,
+                    showShortcutEnginePicker = false,
+                    searchEngineSource = SearchEngineSource.Default(
+                        requireComponents.search.searchEngineManager.getDefaultSearchEngine(requireContext())
+                    ),
+                    showSuggestions = Settings.getInstance(requireContext()).showSearchSuggestions,
+                    showVisitedSitesBookmarks = Settings.getInstance(requireContext()).shouldShowVisitedSitesBookmarks,
+                    session = session
+                )
             )
-        )
+        }
 
         searchInteractor = SearchInteractor(
             activity as HomeActivity,

app/src/test/java/org/mozilla/fenix/components/StoreProviderTest.kt

@@ -0,0 +1,73 @@
+/* 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.components
+
+import androidx.fragment.app.Fragment
+import androidx.fragment.app.testing.launchFragmentInContainer
+import mozilla.components.lib.state.Action
+import mozilla.components.lib.state.State
+import mozilla.components.lib.state.Store
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertFalse
+import org.junit.Assert.assertTrue
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mozilla.fenix.TestApplication
+import org.robolectric.RobolectricTestRunner
+import org.robolectric.annotation.Config
+
+@RunWith(RobolectricTestRunner::class)
+@Config(application = TestApplication::class)
+class StoreProviderTest {
+
+    private class BasicState : State
+
+    private val basicStore = Store(BasicState()) { state, _: Action -> state }
+
+    @Test
+    fun `factory returns store provider`() {
+        var createCalled = false
+        val factory = StoreProviderFactory {
+            createCalled = true
+            basicStore
+        }
+
+        assertFalse(createCalled)
+
+        assertEquals(basicStore, factory.create(StoreProvider::class.java).store)
+
+        assertTrue(createCalled)
+    }
+
+    @Test
+    fun `get returns store`() {
+        val scenario = launchFragmentInContainer { Fragment() }
+        scenario.onFragment {
+            val store = StoreProvider.get(it) { basicStore }
+            assertEquals(basicStore, store)
+        }
+    }
+
+    @Test
+    fun `get only calls createStore if needed`() {
+        val scenario = launchFragmentInContainer { Fragment() }
+        var createCalled = false
+        val createStore = {
+            createCalled = true
+            basicStore
+        }
+
+        scenario.onFragment {
+            StoreProvider.get(it, createStore)
+        }
+        assertTrue(createCalled)
+
+        createCalled = false
+        scenario.onFragment {
+            StoreProvider.get(it, createStore)
+        }
+        assertFalse(createCalled)
+    }
+}