Merge "Make TestRules for hostside tests" into udc-dev

5 files changed, 144 insertions, 96 deletions

diff --git a/tests/hostside/safetycenter/src/android/safetycenter/hostside/ITestDeviceExtensions.kt b/tests/hostside/safetycenter/src/android/safetycenter/hostside/ITestDeviceExtensions.kt
deleted file mode 100644
index f8d4d64c0..000000000
--- a/tests/hostside/safetycenter/src/android/safetycenter/hostside/ITestDeviceExtensions.kt
+++ /dev/null
@@ -1,47 +0,0 @@
-/*
- * Copyright (C) 2023 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.safetycenter.hostside
-
-import android.cts.statsdatom.lib.DeviceUtils
-import com.android.tradefed.device.ITestDevice
-import com.android.tradefed.util.CommandStatus
-import java.io.IOException
-
-/** Runs a test in the Safety Center hostside test helper app */
-internal fun ITestDevice.runTest(testClassName: String, testMethodName: String) {
-    DeviceUtils.runDeviceTests(this, HelperApp.PACKAGE_NAME, testClassName, testMethodName)
-}
-
-/** Checks whether this [ITestDevice] supports Safety Center. */
-internal fun ITestDevice.supportsSafetyCenter(): Boolean =
-    executeShellCommandOrThrow("cmd safety_center supported").toBoolean()
-
-/** Checks whether Safety Center is currently enabled on this [ITestDevice]. */
-internal fun ITestDevice.isSafetyCenterEnabled(): Boolean =
-    executeShellCommandOrThrow("cmd safety_center enabled").toBoolean()
-
-/** Returns the package name of Safety Center on this [ITestDevice]. */
-internal fun ITestDevice.getSafetyCenterPackageName(): String =
-    executeShellCommandOrThrow("cmd safety_center package-name")
-
-private fun ITestDevice.executeShellCommandOrThrow(command: String): String {
-    val result = executeShellV2Command(command)
-    if (result.status != CommandStatus.SUCCESS) {
-        throw IOException("$command exited with status ${result.exitCode}")
-    }
-    return result.stdout.trim()
-}
diff --git a/tests/hostside/safetycenter/src/android/safetycenter/hostside/SafetyCenterInteractionLoggingHostTest.kt b/tests/hostside/safetycenter/src/android/safetycenter/hostside/SafetyCenterInteractionLoggingHostTest.kt
index c49e20c15..5ef8ed84a 100644
--- a/tests/hostside/safetycenter/src/android/safetycenter/hostside/SafetyCenterInteractionLoggingHostTest.kt
+++ b/tests/hostside/safetycenter/src/android/safetycenter/hostside/SafetyCenterInteractionLoggingHostTest.kt
@@ -17,8 +17,9 @@
 package android.safetycenter.hostside
 
 import android.cts.statsdatom.lib.ConfigUtils
-import android.cts.statsdatom.lib.DeviceUtils
 import android.cts.statsdatom.lib.ReportUtils
+import android.safetycenter.hostside.rules.HelperAppRule
+import android.safetycenter.hostside.rules.RequireSafetyCenterRule
 import com.android.os.AtomsProto.Atom
 import com.android.os.AtomsProto.SafetyCenterInteractionReported
 import com.android.os.AtomsProto.SafetyCenterInteractionReported.Action
@@ -27,56 +28,45 @@ import com.android.tradefed.testtype.DeviceJUnit4ClassRunner
 import com.android.tradefed.testtype.junit4.BaseHostJUnit4Test
 import com.google.common.truth.Truth.assertThat
 import org.junit.After
-import org.junit.Assume.assumeTrue
 import org.junit.Before
 import org.junit.Ignore
+import org.junit.Rule
 import org.junit.Test
