Add test utilities for shell permission identity
Add utilities allowing to run lambdas with shell permissions, in kotlin
and java tests.
Bug: 168868607
Test: m
Change-Id: Iaab38d498338b9e95d263789119cd0813983087d
diff --git a/staticlibs/devicetests/com/android/testutils/TestPermissionUtil.kt b/staticlibs/devicetests/com/android/testutils/TestPermissionUtil.kt
new file mode 100644
index 0000000..f557f18
--- /dev/null
+++ b/staticlibs/devicetests/com/android/testutils/TestPermissionUtil.kt
@@ -0,0 +1,77 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+@file:JvmName("TestPermissionUtil")
+
+package com.android.testutils
+
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.testutils.ExceptionUtils.ThrowingRunnable
+import com.android.testutils.ExceptionUtils.ThrowingSupplier
+
+/**
+ * Run the specified [task] with the specified [permissions] obtained through shell
+ * permission identity.
+ *
+ * Passing in an empty list of permissions can grant all shell permissions, but this is
+ * discouraged as it also causes the process to temporarily lose non-shell permissions.
+ */
+fun <T> runAsShell(vararg permissions: String, task: () -> T): T {
+ val autom = InstrumentationRegistry.getInstrumentation().uiAutomation
+ autom.adoptShellPermissionIdentity(*permissions)
+ try {
+ return task()
+ } finally {
+ autom.dropShellPermissionIdentity()
+ }
+}
+
+/**
+ * Convenience overload of [runAsShell] that uses a [ThrowingSupplier] for Java callers, when
+ * only one/two/three permissions are needed.
+ */
+@JvmOverloads
+fun <T> runAsShell(
+ perm1: String,
+ perm2: String = "",
+ perm3: String = "",
+ supplier: ThrowingSupplier<T>
+): T = runAsShell(*getNonEmptyVarargs(perm1, perm2, perm3)) { supplier.get() }
+
+/**
+ * Convenience overload of [runAsShell] that uses a [ThrowingRunnable] for Java callers, when
+ * only one/two/three permissions are needed.
+ */
+@JvmOverloads
+fun runAsShell(
+ perm1: String,
+ perm2: String = "",
+ perm3: String = "",
+ runnable: ThrowingRunnable
+): Unit = runAsShell(*getNonEmptyVarargs(perm1, perm2, perm3)) { runnable.run() }
+
+/**
+ * Get an array containing the first consecutive non-empty arguments out of three arguments.
+ *
+ * The first argument is assumed to be non-empty.
+ */
+private fun getNonEmptyVarargs(arg1: String, arg2: String, arg3: String): Array<String> {
+ return when {
+ arg2 == "" -> arrayOf(arg1)
+ arg3 == "" -> arrayOf(arg1, arg2)
+ else -> arrayOf(arg1, arg2, arg3)
+ }
+}
\ No newline at end of file