Merge "Add an initial set of tests for Trust." into tm-dev

tests/TrustTests/Android.bp

@@ -0,0 +1,39 @@
+// Copyright (C) 2022 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.
+
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+android_test {
+    name: "TrustTests",
+    srcs: [
+        "src/**/*.kt",
+    ],
+    static_libs: [
+        "androidx.test.rules",
+        "androidx.test.ext.junit",
+        "androidx.test.uiautomator",
+        "truth-prebuilt",
+    ],
+    libs: [
+        "android.test.runner",
+        "android.test.base",
+    ],
+    test_suites: [
+        "device-tests",
+    ],
+    platform_apis: true,
+    certificate: "platform",
+}

tests/TrustTests/AndroidManifest.xml

@@ -0,0 +1,75 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2022 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.
+  -->
+
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+          package="android.trust.test"
+          android:targetSandboxVersion="2">
+
+    <uses-permission android:name="android.permission.ACCESS_KEYGUARD_SECURE_STORAGE" />
+    <uses-permission android:name="android.permission.BIND_DEVICE_ADMIN" />
+    <uses-permission android:name="android.permission.CONTROL_KEYGUARD" />
+    <uses-permission android:name="android.permission.DEVICE_POWER" />
+    <uses-permission android:name="android.permission.INTERACT_ACROSS_USERS_FULL" />
+    <uses-permission android:name="android.permission.PROVIDE_TRUST_AGENT" />
+    <uses-permission android:name="android.permission.TRUST_LISTENER" />
+
+    <application>
+        <uses-library android:name="android.test.runner"/>
+        <activity android:name="android.trust.TrustTestActivity"
+                  android:exported="true">
+            <intent-filter>
+                <action android:name="android.intent.action.MAIN"/>
+                <category android:name="android.intent.category.LAUNCHER"/>
+            </intent-filter>
+        </activity>
+
+        <service
+            android:name=".UserUnlockRequestTrustAgent"
+            android:exported="true"
+            android:label="Test Agent"
+            android:permission="android.permission.BIND_TRUST_AGENT">
+            <intent-filter>
+                <action android:name="android.service.trust.TrustAgentService" />
+            </intent-filter>
+        </service>
+
+        <service
+            android:name=".LockUserTrustAgent"
+            android:exported="true"
+            android:label="Test Agent"
+            android:permission="android.permission.BIND_TRUST_AGENT">
+            <intent-filter>
+                <action android:name="android.service.trust.TrustAgentService" />
+            </intent-filter>
+        </service>
+
+        <service
+            android:name=".GrantAndRevokeTrustAgent"
+            android:exported="true"
+            android:label="Test Agent"
+            android:permission="android.permission.BIND_TRUST_AGENT">
+            <intent-filter>
+                <action android:name="android.service.trust.TrustAgentService" />
+            </intent-filter>
+        </service>
+    </application>
+
+    <!--  self-instrumenting test package. -->
+    <instrumentation android:name="androidx.test.runner.AndroidJUnitRunner"
+                     android:targetPackage="android.trust.test">
+    </instrumentation>
+</manifest>

tests/TrustTests/AndroidTest.xml

@@ -0,0 +1,28 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2022 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.
+  -->
+<configuration description="TrustTests configuration">
+    <option name="test-tag" value="TrustTests" />
+    <target_preparer class="com.android.tradefed.targetprep.suite.SuiteApkInstaller">
+        <option name="cleanup-apks" value="true" />
+        <option name="test-file-name" value="TrustTests.apk" />
+    </target_preparer>
+    <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
+        <option name="package" value="android.trust.test" />
+        <option name="runner" value="androidx.test.runner.AndroidJUnitRunner" />
+        <option name="hidden-api-checks" value="false" />
+    </test>
+</configuration>

tests/TrustTests/README.md

