1 /*
2  * Copyright 2023 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 androidx.datastore.testapp.multiprocess
18 
19 import androidx.datastore.testapp.twoWayIpc.TwoWayIpcConnection
20 import androidx.datastore.testapp.twoWayIpc.TwoWayIpcService
21 import androidx.datastore.testapp.twoWayIpc.TwoWayIpcService2
22 import androidx.test.platform.app.InstrumentationRegistry
23 import java.util.concurrent.atomic.AtomicBoolean
24 import kotlin.time.Duration.Companion.seconds
25 import kotlinx.coroutines.CoroutineScope
26 import kotlinx.coroutines.Dispatchers
27 import kotlinx.coroutines.Job
28 import kotlinx.coroutines.async
29 import kotlinx.coroutines.awaitAll
30 import kotlinx.coroutines.cancel
31 import kotlinx.coroutines.runBlocking
32 import kotlinx.coroutines.sync.Mutex
33 import kotlinx.coroutines.sync.withLock
34 import kotlinx.coroutines.withTimeout
35 import org.junit.rules.TestWatcher
36 import org.junit.runner.Description
37 
38 /**
39  * Used for testing multi-process cases while also maintaining resources so that services are
40  * properly closed after test.
41  */
42 class MultiProcessTestRule : TestWatcher() {
43     private val didRunTest = AtomicBoolean(false)
44     private val context = InstrumentationRegistry.getInstrumentation().context
45 
46     // use a real scope, it is too hard to use a TestScope when we cannot control the IPC
47     val datastoreScope = CoroutineScope(Dispatchers.IO + Job())
48     private val connectionsMutex = Mutex()
49     private val connections = mutableListOf<TwoWayIpcConnection>()
50     private val availableServiceClasses =
51         mutableListOf<Class<out TwoWayIpcService>>(
52             TwoWayIpcService::class.java,
53             TwoWayIpcService2::class.java
54         )
55 
runTestnull56     fun runTest(block: suspend CoroutineScope.() -> Unit) {
57         // don't use datastore scope here as it will not finish by itself.
58         runBlocking {
59             check(didRunTest.compareAndSet(false, true)) { "Cannot call runTest multiple times" }
60             try {
61                 withTimeout(TEST_TIMEOUT) { block() }
62             } finally {
63                 connections.map { async { it.disconnect() } }.awaitAll()
64             }
65         }
66     }
67 
createConnectionnull68     suspend fun createConnection(): TwoWayIpcConnection {
69         val connection =
70             connectionsMutex.withLock {
71                 val klass =
72                     availableServiceClasses.removeFirstOrNull()
73                         ?: error(
74                             "Cannot create more services," +
75                                 "you can declare more in the manifest if needed"
76                         )
77                 TwoWayIpcConnection(context, klass).also { connections.add(it) }
78             }
79         connection.connect()
80         return connection
81     }
82 
finishednull83     override fun finished(description: Description) {
84         super.finished(description)
85         datastoreScope.cancel()
86     }
87 
88     companion object {
89         val TEST_TIMEOUT = 10.seconds
90     }
91 }
92