Dagger + ViewModel + SavedStateHandle

ApplicationComponent.kt

@Singleton
@Component(modules = [AssistedInjectionModule::class])
interface ApplicationComponent {
    fun mySavedStateViewModelFactory(): MySavedStateViewModel.Factory
}

AssistedInjectionModule.java

@AssistedModule
@Module(includes = {AssistedInject_AssistedInjectionModule.class})
public class AssistedInjectionModule {
}

MySavedStateFragment.kt

import androidx.lifecycle.observe

class MySavedStateFragment : Fragment(R.layout.my_saved_state_fragment) {
    private val viewModel by navGraphSavedStateViewModels(R.id.registration_graph) { handle ->
        Injector.get().mySavedStateViewModelFactory().create(handle)
    }
  
    override fun onViewCreated(view: View, savedInstanceState: Bundle?) {
        super.onViewCreated(view, savedInstanceState)
      
        val binding = MySavedStateFragmentBinding.bind(view)
      
        viewModel.blah.observe(viewLifecycleOwner) { blah ->
            // ...
        }
    }
}

MySavedStateViewModel.kt

class MySavedStateViewModel @AssistedInject constructor(
    private val authenticationManager: AuthenticationManager,
    @Assisted private val savedStateHandle: SavedStateHandle
) : ViewModel() {
    @AssistedInject.Factory
    interface Factory {
        fun create(savedStateHandle: SavedStateHandle): MySavedStateViewModel
    }
}

ViewModelUtils.kt

/**
 * Create a shared ViewModel instance across a given navGraphId.
 *
 * @receiver Fragment launches ViewModel
 * @param navGraphId Int defines ViewModel lifecycle
 * @param creator Function1<SavedStateHandle, T>
 * @return Lazy<T>
 */
inline fun <reified T : ViewModel> Fragment.navGraphSavedStateViewModels(
        @IdRes navGraphId: Int,
        crossinline creator: (SavedStateHandle) -> T
): Lazy<T> {
    // Wrapped in lazy to not search the NavController each time we want the backStackEntry
    val backStackEntry by lazy { findNavController().getBackStackEntry(navGraphId) }

    return createViewModelLazy(
            viewModelClass = T::class,
            storeProducer = { backStackEntry.viewModelStore },
            factoryProducer = {
                backStackEntry.createAbstractSavedStateViewModelFactory(
                        arguments = backStackEntry.arguments ?: Bundle(),
                        creator = creator
                )
            })
}

/**
 * Create a single instance of the ViewModel with Saved State enabled.
 *
 * @receiver Fragment launches ViewModel
 * @param creator Function1<SavedStateHandle, T>
 * @return Lazy<T>
 */
inline fun <reified T : ViewModel> Fragment.fragmentSavedStateViewModels(
        crossinline creator: (SavedStateHandle) -> T
): Lazy<T> {
    return createViewModelLazy(
            viewModelClass = T::class,
            storeProducer = { viewModelStore },
            factoryProducer = {
                createAbstractSavedStateViewModelFactory(
                        arguments = arguments ?: Bundle(),
                        creator = creator
                )
            })
}

/**
 * Factory function to create a single instance of the ViewModel.
 *
 * @receiver SavedStateRegistryOwner
 * @param arguments Bundle?
 * @param creator Function1<SavedStateHandle, T>
 * @return ViewModelProvider.Factory
 */
inline fun <reified T : ViewModel> SavedStateRegistryOwner.createAbstractSavedStateViewModelFactory(
        arguments: Bundle,
        crossinline creator: (SavedStateHandle) -> T
): ViewModelProvider.Factory {
    return object : AbstractSavedStateViewModelFactory(this, arguments) {
        @Suppress("UNCHECKED_CAST")
        override fun <T : ViewModel?> create(
                key: String,
                modelClass: Class<T>,
                handle: SavedStateHandle
        ): T = creator(handle) as T
    }
}

/**
 * Configure CoroutineScope injection for production and testing.
 *
 * @receiver ViewModel provides viewModelScope for production
 * @param coroutineScope null for production, injects TestCoroutineScope for unit tests
 * @return CoroutineScope to launch coroutines on
 */
fun ViewModel.getViewModelScope(coroutineScope: CoroutineScope?) =
        if (coroutineScope == null) this.viewModelScope
        else coroutineScope

/**
 * Configure Dispatchers injection for production and testing.
 */
interface DispatcherProvider {
    fun main(): CoroutineDispatcher = Dispatchers.Main
    fun default(): CoroutineDispatcher = Dispatchers.Default
    fun io(): CoroutineDispatcher = Dispatchers.IO
    fun unconfined(): CoroutineDispatcher = Dispatchers.Unconfined
}

class DefaultDispatcherProvider : DispatcherProvider