@@ -0,0 +1,40 @@
+# TrustTests framework tests
+
+These tests test the "trust" part of the platform primarily implemented via TrustManagerService in
+the system server and TrustAgentService in system apps.
+
+Tests are separated into separate files based on major groupings. When creating new tests, find a
+_closely_ matching existing test file or create a new test file. Prefer many test files over large
+test files.
+
+Each test file has its own trust agent. To create a new trust agent:
+
+1. Create a new class extending from `BaseTrustAgentService` class in your test file
+2. Add a new `<service>` stanza to `AndroidManifest.xml` in this directory for the new agent
+   following the pattern fo the existing agents.
+
+To run:
+
+```atest TrustTests```
+
+## Testing approach:
+
+1. Test the agent service as a black box; avoid inspecting internal state of the service or
+   modifying the system code outside of this directory.
+2. The primary interface to the system is through these three points:
+    1. `TrustAgentService`, your agent created by the `TrustAgentRule` and accessible via
+       the `agent` property of the rule.
+        1. Call command methods (e.g. `grantTrust`) directly on the agent
+        2. Listen to events (e.g. `onUserRequestedUnlock`) by implementing the method in
+           your test's agent class and tracking invocations. See `UserUnlockRequestTest` for an
+           example.
+    2. `TrustManager` which is the interface the rest of the system (e.g. SystemUI) has to the
+       service.
+        1. Through this API, simulate system events that the service cares about
+           (e.g. `reportUnlockAttempt`).
+    3. `TrustListener` which is the interface the rest of the system (e.g. SystemUI) uses to receive
+       events from the service.
+        1. Through this, verify behavior that affects the rest of the system. For example,
+           see `LockStateTrackingRule`.
+3. To re-use code between tests, prefer creating new rules alongside the existing rules or adding
+   functionality to a _closely_ matching existing rule.

tests/TrustTests/src/android/trust/BaseTrustAgentService.kt

@@ -0,0 +1,47 @@
+/*
+ * Copyright (C) 2022 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.
+ */
+package android.trust
+
+import android.service.trust.TrustAgentService
+import android.util.Log
+import kotlin.reflect.KClass
+
+/**
+ * Base class for test trust agents.
+ */
+abstract class BaseTrustAgentService : TrustAgentService() {
+
+    override fun onCreate() {
+        super.onCreate()
+        Log.d(TAG, "${this::class.simpleName} created")
+        instances[this::class] = this
+    }
+
+    override fun onDestroy() {
+        super.onDestroy()
+        instances.remove(this::class)
+    }
+
+    companion object {
+        private val instances =
+            mutableMapOf<KClass<out BaseTrustAgentService>, BaseTrustAgentService>()
+        private const val TAG = "BaseTrustAgentService"
+
+        fun instance(serviceClass: KClass<out BaseTrustAgentService>): BaseTrustAgentService? {
+            return instances[serviceClass]!!
+        }
+    }
+}

tests/TrustTests/src/android/trust/TrustTestActivity.kt

@@ -0,0 +1,30 @@
+/*
+ * Copyright (C) 2022 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.
+ */
+package android.trust
+
+import android.app.Activity
+import android.os.Bundle
+
+/**
+ * Activity for testing Trust.
+ */
+class TrustTestActivity : Activity() {
+
+    public override fun onCreate(icicle: Bundle?) {
+        super.onCreate(icicle)
+        setTurnScreenOn(true)
+    }
+}

tests/TrustTests/src/android/trust/test/GrantAndRevokeTrustTest.kt

