contract.kt

@UseExperimental(ExperimentalContracts::class)
fun <T : Activity> ActivityScenario<T>.onActivity2(block: (T) -> Unit) {
  contract {
    callsInPlace(block, InvocationKind.EXACTLY_ONCE)
  }
  onActivity {
    block(it)
  }
}

test.kt

import kotlin.contracts.ExperimentalContracts
import kotlin.contracts.InvocationKind
import kotlin.contracts.contract

@RunWith(AndroidJUnit4::class)
internal class OnActivityContractTest {
  @Test
  fun onActivityContract() {
    val scenario = ActivityScenario.launch(Activity::class.java)
    val activity: Activity
    scenario.onActivity2 {
      activity = it
    }
    println(activity)
  }
}