Refactor TrustTests waiting into common function.

Bug: 225231929
Test: atest TrustTests
Change-Id: I94e7da2fbe080027d5b2cbb7ecf0244aa149f645

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

@@ -41,7 +41,7 @@ abstract class BaseTrustAgentService : TrustAgentService() {
         private const val TAG = "BaseTrustAgentService"
 
         fun instance(serviceClass: KClass<out BaseTrustAgentService>): BaseTrustAgentService? {
-            return instances[serviceClass]!!
+            return instances[serviceClass]
         }
     }
 }

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

@@ -57,7 +57,6 @@ class LockUserTest {
 
     companion object {
         private const val TAG = "LockUserTest"
-        private fun await() = Thread.sleep(250)
     }
 }
 

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

@@ -22,7 +22,6 @@ 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
@@ -53,37 +52,12 @@ class LockStateTrackingRule : TestRule {
     }
 
     fun assertLocked() {
-        val maxWaits = 50
-        var waitCount = 0
-
-        // First verify we get the call in LockState via TrustListener
-        while ((lockState.locked == false) && waitCount < maxWaits) {
-            Log.i(TAG, "phone still unlocked (TrustListener), wait 50ms more ($waitCount)")
-            Thread.sleep(50)
-            waitCount++
-        }
-        assertThat(lockState.locked).isTrue()
-
-        // TODO(b/225231929): refactor checks into one loop and re-use for assertUnlocked
-        // Then verify we get the window manager locked
-        while (!windowManager.isKeyguardLocked && waitCount < maxWaits) {
-            Log.i(TAG, "phone still unlocked (WindowManager), wait 50ms more ($waitCount)")
-            Thread.sleep(50)
-            waitCount++
-        }
-        assertThat(windowManager.isKeyguardLocked).isTrue()
+        wait("un-locked per TrustListener") { lockState.locked == true }
+        wait("keyguard lock") { windowManager.isKeyguardLocked }
     }
 
     fun assertUnlocked() {
-        val maxWaits = 50
-        var waitCount = 0
-
-        while ((lockState.locked == true) && waitCount < maxWaits) {
-            Log.i(TAG, "phone still unlocked, wait 50ms more ($waitCount)")
-            Thread.sleep(50)
-            waitCount++
-        }
-        assertThat(lockState.locked).isFalse()
+        wait("locked per TrustListener") { lockState.locked == false }
     }
 
     inner class Listener : TrustListener {

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

@@ -42,7 +42,7 @@ class ScreenLockRule : TestRule {
     override fun apply(base: Statement, description: Description) = object : Statement() {
         override fun evaluate() {
             verifyNoScreenLockAlreadySet()
-            verifyKeyguardDismissed()
+            dismissKeyguard()
             setScreenLock()
             setLockOnPowerButton()
 
@@ -51,7 +51,7 @@ class ScreenLockRule : TestRule {
             } finally {
                 removeScreenLock()
                 revertLockOnPowerButton()
-                verifyKeyguardDismissed()
+                dismissKeyguard()
             }
         }
     }
@@ -62,30 +62,21 @@ class ScreenLockRule : TestRule {
                 .isFalse()
     }
 
-    private fun verifyKeyguardDismissed() {
-        val maxWaits = 30
-        var waitCount = 0
-
-        while (windowManager.isKeyguardLocked && waitCount < maxWaits) {
-            Log.i(TAG, "Keyguard still showing; attempting to dismiss and wait 50ms ($waitCount)")
+    fun dismissKeyguard() {
+        wait("keyguard dismissed") { count ->
             windowManager.dismissKeyguard(null, null)
 
             // Sometimes, bouncer gets shown due to a race, so we have to put display to sleep
             // and wake it back up to get it to go away
-            if (waitCount >= 10 && waitCount % 5 == 0) {
-                Log.i(TAG, "Escalation: attempting screen off/on to get rid of bouncer (+500ms)")
+            if (count >= 10 && count % 5 == 0) {
+                Log.i(TAG, "Escalation: attempting screen off/on to get rid of bouncer")
                 uiDevice.sleep()
                 Thread.sleep(250)
                 uiDevice.wakeUp()
-                Thread.sleep(250)
             }
 
-            Thread.sleep(50)
-            waitCount++
+            !windowManager.isKeyguardLocked
         }
-        assertWithMessage("Keyguard should be unlocked")
-                .that(windowManager.isKeyguardLocked)
-                .isFalse()
     }
 
     private fun setScreenLock() {
@@ -94,9 +85,7 @@ class ScreenLockRule : TestRule {
                 LockscreenCredential.createNone(),
                 context.userId
         )
-        assertWithMessage("Screen Lock should now be set")
-                .that(lockPatternUtils.isSecure(context.userId))
-                .isTrue()
+        wait("screen lock set") { lockPatternUtils.isSecure(context.userId) }
         Log.i(TAG, "Device PIN set to $PIN")
     }
 
@@ -110,21 +99,16 @@ class ScreenLockRule : TestRule {
                 LockscreenCredential.createNone(),
                 LockscreenCredential.createPin(PIN),
                 context.userId)
-        Thread.sleep(100)
+        Log.i(TAG, "Removing screen lock")
         assertWithMessage("Lock screen credential should be unset")
                 .that(lockCredentialUnset)
                 .isTrue()
 
         lockPatternUtils.setLockScreenDisabled(true, context.userId)
-        Thread.sleep(100)
-        assertWithMessage("Lockscreen needs to be disabled")
-                .that(lockPatternUtils.isLockScreenDisabled(context.userId))
-                .isTrue()
-
-        // this is here because somehow it helps the keyguard not get stuck
-        uiDevice.sleep()
-        Thread.sleep(500) // delay added to avoid initiating camera by double clicking power
-        uiDevice.wakeUp()
+        wait("screen lock un-set") {
+            lockPatternUtils.isLockScreenDisabled(context.userId)
+        }
+        wait("screen lock insecure") { !lockPatternUtils.isSecure(context.userId) }
     }
 
     private fun revertLockOnPowerButton() {

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

@@ -50,7 +50,6 @@ class TrustAgentRule<T : BaseTrustAgentService>(
             verifyTrustServiceRunning()
             unlockDeviceWithCredential()
             enableTrustAgent()
-            waitForEnablement()
 
             try {
                 verifyAgentIsRunning()
@@ -80,15 +79,10 @@ class TrustAgentRule<T : BaseTrustAgentService>(
         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()
+        wait("${serviceClass.simpleName} to be running") {
+            BaseTrustAgentService.instance(serviceClass) != null
+        }
     }
 
     private fun disableTrustAgent() {
@@ -112,6 +106,5 @@ class TrustAgentRule<T : BaseTrustAgentService>(
             TrustAgentRule(T::class)
 
         private const val TAG = "TrustAgentRule"
-        private val WAIT_TIME = 1000L
     }
 }

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

@@ -0,0 +1,51 @@
+/*
+ * 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.util.Log
+import com.google.common.truth.Truth.assertWithMessage
+
+private const val TAG = "TrustTestUtils"
+
+/**
+ * Waits for [conditionFunction] to be true with a failed assertion if it is not after [maxWait]
+ * ms.
+ *
+ * The condition function can perform additional logic (for example, logging or attempting to make
+ * the condition become true).
+ *
+ * @param conditionFunction function which takes the attempt count & returns whether the condition
+ *                          is met
+ */
+internal fun wait(
+    description: String? = null,
+    maxWait: Long = 1500L,
+    rate: Long = 50L,
+    conditionFunction: (count: Int) -> Boolean
+) {
+    var waited = 0L
+    var count = 0
+    while (!conditionFunction.invoke(count)) {
+        assertWithMessage("Condition exceeded maximum wait time of $maxWait ms: $description")
+            .that(waited <= maxWait)
+            .isTrue()
+        waited += rate
+        count++
+        Log.i(TAG, "Waiting for $description ($waited/$maxWait) #$count")
+        Thread.sleep(rate)
+    }
+}