1 /* 2 * Copyright (C) 2024 The Android Open Source Project 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 package com.android.test.tracing.coroutines 18 19 import android.platform.test.annotations.EnableFlags 20 import com.android.app.tracing.coroutines.createCoroutineTracingContext 21 import com.android.app.tracing.coroutines.flow.stateInTraced 22 import com.android.app.tracing.coroutines.launchTraced 23 import com.android.systemui.Flags.FLAG_COROUTINE_TRACING 24 import java.util.concurrent.Executor 25 import kotlinx.coroutines.CoroutineDispatcher 26 import kotlinx.coroutines.CoroutineScope 27 import kotlinx.coroutines.asExecutor 28 import kotlinx.coroutines.cancel 29 import kotlinx.coroutines.channels.awaitClose 30 import kotlinx.coroutines.delay 31 import kotlinx.coroutines.flow.Flow 32 import kotlinx.coroutines.flow.MutableStateFlow 33 import kotlinx.coroutines.flow.SharingStarted 34 import kotlinx.coroutines.flow.StateFlow 35 import kotlinx.coroutines.flow.callbackFlow 36 import kotlinx.coroutines.flow.combine 37 import kotlinx.coroutines.flow.distinctUntilChanged 38 import kotlinx.coroutines.flow.map 39 import kotlinx.coroutines.flow.onEach 40 import kotlinx.coroutines.flow.onStart 41 import kotlinx.coroutines.job 42 import org.junit.Test 43 44 data class ExampleInfo(val a: Int, val b: Boolean, val c: String) 45 46 interface ExampleStateTracker { 47 val info: ExampleInfo 48 addCallbacknull49 fun addCallback(callback: Callback, executor: Executor) 50 51 fun removeCallback(callback: Callback) 52 53 interface Callback { 54 fun onInfoChanged(newInfo: ExampleInfo) 55 } 56 } 57 58 interface ExampleRepository { 59 val currentInfo: Flow<ExampleInfo> 60 val otherState: StateFlow<Boolean> 61 val combinedState: StateFlow<Boolean> // true when otherState == true and current.b == true 62 } 63 64 class ExampleStateTrackerImpl : ExampleStateTracker { 65 private var _info = ExampleInfo(0, false, "Initial") 66 override val info: ExampleInfo 67 get() = _info 68 69 val callbacks = mutableListOf<Pair<ExampleStateTracker.Callback, Executor>>() 70 addCallbacknull71 override fun addCallback(callback: ExampleStateTracker.Callback, executor: Executor) { 72 callbacks.add(Pair(callback, executor)) 73 } 74 removeCallbacknull75 override fun removeCallback(callback: ExampleStateTracker.Callback) { 76 callbacks.removeIf { it.first == callback } 77 } 78 forceUpdatenull79 fun forceUpdate(a: Int, b: Boolean, c: String) { 80 _info = ExampleInfo(a, b, c) 81 callbacks.forEach { it.second.execute { it.first.onInfoChanged(_info) } } 82 } 83 } 84 85 private class ExampleRepositoryImpl( 86 private val testBase: TestBase, 87 private val bgScope: CoroutineScope, 88 private val tracker: ExampleStateTrackerImpl, 89 ) : ExampleRepository { 90 @OptIn(ExperimentalStdlibApi::class) 91 override val currentInfo: StateFlow<ExampleInfo> = <lambda>null92 callbackFlow { 93 channel.trySend(tracker.info) 94 val callback = 95 object : ExampleStateTracker.Callback { 96 override fun onInfoChanged(newInfo: ExampleInfo) { 97 channel.trySend(newInfo) 98 } 99 } 100 tracker.addCallback( 101 callback, 102 bgScope.coroutineContext[CoroutineDispatcher]!!.asExecutor(), 103 ) 104 awaitClose { tracker.removeCallback(callback) } 105 } <lambda>null106 .onEach { testBase.expect("1^currentInfo") } 107 .stateInTraced( 108 "currentInfo", 109 bgScope, 110 SharingStarted.Eagerly, 111 initialValue = tracker.info, 112 ) 113 114 override val otherState = MutableStateFlow(false) 115 116 /** flow that emits true only when currentInfo.b == true and otherState == true */ 117 override val combinedState: StateFlow<Boolean> 118 get() = 119 combine(currentInfo, otherState, ::Pair) <lambda>null120 .map { it.first.b && it.second } 121 .distinctUntilChanged() <lambda>null122 .onEach { testBase.expect("2^combinedState:1^:2^") } <lambda>null123 .onStart { emit(false) } 124 .stateInTraced( 125 "combinedState", 126 scope = bgScope, 127 started = SharingStarted.WhileSubscribed(), 128 initialValue = false, 129 ) 130 } 131 132 @EnableFlags(FLAG_COROUTINE_TRACING) 133 class CallbackFlowTracingTest : TestBase() { 134 <lambda>null135 private val bgScope: CoroutineScope by lazy { 136 CoroutineScope( 137 createCoroutineTracingContext("bg", testMode = true) + 138 bgThread1 + 139 scope.coroutineContext.job 140 ) 141 } 142 143 @Test callbackFlownull144 fun callbackFlow() { 145 val exampleTracker = ExampleStateTrackerImpl() 146 val repository = ExampleRepositoryImpl(this, bgScope, exampleTracker) 147 runTest(totalEvents = 15) { 148 launchTraced("collectCombined") { 149 // upstream flow already has tracing, so tracing with a collect call here would be 150 // redundant. That's why we call `collect` instead of `collectTraced` 151 repository.combinedState.collect { 152 expect( 153 "1^main:1^collectCombined", 154 "collect:combinedState", 155 "emit:combinedState", 156 ) 157 } 158 } 159 delay(10) 160 expect("1^main") 161 delay(10) 162 exampleTracker.forceUpdate(1, false, "A") // <-- no change 163 delay(10) 164 repository.otherState.value = true // <-- no change 165 delay(10) 166 exampleTracker.forceUpdate(2, true, "B") // <-- should update `combinedState` 167 delay(10) 168 repository.otherState.value = false // <-- should update `combinedState` 169 delay(10) 170 exampleTracker.forceUpdate(3, false, "C") // <-- no change 171 delay(10) 172 exampleTracker.forceUpdate(4, true, "D") // <-- no change 173 delay(10) 174 repository.otherState.value = true // <-- should update `combinedState` 175 delay(10) 176 expect("1^main") 177 cancel("Cancelled normally for test") 178 } 179 bgScope.cancel("Cancelled normally for test") 180 } 181 } 182