@@ -0,0 +1,95 @@
+/*
+ * Copyright (C) 2022 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.
+ */
+
+package android.trust.test
+
+import android.trust.BaseTrustAgentService
+import android.trust.TrustTestActivity
+import android.trust.test.lib.LockStateTrackingRule
+import android.trust.test.lib.ScreenLockRule
+import android.trust.test.lib.TrustAgentRule
+import androidx.test.ext.junit.rules.ActivityScenarioRule
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
+import androidx.test.uiautomator.UiDevice
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.rules.RuleChain
+import org.junit.runner.RunWith
+
+/**
+ * Test for testing revokeTrust & grantTrust for non-renewable trust.
+ *
+ * atest TrustTests:GrantAndRevokeTrustTest
+ */
+@RunWith(AndroidJUnit4::class)
+class GrantAndRevokeTrustTest {
+    private val uiDevice = UiDevice.getInstance(getInstrumentation())
+    private val activityScenarioRule = ActivityScenarioRule(TrustTestActivity::class.java)
+    private val lockStateTrackingRule = LockStateTrackingRule()
+    private val trustAgentRule = TrustAgentRule<GrantAndRevokeTrustAgent>()
+
+    @get:Rule
+    val rule: RuleChain = RuleChain
+        .outerRule(activityScenarioRule)
+        .around(ScreenLockRule())
+        .around(lockStateTrackingRule)
+        .around(trustAgentRule)
+
+    @Before
+    fun manageTrust() {
+        trustAgentRule.agent.setManagingTrust(true)
+    }
+
+    // This test serves a baseline for Grant tests, verifying that the default behavior of the
+    // device is to lock when put to sleep
+    @Test
+    fun sleepingDeviceWithoutGrantLocksDevice() {
+        uiDevice.sleep()
+        await()
+
+        lockStateTrackingRule.assertLocked()
+    }
+
+    @Test
+    fun grantKeepsDeviceUnlocked() {
+        trustAgentRule.agent.grantTrust(GRANT_MESSAGE, 10000, 0)
+        uiDevice.sleep()
+        await()
+
+        lockStateTrackingRule.assertUnlocked()
+    }
+
+    @Test
+    fun grantKeepsDeviceUnlocked_untilRevoked() {
+        trustAgentRule.agent.grantTrust(GRANT_MESSAGE, 0, 0)
+        await()
+        uiDevice.sleep()
+        trustAgentRule.agent.revokeTrust()
+        await()
+
+        lockStateTrackingRule.assertLocked()
+    }
+
+    companion object {
+        private const val TAG = "GrantAndRevokeTrustTest"
+        private const val GRANT_MESSAGE = "granted by test"
+        private fun await() = Thread.sleep(250)
+    }
+}
+
+class GrantAndRevokeTrustAgent : BaseTrustAgentService()

tests/TrustTests/src/android/trust/test/LockUserTest.kt

@@ -0,0 +1,68 @@
+/*
+ * Copyright (C) 2022 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.
+ */
+
+package android.trust.test
+
+import android.trust.BaseTrustAgentService
+import android.trust.TrustTestActivity
+import android.trust.test.lib.LockStateTrackingRule
+import android.trust.test.lib.ScreenLockRule
+import android.trust.test.lib.TrustAgentRule
+import android.util.Log
+import androidx.test.ext.junit.rules.ActivityScenarioRule
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.google.common.truth.Truth.assertThat
+import org.junit.Ignore
+import org.junit.Rule
+import org.junit.Test
+import org.junit.rules.RuleChain
+import org.junit.runner.RunWith
+
+/**
+ * Test for testing lockUser.
+ *
+ * atest TrustTests:LockUserTest
+ */
+@RunWith(AndroidJUnit4::class)
+class LockUserTest {
+    private val activityScenarioRule = ActivityScenarioRule(TrustTestActivity::class.java)
+    private val lockStateTrackingRule = LockStateTrackingRule()
+    private val trustAgentRule = TrustAgentRule<LockUserTrustAgent>()
+
+    @get:Rule
+    val rule: RuleChain = RuleChain
+        .outerRule(activityScenarioRule)
+        .around(ScreenLockRule())
+        .around(lockStateTrackingRule)
+        .around(trustAgentRule)
+
+    @Ignore("Causes issues with subsequent tests") // TODO: Enable test
+    @Test
+    fun lockUser_locksTheDevice() {
+        Log.i(TAG, "Locking user")
+        trustAgentRule.agent.lockUser()
+        await()
+
+        assertThat(lockStateTrackingRule.lockState.locked).isTrue()
+    }
+
+    companion object {
+        private const val TAG = "LockUserTest"
+        private fun await() = Thread.sleep(250)
+    }
+}
+
+class LockUserTrustAgent : BaseTrustAgentService()

