Magnet Scoping / ViewModel injection

BaseMagnetActivity.kt

open class BaseMagnetActivity : AppCompatActivity(), ScopeOwner {

    private val scopeModel: ScopeModel by lazy { ScopeModel.setup(this) }
    private var testScope: Scope? = null

    override var scope: Scope
        get() = testScope ?: scopeModel.scope
        @VisibleForTesting
        set(value) {
            testScope = value
        }
}

BaseMagnetApplication.kt

const val APPLICATION = "application"

open class BaseMagnetApplication : Application(), ScopeOwner {

    @Suppress("UnsafeCast")
    override val scope: Scope by lazy {
        Magnet.createRootScope().apply {
            bind(this@BaseMagnetApplication as Application)
            bind(this@BaseMagnetApplication as Context, APPLICATION)
        }
    }
}

BaseMagnetFragment.kt

open class BaseMagnetFragment : Fragment(), ScopeOwner {

    private val scopeModel: ScopeModel by lazy { ScopeModel.setup(this) }
    private var testScope: Scope? = null

    override var scope: Scope
        get() = testScope ?: scopeModel.scope
        @VisibleForTesting
        set(value) {
            testScope = value
        }
}

LifecycleOwnerAware.kt

@Instance(type = ActivityAware::class, scoping = Scoping.DIRECT)
class ActivityAware : LifecycleOwnerAware<FragmentActivity>()

@Instance(type = FragmentAware::class, scoping = Scoping.DIRECT)
class FragmentAware : LifecycleOwnerAware<Fragment>()

open class LifecycleOwnerAware<T : LifecycleOwner> {
    private val subject = BehaviorSubject.createDefault<Instance>(Instance.Absent)

    @Suppress("UNCHECKED_CAST")
    fun <R> with(work: T.() -> R): R =
        when (val value = subject.value) {
            null,
            is Instance.Absent -> error("no instance present")
            is Instance.Exists<*> -> (value.instance as T).run(work)
        }

    @Suppress("UNCHECKED_CAST")
    fun observe(): Observable<T> =
        subject
            .filter { it is Instance.Exists<*> }
            .map { (it as Instance.Exists<*>).instance as T }
            .hide()

    internal fun setInstance(instance: T) {
        subject.onNext(Instance.Exists(instance))
    }

    internal fun clearInstance() {
        subject.onNext(Instance.Absent)
    }

    private sealed class Instance {
        class Exists<T>(val instance: T) : Instance()
        object Absent : Instance()
    }
}

MyExampleFragment.kt

class MyExampleFragment : BaseFragment() {

    private lateinit var viewModel: MyViewModel
  
    override fun onCreate(savedInstanceState: Bundle?) {
        super.onCreate(savedInstanceState)

        viewModel = scope.getSingle()
    }
}

MyExampleFragmentTest.kt

class MyExampleFragmentTest {

    @Mock
    private lateinit var myViewModel: MyViewModel

    private val stateSubject = PublishSubject.create<MyState>()
    private val effectSubject = PublishSubject.create<MyEffect>()

    private val testScope = Magnet.createRootScope()

    @Before
    fun setup() {
        MockitoAnnotations.initMocks(this)
        testScope.apply {
            bind(MyViewModel::class.java, myViewModel)
        }

        myViewModel.stub {
            on { observeState() } doReturn stateSubject
            on { observeEffects() } doReturn effectSubject
        }
    }

    @Test
    fun shouldRenderScreen() {
        launchFragment()

        stateSubject.onNext(SOME_FIXTURE_STATE)

        // verify now with Espresso
    }

    // needs JetPack's fragment-testing artifact
    private fun launchFragment(): FragmentScenario<MyExampleFragment> =
        launchFragmentInContainer {
            MyExampleFragment().apply {
                this.scope = testScope
            }
        }

    companion object {
        private val SOME_FIXTURE_STATE = MyState(...)
    }
}

MyViewModel.kt

sealed class MyState { ... }
sealed class MyEffect { ... }

