• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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