+import org.junit.rules.RuleChain
 import org.junit.runner.RunWith
 
 /** Host-side tests for Safety Center statsd logging. */
 @RunWith(DeviceJUnit4ClassRunner::class)
 class SafetyCenterInteractionLoggingHostTest : BaseHostJUnit4Test() {
 
-    private val shouldRunTests: Boolean by lazy {
-        // Device is not available when the test is first constructed
-        // TODO(b/239682646): These tests should enable safety center instead of only running when
-        // it's already enabled.
-        device.supportsSafetyCenter() && device.isSafetyCenterEnabled()
-    }
+    private val safetyCenterRule = RequireSafetyCenterRule(this)
+    private val helperAppRule = HelperAppRule(this, HelperApp.APK_NAME, HelperApp.PACKAGE_NAME)
 
-    @Before
-    fun assumeDeviceSupportsSafetyCenterToRunTests() {
-        assumeTrue(shouldRunTests)
-    }
+    @Rule
+    @JvmField
+    val rules: RuleChain = RuleChain.outerRule(safetyCenterRule).around(helperAppRule)
 
     @Before
     fun setUp() {
-        if (!shouldRunTests) return
-
         ConfigUtils.removeConfig(device)
         ReportUtils.clearReports(device)
         ConfigUtils.uploadConfigForPushedAtom(
             device,
-            device.getSafetyCenterPackageName(),
+            safetyCenterRule.getSafetyCenterPackageName(),
             Atom.SAFETY_CENTER_INTERACTION_REPORTED_FIELD_NUMBER
         )
-        DeviceUtils.installTestApp(device, HelperApp.APK_NAME, HelperApp.PACKAGE_NAME, build)
-
         // TODO(b/239682646): Consider adding a target preparer that unlocks the device (like CTS)
     }
 
     @After
     fun tearDown() {
-        if (!shouldRunTests) return
-
         ConfigUtils.removeConfig(device)
         ReportUtils.clearReports(device)
-        DeviceUtils.uninstallTestApp(device, HelperApp.PACKAGE_NAME)
     }
 
     @Test
     fun openSafetyCenter_recordsSafetyCenterViewedEvent() {
-        device.runTest(TEST_CLASS_NAME, testMethodName = "openSafetyCenter")
+        helperAppRule.runTest(TEST_CLASS_NAME, testMethodName = "openSafetyCenter")
 
         val safetyCenterViewedAtoms = getInteractionReportedAtoms(Action.SAFETY_CENTER_VIEWED)
 
@@ -85,7 +75,7 @@ class SafetyCenterInteractionLoggingHostTest : BaseHostJUnit4Test() {
 
     @Test
     fun sendNotification_recordsNotificationPostedEvent() {
-        device.runTest(
+        helperAppRule.runTest(
             testClassName = ".SafetyCenterNotificationLoggingHelperTests",
             testMethodName = "sendNotification"
         )
@@ -99,7 +89,7 @@ class SafetyCenterInteractionLoggingHostTest : BaseHostJUnit4Test() {
 
     @Test
     fun openSubpageFromIntentExtra_recordsEventWithUnknownNavigationSource() {
-        device.runTest(TEST_CLASS_NAME, testMethodName = "openSubpageFromIntentExtra")
+        helperAppRule.runTest(TEST_CLASS_NAME, testMethodName = "openSubpageFromIntentExtra")
 
         val safetyCenterViewedAtoms = getInteractionReportedAtoms(Action.SAFETY_CENTER_VIEWED)
 
@@ -116,7 +106,7 @@ class SafetyCenterInteractionLoggingHostTest : BaseHostJUnit4Test() {
     @Ignore
     // TODO(b/278202773): Fix/de-flake this test
     fun openSubpageFromHomepage_recordsEventWithSafetyCenterNavigationSource() {
-        device.runTest(TEST_CLASS_NAME, testMethodName = "openSubpageFromHomepage")
+        helperAppRule.runTest(TEST_CLASS_NAME, testMethodName = "openSubpageFromHomepage")
 
         val safetyCenterViewedAtoms = getInteractionReportedAtoms(Action.SAFETY_CENTER_VIEWED)
 
@@ -132,7 +122,7 @@ class SafetyCenterInteractionLoggingHostTest : BaseHostJUnit4Test() {
     @Ignore
     // TODO(b/278202773): Fix/de-flake this test
     fun openSubpageFromSettingsSearch_recordsEventWithSettingsNavigationSource() {
-        device.runTest(TEST_CLASS_NAME, testMethodName = "openSubpageFromSettingsSearch")
+        helperAppRule.runTest(TEST_CLASS_NAME, testMethodName = "openSubpageFromSettingsSearch")
 
         val safetyCenterViewedAtoms = getInteractionReportedAtoms(Action.SAFETY_CENTER_VIEWED)
 
diff --git a/tests/hostside/safetycenter/src/android/safetycenter/hostside/SafetySourceStateCollectedLoggingHostTest.kt b/tests/hostside/safetycenter/src/android/safetycenter/hostside/SafetySourceStateCollectedLoggingHostTest.kt
index ad00450d6..8c00ef35c 100644
--- a/tests/hostside/safetycenter/src/android/safetycenter/hostside/SafetySourceStateCollectedLoggingHostTest.kt
+++ b/tests/hostside/safetycenter/src/android/safetycenter/hostside/SafetySourceStateCollectedLoggingHostTest.kt
@@ -17,62 +17,52 @@
 package android.safetycenter.hostside
 
 import android.cts.statsdatom.lib.ConfigUtils
-import android.cts.statsdatom.lib.DeviceUtils
 import android.cts.statsdatom.lib.ReportUtils
+import android.safetycenter.hostside.rules.HelperAppRule
+import android.safetycenter.hostside.rules.RequireSafetyCenterRule
 import com.android.os.AtomsProto.Atom
 import com.android.os.AtomsProto.SafetySourceStateCollected
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner
 import com.android.tradefed.testtype.junit4.BaseHostJUnit4Test
 import com.google.common.truth.Truth.assertThat
 import org.junit.After
-import org.junit.Assume.assumeTrue
 import org.junit.Before
+import org.junit.Rule
 import org.junit.Test
+import org.junit.rules.RuleChain
 import org.junit.runner.RunWith
 
 /** Host-side tests for Safety Center statsd logging. */
 @RunWith(DeviceJUnit4ClassRunner::class)
 class SafetySourceStateCollectedLoggingHostTest : BaseHostJUnit4Test() {
 
-    // Use lazy here because device is not available when the test is first constructed
-    private val shouldRunTests: Boolean by lazy {
-        // TODO(b/239682646): These tests should enable Safety Center
-        device.supportsSafetyCenter() && device.isSafetyCenterEnabled()
-    }
+    private val safetyCenterRule = RequireSafetyCenterRule(this)
+    private val helperAppRule = HelperAppRule(this, HelperApp.APK_NAME, HelperApp.PACKAGE_NAME)
 
-    @Before
-    fun assumeDeviceSupportsSafetyCenterToRunTests() {
-        assumeTrue(shouldRunTests)
-    }
+    @Rule
+    @JvmField
+    val rules: RuleChain = RuleChain.outerRule(safetyCenterRule).around(helperAppRule)
 
     @Before
     fun setUp() {
-        if (!shouldRunTests) return
-
         ConfigUtils.removeConfig(device)
         ReportUtils.clearReports(device)
-
-        val config = ConfigUtils.createConfigBuilder(device.getSafetyCenterPackageName())
+        val config = ConfigUtils.createConfigBuilder(safetyCenterRule.getSafetyCenterPackageName())
         ConfigUtils.addGaugeMetric(config, Atom.SAFETY_STATE_FIELD_NUMBER)
         ConfigUtils.addEventMetric(config, Atom.SAFETY_SOURCE_STATE_COLLECTED_FIELD_NUMBER)
         ConfigUtils.uploadConfig(device, config)
-        DeviceUtils.installTestApp(device, HelperApp.APK_NAME, HelperApp.PACKAGE_NAME, build)
-
         // TODO(b/239682646): Consider adding a target preparer that unlocks the device (like CTS)
     }
 
     @After
     fun tearDown() {
-        if (!shouldRunTests) return
-
         ReportUtils.clearReports(device)
         ConfigUtils.removeConfig(device)
-        DeviceUtils.uninstallTestApp(device, HelperApp.PACKAGE_NAME)
     }
 
     @Test
     fun triggerStatsPull_atomsPushedForAllSources() {
-        device.runTest(TEST_CLASS_NAME, "triggerStatsPull")
+        helperAppRule.runTest(TEST_CLASS_NAME, "triggerStatsPull")
 
         val sourceStateAtoms = getSafetySourceStateCollectedAtoms()
 
@@ -86,7 +76,7 @@ class SafetySourceStateCollectedLoggingHostTest : BaseHostJUnit4Test() {
 
     @Test
     fun triggerStatsPull_atomsHaveCollectionTypeAutomatic() {
-        device.runTest(TEST_CLASS_NAME, "triggerStatsPull")
+        helperAppRule.runTest(TEST_CLASS_NAME, "triggerStatsPull")
 
         val sourceStateAtoms = getSafetySourceStateCollectedAtoms()
 
@@ -96,7 +86,7 @@ class SafetySourceStateCollectedLoggingHostTest : BaseHostJUnit4Test() {
 
     @Test
     fun setSafetySourceData_atomPushedForThatSource() {
-        device.runTest(TEST_CLASS_NAME, "setSafetySourceData_source1")
+        helperAppRule.runTest(TEST_CLASS_NAME, "setSafetySourceData_source1")
 
         val sourceStateAtoms = getSafetySourceStateCollectedAtoms()
 
@@ -106,7 +96,7 @@ class SafetySourceStateCollectedLoggingHostTest : BaseHostJUnit4Test() {
 
     @Test
     fun setSafetySourceData_atomHasCollectionTypeSourceUpdated() {
-        device.runTest(TEST_CLASS_NAME, "setSafetySourceData_source1")
+        helperAppRule.runTest(TEST_CLASS_NAME, "setSafetySourceData_source1")
 
         val sourceStateAtoms = getSafetySourceStateCollectedAtoms()
 
diff --git a/tests/hostside/safetycenter/src/android/safetycenter/hostside/rules/HelperAppRule.kt b/tests/hostside/safetycenter/src/android/safetycenter/hostside/rules/HelperAppRule.kt
new file mode 100644
index 000000000..cf46a8443
--- /dev/null
+++ b/tests/hostside/safetycenter/src/android/safetycenter/hostside/rules/HelperAppRule.kt
@@ -0,0 +1,57 @@
+/*
+ * Copyright (C) 2023 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.safetycenter.hostside.rules
+
+import android.cts.statsdatom.lib.DeviceUtils
+import com.android.tradefed.testtype.junit4.BaseHostJUnit4Test
+import org.junit.rules.TestRule
+import org.junit.runner.Description
+import org.junit.runners.model.Statement
+
+/**
+ * JUnit rule for host side tests that installs a helper app with the given [apkName] and
+ * [packageName].
+ */
+class HelperAppRule(
+    private val hostTestClass: BaseHostJUnit4Test,
+    private val apkName: String,
+    private val packageName: String
+) : TestRule {
+
+    override fun apply(base: Statement, description: Description): Statement {
+        return object : Statement() {
+            override fun evaluate() {
+                try {
+                    DeviceUtils.installTestApp(
+                        hostTestClass.device,
+                        apkName,
+                        packageName,
+                        hostTestClass.build
+                    )
+                    base.evaluate()
+                } finally {
+                    DeviceUtils.uninstallTestApp(hostTestClass.device, packageName)
+                }
+            }
+        }
+    }
+
+    /** Runs the specified test in the helper app. */
+    fun runTest(testClassName: String, testMethodName: String) {
+        DeviceUtils.runDeviceTests(hostTestClass.device, packageName, testClassName, testMethodName)
+    }
+}
diff --git a/tests/hostside/safetycenter/src/android/safetycenter/hostside/rules/RequireSafetyCenterRule.kt b/tests/hostside/safetycenter/src/android/safetycenter/hostside/rules/RequireSafetyCenterRule.kt
new file mode 100644
index 000000000..809fe5f0f
--- /dev/null
+++ b/tests/hostside/safetycenter/src/android/safetycenter/hostside/rules/RequireSafetyCenterRule.kt
@@ -0,0 +1,58 @@
+/*
+ * Copyright (C) 2023 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.safetycenter.hostside.rules
+
+import com.android.tradefed.testtype.junit4.BaseHostJUnit4Test
+import com.android.tradefed.util.CommandStatus
+import java.io.IOException
+import org.junit.Assume.assumeTrue
+import org.junit.rules.TestRule
+import org.junit.runner.Description
+import org.junit.runners.model.Statement
+
+/** JUnit rule for host side tests that requires Safety Center to be supported and enabled. */
+class RequireSafetyCenterRule(private val hostTestClass: BaseHostJUnit4Test) : TestRule {
+
+    private val safetyCenterSupported: Boolean by lazy {
+        executeShellCommandOrThrow("cmd safety_center supported").toBoolean()
+    }
+    private val safetyCenterEnabled: Boolean by lazy {
+        executeShellCommandOrThrow("cmd safety_center enabled").toBoolean()
+    }
+
+    override fun apply(base: Statement, description: Description): Statement {
+        return object : Statement() {
+            override fun evaluate() {
+                assumeTrue("Test device does not support Safety Center", safetyCenterSupported)
+                assumeTrue("Safety Center is not enabled on test device", safetyCenterEnabled)
+                base.evaluate()
+            }
+        }
+    }
+
+    /** Returns the package name of Safety Center on the test device. */
+    fun getSafetyCenterPackageName(): String =
+        executeShellCommandOrThrow("cmd safety_center package-name")
+
+    private fun executeShellCommandOrThrow(command: String): String {
+        val result = hostTestClass.device.executeShellV2Command(command)
+        if (result.status != CommandStatus.SUCCESS) {
+            throw IOException("$command exited with status ${result.exitCode}")
+        }
+        return result.stdout.trim()
+    }
+}