tests/TrustTests/src/android/trust/test/UserUnlockRequestTest.kt

@@ -0,0 +1,81 @@
+/*
+ * Copyright (C) 2022 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.
+ */
+package android.trust.test
+
+import android.app.trust.TrustManager
+import android.content.Context
+import android.trust.BaseTrustAgentService
+import android.trust.TrustTestActivity
+import android.trust.test.lib.ScreenLockRule
+import android.trust.test.lib.TrustAgentRule
+import android.util.Log
+import androidx.test.core.app.ApplicationProvider.getApplicationContext
+import androidx.test.ext.junit.rules.ActivityScenarioRule
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.google.common.truth.Truth.assertThat
+import org.junit.Rule
+import org.junit.Test
+import org.junit.rules.RuleChain
+import org.junit.runner.RunWith
+
+/**
+ * Test for testing the user unlock trigger.
+ *
+ * atest TrustTests:UserUnlockRequestTest
+ */
+@RunWith(AndroidJUnit4::class)
+class UserUnlockRequestTest {
+    private val context: Context = getApplicationContext()
+    private val trustManager = context.getSystemService(TrustManager::class.java) as TrustManager
+    private val userId = context.userId
+    private val activityScenarioRule = ActivityScenarioRule(TrustTestActivity::class.java)
+    private val trustAgentRule = TrustAgentRule<UserUnlockRequestTrustAgent>()
+
+    @get:Rule
+    val rule: RuleChain = RuleChain
+        .outerRule(activityScenarioRule)
+        .around(ScreenLockRule())
+        .around(trustAgentRule)
+
+    @Test
+    fun reportUserRequestedUnlock_propagatesToAgent() {
+        val oldCount = trustAgentRule.agent.onUserRequestedUnlockCallCount
+        trustManager.reportUserRequestedUnlock(userId)
+        await()
+
+        assertThat(trustAgentRule.agent.onUserRequestedUnlockCallCount)
+            .isEqualTo(oldCount + 1)
+    }
+
+    companion object {
+        private const val TAG = "UserUnlockRequestTest"
+        private fun await() = Thread.sleep(250)
+    }
+}
+
+class UserUnlockRequestTrustAgent : BaseTrustAgentService() {
+    var onUserRequestedUnlockCallCount: Long = 0
+        private set
+
+    override fun onUserRequestedUnlock() {
+        Log.i(TAG, "onUserRequestedUnlock")
+        onUserRequestedUnlockCallCount++
+    }
+
+    companion object {
+        private const val TAG = "UserUnlockRequestTrustAgent"
+    }
+}

tests/TrustTests/src/android/trust/test/lib/LockStateTrackingRule.kt

@@ -0,0 +1,83 @@
+/*
+ * Copyright (C) 2022 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.
+ */
+
+package android.trust.test.lib
+
+import android.app.trust.TrustManager
+import android.app.trust.TrustManager.TrustListener
+import android.content.Context
+import android.util.Log
+import android.view.WindowManagerGlobal
+import androidx.test.core.app.ApplicationProvider.getApplicationContext
+import com.google.common.truth.Truth.assertThat
+import org.junit.rules.TestRule
+import org.junit.runner.Description
+import org.junit.runners.model.Statement
+
+/**
+ * Rule for tracking the lock state of the device based on events emitted to [TrustListener].
+ */
+class LockStateTrackingRule : TestRule {
+    private val context: Context = getApplicationContext()
+    private val windowManager = WindowManagerGlobal.getWindowManagerService()
+
+    @Volatile lateinit var lockState: LockState
+        private set
+
+    override fun apply(base: Statement, description: Description) = object : Statement() {
+        override fun evaluate() {
+            lockState = LockState(locked = windowManager.isKeyguardLocked)
+            val trustManager = context.getSystemService(TrustManager::class.java) as TrustManager
+            val listener = Listener()
+
+            trustManager.registerTrustListener(listener)
+            try {
+                base.evaluate()
+            } finally {
+                trustManager.unregisterTrustListener(listener)
+            }
+        }
+    }
+
+    fun assertLocked() = assertThat(lockState.locked).isTrue()
+    fun assertUnlocked() = assertThat(lockState.locked).isFalse()
+
+    inner class Listener : TrustListener {
+        override fun onTrustChanged(
+            enabled: Boolean,
+            userId: Int,
+            flags: Int,
+            trustGrantedMessages: MutableList<String>
+        ) {
+            Log.d(TAG, "Device became trusted=$enabled")
+            lockState = lockState.copy(locked = !enabled)
+        }
+
+        override fun onTrustManagedChanged(enabled: Boolean, userId: Int) {
+        }
+
+        override fun onTrustError(message: CharSequence) {
+        }
+    }
+
+    data class LockState(
+        val locked: Boolean? = null
+    )
+
+    companion object {
+        private const val TAG = "LockStateTrackingRule"
+    }
+}