@Instance(
    type = MyViewModel::class,
    factory = ViewModelFactory::class,
    scoping = Scoping.UNSCOPED
)
class MyViewModel : ViewModel() {
  
  fun observeState(): Observable<MyState> = ...
  
  fun observeEffects(): Observable<MyEffect> = ...
  
}

ScopeModel.kt

class ScopeModel(context: Context) : ViewModel(), DefaultLifecycleObserver, ScopeOwner {

    override val scope: Scope by lazy {
        context.createSubscope {}
    }

    override fun onCreate(owner: LifecycleOwner) {
        scope.apply {
            if (owner is FragmentActivity) {
                getSingle(ActivityAware::class.java).setInstance(owner)
            }
            if (owner is Fragment) {
                getSingle(FragmentAware::class.java).setInstance(owner)
            }
        }
    }

    override fun onDestroy(owner: LifecycleOwner) {
        scope.apply {
            if (owner is FragmentActivity) {
                getSingle(ActivityAware::class.java).clearInstance()
            }
            if (owner is Fragment) {
                getSingle(FragmentAware::class.java).clearInstance()
            }
        }
    }

    @Suppress("UNCHECKED_CAST")
    companion object {
        fun setup(activity: FragmentActivity): ScopeModel {
            val model = ViewModelProviders.of(activity, object : ViewModelProvider.Factory {
                override fun <T : ViewModel?> create(modelClass: Class<T>): T =
                    ScopeModel(activity.application) as T
            }).get(ScopeModel::class.java)
            activity.lifecycle.addObserver(model)
            return model
        }

        fun setup(fragment: Fragment): ScopeModel {
            val model = ViewModelProviders.of(fragment, object : ViewModelProvider.Factory {
                override fun <T : ViewModel?> create(modelClass: Class<T>): T =
                    ScopeModel(fragment.requireActivity()) as T
            }).get(ScopeModel::class.java)
            fragment.lifecycle.addObserver(model)
            return model
        }
    }
}

ScopeOwner.kt

interface ScopeOwner {
    val scope: Scope
}

fun Context?.createSubscope(initializer: Scope.() -> Unit): Scope {
    val parentScopeOwner = this as? ScopeOwner ?: error("$this is expected to implement ScopeOwner")
    return parentScopeOwner.scope.createSubscope { initializer(this) }
}

ViewModelFactory.kt

const val VIEW_MODEL_KEY = "view-model-key"

class ViewModelFactory<T : ViewModel> : Factory<T> {
    override fun create(
        scope: Scope,
        type: Class<T>,
        classifier: String,
        scoping: Scoping,
        instantiator: Factory.Instantiator<T>
    ): T {
        val key = scope.getOptional(String::class.java, VIEW_MODEL_KEY)
        if (key != null && scoping != Scoping.UNSCOPED) {
            error("ViewModel '$type' with key '$key' must be declared with Scoping.UNSCOPED")
        }
        val androidViewModelFactory = AndroidViewModelFactory(scope, instantiator)
        val provider = getProviderFor(scope, androidViewModelFactory)
        return if (key == null) provider.get(type) else provider.get(key, type)
    }

    private fun getProviderFor(scope: Scope, viewModelProviderFactory: ViewModelProvider.Factory): ViewModelProvider {
        // fragment's have precedence over activities in the given scope
        val fragmentAware = scope.getOptional(FragmentAware::class.java)
        if (fragmentAware != null) {
            return fragmentAware.with {
                ViewModelProviders.of(this, viewModelProviderFactory)
            }
        }
        val activityAware = scope.getSingle(ActivityAware::class.java)
        return activityAware.with {
            ViewModelProviders.of(this, viewModelProviderFactory)
        }
    }

    private class AndroidViewModelFactory<T>(
        private var scope: Scope,
        private var instantiator: Factory.Instantiator<T>
    ) : ViewModelProvider.Factory {
        @Suppress("UNCHECKED_CAST")
        override fun <T : ViewModel?> create(modelClass: Class<T>): T =
            instantiator.instantiate(scope) as T
    }
}