ActivityLifecycleTestRule.kt

/*
 * Copyright 2023 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.example.android.architecture.blueprints.todoapp

import androidx.activity.ComponentActivity
import androidx.test.runner.lifecycle.ActivityLifecycleCallback
import androidx.test.runner.lifecycle.ActivityLifecycleMonitorRegistry
import androidx.test.runner.lifecycle.Stage
import org.junit.rules.TestRule
import org.junit.runner.Description
import org.junit.runners.model.Statement

/**
 * A JUnit [TestRule] that monitors the lifecycle of an activity and executes a given function
 * when an instance of the activity's class is in the CREATED stage.
 *
 * This rule can be used to set the composable content of an activity and to make sure that it's
 * set again after activity recreation. It doesn't replace `ComposeTestRule`.
 *
 * Note that this rule might not work with multiple instances of the same activity class running.
 * 
 * Example usage: 
 * ```
 * @get:Rule
 * val monitorRule = createActivityLifecycleTestRule<ComponentActivity> { activity ->
 *     val viewModel = ...
 *     activity.setContent {
 *         MyScreen(
 *             viewModel = viewModel,
 *         )
 *     }
 * }
 *
 * @get:Rule
 * val composeTestRule = createAndroidComposeRule<ComponentActivity>()
 * 
 * private val activity: ComponentActivity
 *     get() = monitorRule.activity ?: throw IllegalStateException("Activity was never created")
 * ```
 */
class ActivityLifecycleTestRule<A>(
    private val activityClass: Class<A>,
    private val onCreate: (activity: A) -> Unit
) : TestRule {

    var activity: A? = null

    override fun apply(base: Statement, description: Description): Statement {
        return object : Statement() {
            override fun evaluate() {
                val callback = ActivityLifecycleCallback { activity, stage ->
                    if (
                        stage == Stage.CREATED
                        && activityClass.isInstance(activity) 
                    ) {
                        val newActivity = activityClass.cast(activity) as A
                        this@ActivityLifecycleTestRule.activity = newActivity
                        onCreate(newActivity)
                    }
                }

                val monitor = ActivityLifecycleMonitorRegistry.getInstance()
                monitor.addLifecycleCallback(callback)
                try {
                    base.evaluate()
                } finally {
                    monitor.removeLifecycleCallback(callback)
                }
            }
        }
    }
}

inline fun <reified A : ComponentActivity> createActivityLifecycleTestRule(
    noinline onCreate: (activity: A) -> Unit
) = ActivityLifecycleTestRule(A::class.java, onCreate)

awesome thanks @JoseAlcerreca