tests/TrustTests/src/android/trust/test/lib/ScreenLockRule.kt

@@ -0,0 +1,105 @@
+/*
+ * Copyright (C) 2022 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.
+ */
+
+package android.trust.test.lib
+
+import android.content.Context
+import android.util.Log
+import android.view.WindowManagerGlobal
+import androidx.test.core.app.ApplicationProvider.getApplicationContext
+import com.android.internal.widget.LockPatternUtils
+import com.android.internal.widget.LockscreenCredential
+import com.google.common.truth.Truth.assertWithMessage
+import org.junit.rules.TestRule
+import org.junit.runner.Description
+import org.junit.runners.model.Statement
+
+/**
+ * Sets a screen lock on the device for the duration of the test.
+ */
+class ScreenLockRule : TestRule {
+    private val context: Context = getApplicationContext()
+    private val windowManager = WindowManagerGlobal.getWindowManagerService()
+    private val lockPatternUtils = LockPatternUtils(context)
+    private var instantLockSavedValue = false
+
+    override fun apply(base: Statement, description: Description) = object : Statement() {
+        override fun evaluate() {
+            verifyNoScreenLockAlreadySet()
+            verifyKeyguardDismissed()
+            setScreenLock()
+            setLockOnPowerButton()
+
+            try {
+                base.evaluate()
+            } finally {
+                removeScreenLock()
+                revertLockOnPowerButton()
+            }
+        }
+    }
+
+    private fun verifyNoScreenLockAlreadySet() {
+        assertWithMessage("Screen Lock must not already be set on device")
+            .that(lockPatternUtils.isSecure(context.userId))
+            .isFalse()
+    }
+
+    private fun verifyKeyguardDismissed() {
+        windowManager.dismissKeyguard(null, null)
+        Thread.sleep(250)
+        assertWithMessage("Keyguard should be unlocked")
+            .that(windowManager.isKeyguardLocked)
+            .isFalse()
+    }
+
+    private fun setScreenLock() {
+        lockPatternUtils.setLockCredential(
+            LockscreenCredential.createPin(PIN),
+            LockscreenCredential.createNone(),
+            context.userId
+        )
+        assertWithMessage("Screen Lock should now be set")
+            .that(lockPatternUtils.isSecure(context.userId))
+            .isTrue()
+        Log.i(TAG, "Device PIN set to $PIN")
+    }
+
+    private fun setLockOnPowerButton() {
+        instantLockSavedValue = lockPatternUtils.getPowerButtonInstantlyLocks(context.userId)
+        lockPatternUtils.setPowerButtonInstantlyLocks(true, context.userId)
+    }
+
+    private fun removeScreenLock() {
+        lockPatternUtils.setLockCredential(
+            LockscreenCredential.createNone(),
+            LockscreenCredential.createPin(PIN),
+            context.userId
+        )
+        Log.i(TAG, "Device PIN cleared; waiting 50 ms then dismissing Keyguard")
+        Thread.sleep(50)
+        windowManager.dismissKeyguard(null, null)
+    }
+
+    private fun revertLockOnPowerButton() {
+        lockPatternUtils.setPowerButtonInstantlyLocks(instantLockSavedValue, context.userId)
+    }
+
+    companion object {
+        private const val TAG = "ScreenLockRule"
+        private const val PIN = "0000"
+    }
+}

