1 /*
<lambda>null2  * Copyright (C) 2022 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.testutils
18 
19 import android.Manifest.permission.READ_DEVICE_CONFIG
20 import android.Manifest.permission.WRITE_DEVICE_CONFIG
21 import android.Manifest.permission.WRITE_ALLOWLISTED_DEVICE_CONFIG
22 import android.provider.DeviceConfig
23 import android.util.Log
24 import com.android.modules.utils.build.SdkLevel
25 import com.android.testutils.FunctionalUtils.ThrowingRunnable
26 import java.util.concurrent.CompletableFuture
27 import java.util.concurrent.Executor
28 import java.util.concurrent.TimeUnit
29 import org.junit.rules.TestRule
30 import org.junit.runner.Description
31 import org.junit.runners.model.Statement
32 
33 private val TAG = DeviceConfigRule::class.simpleName
34 
35 private const val TIMEOUT_MS = 20_000L
36 
37 /**
38  * A [TestRule] that helps set [DeviceConfig] for tests and clean up the test configuration
39  * automatically on teardown.
40  *
41  * The rule can also optionally retry tests when they fail following an external change of
42  * DeviceConfig before S; this typically happens because device config flags are synced while the
43  * test is running, and DisableConfigSyncTargetPreparer is only usable starting from S.
44  *
45  * @param retryCountBeforeSIfConfigChanged if > 0, when the test fails before S, check if
46  *        the configs that were set through this rule were changed, and retry the test
47  *        up to the specified number of times if yes.
48  */
49 class DeviceConfigRule @JvmOverloads constructor(
50     val retryCountBeforeSIfConfigChanged: Int = 0
51 ) : TestRule {
52     // Maps (namespace, key) -> value
53     private val originalConfig = mutableMapOf<Pair<String, String>, String?>()
54     private val usedConfig = mutableMapOf<Pair<String, String>, String?>()
55 
56     /**
57      * Actions to be run after cleanup of the config, for the current test only.
58      */
59     private val currentTestCleanupActions = mutableListOf<ThrowingRunnable>()
60 
61     override fun apply(base: Statement, description: Description): Statement {
62         return TestValidationUrlStatement(base, description)
63     }
64 
65     private inner class TestValidationUrlStatement(
66         private val base: Statement,
67         private val description: Description
68     ) : Statement() {
69         override fun evaluate() {
70             var retryCount = if (SdkLevel.isAtLeastS()) 1 else retryCountBeforeSIfConfigChanged + 1
71             while (retryCount > 0) {
72                 retryCount--
73                 tryTest {
74                     base.evaluate()
75                     // Can't use break/return out of a loop here because this is a tryTest lambda,
76                     // so set retryCount to exit instead
77                     retryCount = 0
78                 }.catch<Throwable> { e -> // junit AssertionFailedError does not extend Exception
79                     if (retryCount == 0) throw e
80                     usedConfig.forEach { (key, value) ->
81                         val currentValue = runAsShell(READ_DEVICE_CONFIG) {
82                             DeviceConfig.getProperty(key.first, key.second)
83                         }
84                         if (currentValue != value) {
85                             Log.w(TAG, "Test failed with unexpected device config change, retrying")
86                             return@catch
87                         }
88                     }
89                     throw e
90                 } cleanupStep {
91                     runAsShell(WRITE_DEVICE_CONFIG, WRITE_ALLOWLISTED_DEVICE_CONFIG) {
92                         originalConfig.forEach { (key, value) ->
93                             Log.i(TAG, "Resetting config \"${key.second}\" to \"$value\"")
94                             DeviceConfig.setProperty(
95                                     key.first, key.second, value, false /* makeDefault */)
96                         }
97                     }
98                 } cleanupStep {
99                     originalConfig.clear()
100                     usedConfig.clear()
101                 } cleanup {
102                     // Fold all cleanup actions into cleanup steps of an empty tryTest, so they are
103                     // all run even if exceptions are thrown, and exceptions are reported properly.
104                     currentTestCleanupActions.fold(tryTest { }) {
105                         tryBlock, action -> tryBlock.cleanupStep { action.run() }
106                     }.cleanup {
107                         currentTestCleanupActions.clear()
108                     }
109                 }
110             }
111         }
112     }
113 
114     /**
115      * Set a configuration key/value. After the test case ends, it will be restored to the value it
116      * had when this method was first called.
117      */
118     fun setConfig(namespace: String, key: String, value: String?): String? {
119         Log.i(TAG, "Setting config \"$key\" to \"$value\"")
120         val readWritePermissions =
121             arrayOf(READ_DEVICE_CONFIG, WRITE_DEVICE_CONFIG, WRITE_ALLOWLISTED_DEVICE_CONFIG)
122 
123         val keyPair = Pair(namespace, key)
124         val existingValue = runAsShell(*readWritePermissions) {
125             DeviceConfig.getProperty(namespace, key)
126         }
127         if (!originalConfig.containsKey(keyPair)) {
128             originalConfig[keyPair] = existingValue
129         }
130         usedConfig[keyPair] = value
131         if (existingValue == value) {
132             // Already the correct value. There may be a race if a change is already in flight,
133             // but if multiple threads update the config there is no way to fix that anyway.
134             Log.i(TAG, "\"$key\" already had value \"$value\"")
135             return value
136         }
137 
138         val future = CompletableFuture<String>()
139         val listener = DeviceConfig.OnPropertiesChangedListener {
140             // The listener receives updates for any change to any key, so don't react to
141             // changes that do not affect the relevant key
142             if (!it.keyset.contains(key)) return@OnPropertiesChangedListener
143             // "null" means absent in DeviceConfig : there is no such thing as a present but
144             // null value, so the following works even if |value| is null.
145             if (it.getString(key, null) == value) {
146                 future.complete(value)
147             }
148         }
149 
150         return tryTest {
151             runAsShell(*readWritePermissions) {
152                 DeviceConfig.addOnPropertiesChangedListener(
153                         namespace,
154                         inlineExecutor,
155                         listener)
156                 DeviceConfig.setProperty(
157                         namespace,
158                         key,
159                         value,
160                         false /* makeDefault */)
161                 // Don't drop the permission until the config is applied, just in case
162                 future.get(TIMEOUT_MS, TimeUnit.MILLISECONDS)
163             }.also {
164                 Log.i(TAG, "Config \"$key\" successfully set to \"$value\"")
165             }
166         } cleanup {
167             DeviceConfig.removeOnPropertiesChangedListener(listener)
168         }
169     }
170 
171     private val inlineExecutor get() = Executor { r -> r.run() }
172 
173     /**
174      * Add an action to be run after config cleanup when the current test case ends.
175      */
176     fun runAfterNextCleanup(action: ThrowingRunnable) {
177         currentTestCleanupActions.add(action)
178     }
179 }
180