KotlinMultiplatform Desktop Test Issue

FlowExt.android.kt

@Composable
actual fun <T> Flow<T>.rememberFlowWithLifecycle(
    lifecycleOwner: LifecycleOwner
): Flow<T> {
    val thisValue = this
    return remember(thisValue, lifecycleOwner) {
        thisValue.flowWithLifecycle(
            lifecycleOwner.lifecycle,
            Lifecycle.State.STARTED
        )
    }
}

FlowExt.desktop.kt

@Composable
actual fun <T> Flow<T>.rememberFlowWithLifecycle(lifecycleOwner: LifecycleOwner): Flow<T> {
    val thisValue = this
    return remember(thisValue, lifecycleOwner) {
        thisValue.flowWithLifecycle(
            lifecycleOwner.lifecycle,
            Lifecycle.State.STARTED
        )
    }
}


// implementation copied from Android implementation which is not yet available in KMP
private fun <T> Flow<T>.flowWithLifecycle(
    lifecycle: Lifecycle,
    minActiveState: Lifecycle.State = Lifecycle.State.STARTED
): Flow<T> =
    callbackFlow {
        lifecycle.repeatOnLifecycle(minActiveState) {
            this@flowWithLifecycle.collect {
                send(it)
            }
        }
        close()
    }

// implementation copied from Android implementation which is not yet available in KMP
@Suppress("LabeledExpression", "CognitiveComplexMethod", "CastNullableToNonNullableType")
private suspend fun Lifecycle.repeatOnLifecycle(
    state: Lifecycle.State,
    block: suspend CoroutineScope.() -> Unit
) {
    require(state !== Lifecycle.State.INITIALIZED) {
        "repeatOnLifecycle cannot start work with the INITIALIZED lifecycle state."
    }

    if (currentState === Lifecycle.State.DESTROYED) {
        return
    }

    // This scope is required to preserve context before we move to Dispatchers.Main
    coroutineScope {
        withContext(Dispatchers.Main.immediate) {
            // Check the current state of the lifecycle as the previous check is not guaranteed
            // to be done on the main thread.
            if (currentState === Lifecycle.State.DESTROYED) return@withContext

            // Instance of the running repeating coroutine
            var launchedJob: Job? = null

            // Registered observer
            var observer: LifecycleEventObserver? = null
            try {
                // Suspend the coroutine until the lifecycle is destroyed or
                // the coroutine is cancelled
                suspendCancellableCoroutine<Unit> { cont ->
                    // Lifecycle observers that executes `block` when the lifecycle reaches certain state, and
                    // cancels when it falls below that state.
                    val startWorkEvent = Lifecycle.Event.upTo(state)
                    val cancelWorkEvent = Lifecycle.Event.downFrom(state)
                    val mutex = Mutex()
                    observer =
                        LifecycleEventObserver { _, event ->
                            if (event == startWorkEvent) {
                                // Launch the repeating work preserving the calling context
                                launchedJob =
                                    this@coroutineScope.launch {
                                        // Mutex makes invocations run serially,
                                        // coroutineScope ensures all child coroutines finish
                                        mutex.withLock {
                                            coroutineScope {
                                                block()
                                            }
                                        }
                                    }
                                return@LifecycleEventObserver
                            }
                            if (event == cancelWorkEvent) {
                                launchedJob?.cancel()
                                launchedJob = null
                            }
                            if (event == Lifecycle.Event.ON_DESTROY) {
                                cont.resume(Unit)
                            }
                        }
                    this@repeatOnLifecycle.addObserver(observer as LifecycleEventObserver)
                }
            } finally {
                launchedJob?.cancel()
                observer?.let {
                    this@repeatOnLifecycle.removeObserver(it)
                }
            }
        }
    }
}

FlowExt.kt

@Composable
fun <T> Flow<T>.collectInLaunchedEffect(
    function: suspend (value: T) -> Unit
) {
    val effectFlow = rememberFlowWithLifecycle()
    LaunchedEffect(effectFlow) {
        effectFlow.collect(function)
    }
}

@Composable
expect fun <T> Flow<T>.rememberFlowWithLifecycle(
    lifecycleOwner: LifecycleOwner = LocalLifecycleOwner.current
): Flow<T>

Screen.kt

@Composable
internal fun Screen(
    onCloseApp: () -> Unit,
    viewModel: ScreenViewModel = koinViewModel<ScreenViewModel>()
) {
    viewModel.effects.collectInLaunchedEffect { effect ->
        when (effect) {
            ScreenEffect.CloseApp -> onCloseApp()
        }
    }
    
    Button(text = "Click me", onClick = { viewModel.onClickButton() })
}

ScreenEffect.kt

sealed class ScreenEffect {
    data object CloseApp : ScreenEffect()
}

ScreenTest.kt

class ScreenTest {
    private val onCloseApp: () -> Unit = mockk(relaxed = true)

    @Test
    fun `should close app on click`() {
        runComposeUiTest {
            setContent {
                Screen(onCloseApp = onCloseApp)
            }

            onNodeWithText("Close App").performClick()

            // this is flaky on Desktop, but always succeeds on Android
            verify { onCloseApp() }
        }
    }
}

ScreenViewModel.kt

class ScreenViewModel : ViewModel() {
    private val _effects = Channel<ScreenEffect>()
    val effects: Flow<ScreenEffect>
        get() = _effects.receiveAsFlow()

    protected fun runEffect(effect: ScreenEffect) {
        viewModelScope.launch {
            _effects.send(effect)
        }
    }
    
    fun onClickButton() {
        runEffect(ScreenEffect.CloseApp)
    }
    
    // more things here
}