tests/TrustTests/src/android/trust/test/lib/TrustAgentRule.kt

@@ -0,0 +1,117 @@
+/*
+ * Copyright (C) 2022 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.
+ */
+
+package android.trust.test.lib
+
+import android.app.trust.TrustManager
+import android.content.ComponentName
+import android.content.Context
+import android.trust.BaseTrustAgentService
+import android.util.Log
+import androidx.test.core.app.ApplicationProvider.getApplicationContext
+import com.android.internal.widget.LockPatternUtils
+import com.google.common.truth.Truth.assertWithMessage
+import org.junit.rules.TestRule
+import org.junit.runner.Description
+import org.junit.runners.model.Statement
+import kotlin.reflect.KClass
+
+/**
+ * Enables a trust agent and causes the system service to bind to it.
+ *
+ * The enabled agent can be accessed during the test via the [agent] property.
+ *
+ * @constructor Creates the rule. Do not use; instead, use [invoke].
+ */
+class TrustAgentRule<T : BaseTrustAgentService>(
+    private val serviceClass: KClass<T>
+) : TestRule {
+    private val context: Context = getApplicationContext()
+    private val trustManager = context.getSystemService(TrustManager::class.java) as TrustManager
+    private val lockPatternUtils = LockPatternUtils(context)
+
+    val agent get() = BaseTrustAgentService.instance(serviceClass) as T
+
+    override fun apply(base: Statement, description: Description) = object : Statement() {
+        override fun evaluate() {
+            verifyTrustServiceRunning()
+            unlockDeviceWithCredential()
+            enableTrustAgent()
+            waitForEnablement()
+
+            try {
+                verifyAgentIsRunning()
+                base.evaluate()
+            } finally {
+                disableTrustAgent()
+            }
+        }
+    }
+
+    private fun verifyTrustServiceRunning() {
+        assertWithMessage("Trust service is not running").that(trustManager).isNotNull()
+    }
+
+    private fun unlockDeviceWithCredential() {
+        Log.d(TAG, "Unlocking device with credential")
+        trustManager.reportUnlockAttempt(true, context.userId)
+    }
+
+    private fun enableTrustAgent() {
+        val componentName = ComponentName(context, serviceClass.java)
+        val userId = context.userId
+        Log.i(TAG, "Enabling trust agent ${componentName.flattenToString()} for user $userId")
+        val agents = mutableListOf(componentName)
+            .plus(lockPatternUtils.getEnabledTrustAgents(userId))
+            .distinct()
+        lockPatternUtils.setEnabledTrustAgents(agents, userId)
+    }
+
+    private fun waitForEnablement() {
+        Log.d(TAG, "Waiting for $WAIT_TIME ms")
+        Thread.sleep(WAIT_TIME)
+        Log.d(TAG, "Done waiting")
+    }
+
+    private fun verifyAgentIsRunning() {
+        assertWithMessage("${serviceClass.simpleName} should be running")
+            .that(BaseTrustAgentService.instance(serviceClass)).isNotNull()
+    }
+
+    private fun disableTrustAgent() {
+        val componentName = ComponentName(context, serviceClass.java)
+        val userId = context.userId
+        Log.i(TAG, "Disabling trust agent ${componentName.flattenToString()} for user $userId")
+        val agents = lockPatternUtils.getEnabledTrustAgents(userId).toMutableList()
+            .distinct()
+            .minus(componentName)
+        lockPatternUtils.setEnabledTrustAgents(agents, userId)
+    }
+
+    companion object {
+        /**
+         * Creates a new rule for the specified agent class. Example usage:
+         * ```
+         *   @get:Rule val rule = TrustAgentRule<MyTestAgent>()
+         * ```
+         */
+        inline operator fun <reified T : BaseTrustAgentService> invoke() =
+            TrustAgentRule(T::class)
+
+        private const val TAG = "TrustAgentRule"
+        private val WAIT_TIME = 1000L
+    }
+}