1 files changed, 21 insertions, 25 deletions
diff --git a/hostsidetests/securitybulletin/src/android/security/cts/CVE_2021_0954.java b/hostsidetests/securitybulletin/src/android/security/cts/CVE_2021_0954.java
index 12287b0151d..95c90d4ec4b 100644
--- a/hostsidetests/securitybulletin/src/android/security/cts/CVE_2021_0954.java
+++ b/hostsidetests/securitybulletin/src/android/security/cts/CVE_2021_0954.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2021 The Android Open Source Project
+ * 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.
@@ -16,43 +16,39 @@
 
 package android.security.cts;
 
-import android.platform.test.annotations.AppModeFull;
+import static org.junit.Assume.assumeNoException;
+
 import android.platform.test.annotations.AsbSecurityTest;
 
+import com.android.sts.common.tradefed.testtype.StsExtraBusinessLogicHostTestBase;
 import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
-import com.android.sts.common.tradefed.testtype.StsExtraBusinessLogicHostTestBase;
 
-import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
 @RunWith(DeviceJUnit4ClassRunner.class)
 public class CVE_2021_0954 extends StsExtraBusinessLogicHostTestBase {
-    private static final String TEST_PKG = "android.security.cts.cve_2021_0954";
-    private static final String TEST_CLASS = TEST_PKG + "." + "DeviceTest";
-    private static final String TEST_APP = "CVE-2021-0954.apk";
-    private ITestDevice device;
-
-    @Before
-    public void setUp() throws Exception {
-        device = getDevice();
-        uninstallPackage(device, TEST_PKG);
-
-        /* Wake up the screen */
-        AdbUtils.runCommandLine("input keyevent KEYCODE_WAKEUP", device);
-        AdbUtils.runCommandLine("input keyevent KEYCODE_MENU", device);
-        AdbUtils.runCommandLine("input keyevent KEYCODE_HOME", device);
-    }
+    private static final String TEST_PKG = "android.security.cts.CVE_2021_0954";
 
-    @AppModeFull
     @AsbSecurityTest(cveBugId = 143559931)
     @Test
     public void testPocCVE_2021_0954() throws Exception {
-        installPackage(TEST_APP);
-        AdbUtils.runCommandLine("pm grant " + TEST_PKG + " android.permission.SYSTEM_ALERT_WINDOW",
-                device);
-        runDeviceTests(TEST_PKG, TEST_CLASS, "testVulnerableActivityPresence");
+        try {
+            ITestDevice device = getDevice();
+            uninstallPackage(device, TEST_PKG);
+
+            /* Wake up the screen */
+            AdbUtils.runCommandLine("input keyevent KEYCODE_WAKEUP", device);
+            AdbUtils.runCommandLine("input keyevent KEYCODE_MENU", device);
+            AdbUtils.runCommandLine("input keyevent KEYCODE_HOME", device);
+
+            installPackage("CVE-2021-0954.apk");
+            AdbUtils.runCommandLine(
+                    "pm grant " + TEST_PKG + " android.permission.SYSTEM_ALERT_WINDOW", device);
+            runDeviceTests(TEST_PKG, TEST_PKG + "." + "DeviceTest", "testOverlayButtonPresence");
+        } catch (Exception e) {
+            assumeNoException(e);
+        }
     }
 }