Replace reboot with NotificationListener in DeviceOwnerTest.

Fixes: 134555169
Test: atest com.android.cts.devicepolicy.DeviceOwnerTest#testAdminActionBookkeeping
Change-Id: Ida648dffbc54f2cb2aaa0eb7deaa4c4ea64745b5
Merged-In: Ida648dffbc54f2cb2aaa0eb7deaa4c4ea64745b5
(cherry picked from commit e4e07e6680700d2a3892851cee15e06995482f84)
diff --git a/hostsidetests/devicepolicy/app/DeviceOwner/AndroidManifest.xml b/hostsidetests/devicepolicy/app/DeviceOwner/AndroidManifest.xml
index 1b5bdac..da569f3 100644
--- a/hostsidetests/devicepolicy/app/DeviceOwner/AndroidManifest.xml
+++ b/hostsidetests/devicepolicy/app/DeviceOwner/AndroidManifest.xml
@@ -82,6 +82,13 @@
 
         <activity android:name="com.android.compatibility.common.util.devicepolicy.provisioning.StartProvisioningActivity"/>
 
+        <service android:name="com.android.cts.deviceowner.NotificationListener"
+                 android:label="Notification Listener"
+                 android:permission="android.permission.BIND_NOTIFICATION_LISTENER_SERVICE">
+            <intent-filter>
+                <action android:name="android.service.notification.NotificationListenerService" />
+            </intent-filter>
+        </service>
     </application>
 
     <instrumentation android:name="androidx.test.runner.AndroidJUnitRunner"
diff --git a/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/AdminActionBookkeepingTest.java b/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/AdminActionBookkeepingTest.java
index 42fb4464..b793408 100644
--- a/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/AdminActionBookkeepingTest.java
+++ b/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/AdminActionBookkeepingTest.java
@@ -15,9 +15,8 @@
  */
 package com.android.cts.deviceowner;
 
-import android.app.admin.DevicePolicyManager;
+import android.app.PendingIntent;
 import android.content.ContentResolver;
-import android.content.Context;
 import android.os.Process;
 import android.provider.Settings;
 
@@ -26,6 +25,7 @@
 import java.security.cert.Certificate;
 import java.security.cert.CertificateFactory;
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
 
 public class AdminActionBookkeepingTest extends BaseDeviceOwnerTest {
     /*
@@ -107,6 +107,10 @@
      * Test: Requesting a bug report should update the corresponding timestamp.
      */
     public void testRequestBugreport() throws Exception {
+        // This test leaves a notification which will block future tests that request bug reports
+        // to fix this - we dismiss the bug report before returning
+        CountDownLatch notificationDismissedLatch = initTestRequestBugreport();
+
         Thread.sleep(1);
         final long previousTimestamp = mDevicePolicyManager.getLastBugReportRequestTime();
 
@@ -118,6 +122,34 @@
         assertTrue(newTimestamp > previousTimestamp);
         assertTrue(newTimestamp >= timeBefore);
         assertTrue(newTimestamp <= timeAfter);
+
+        cleanupTestRequestBugreport(notificationDismissedLatch);
+    }
+
+    private CountDownLatch initTestRequestBugreport() {
+        CountDownLatch notificationDismissedLatch = new CountDownLatch(1);
+        NotificationListener.getInstance().addListener((sbt) -> {
+            // The notification we are looking for is the one which confirms the bug report is
+            // ready and asks for consent to send it
+            if (sbt.getPackageName().equals("android") &&
+                    sbt.getTag().equals("DevicePolicyManager") &&
+                    sbt.getNotification().actions != null) {
+                try {
+                    // The first action is to decline
+                    sbt.getNotification().actions[0].actionIntent.send();
+                    notificationDismissedLatch.countDown();
+                } catch (PendingIntent.CanceledException e) {
+                    fail("Could not dismiss bug report notification");
+                }
+            }
+        });
+        return notificationDismissedLatch;
+    }
+
+    private void cleanupTestRequestBugreport(CountDownLatch notificationDismissedLatch)
+            throws Exception {
+        notificationDismissedLatch.await();
+        NotificationListener.getInstance().clearListeners();
     }
 
     /**
diff --git a/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/NotificationListener.java b/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/NotificationListener.java
new file mode 100644
index 0000000..f224221
--- /dev/null
+++ b/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/NotificationListener.java
@@ -0,0 +1,70 @@
+/*
+ * Copyright (C) 2019 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 com.android.cts.deviceowner;
+
+import android.service.notification.NotificationListenerService;
+import android.service.notification.StatusBarNotification;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+
+/**
+ * A {@link NotificationListenerService} that allows tests to register to respond to notifications.
+ */
+public class NotificationListener extends NotificationListenerService {
+
+    private static final int TIMEOUT_SECONDS = 120;
+
+    private static NotificationListener instance;
+    private static CountDownLatch connectedLatch = new CountDownLatch(1);
+
+    public static NotificationListener getInstance() {
+        try {
+            connectedLatch.await(TIMEOUT_SECONDS, TimeUnit.SECONDS);
+        } catch (InterruptedException e) {
+            throw new RuntimeException("NotificationListener not connected.", e);
+        }
+
+        return instance;
+    }
+
+    private List<Consumer<StatusBarNotification>> mListeners = new ArrayList<>();
+
+    public void addListener(Consumer<StatusBarNotification> listener) {
+        mListeners.add(listener);
+    }
+
+    public void clearListeners() {
+        mListeners.clear();
+    }
+
+    @Override
+    public void onNotificationPosted(StatusBarNotification sbn) {
+        for (Consumer<StatusBarNotification> listener : mListeners) {
+            listener.accept(sbn);
+        }
+    }
+
+    @Override
+    public void onListenerConnected() {
+        instance = this;
+        connectedLatch.countDown();
+    }
+}
diff --git a/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/DeviceOwnerTest.java b/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/DeviceOwnerTest.java
index 2b460ce..d223f03 100644
--- a/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/DeviceOwnerTest.java
+++ b/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/DeviceOwnerTest.java
@@ -113,6 +113,8 @@
             }
 
             getDevice().executeShellCommand(" mkdir " + TEST_UPDATE_LOCATION);
+            // Enable the notification listener
+            getDevice().executeShellCommand("cmd notification allow_listener com.android.cts.deviceowner/com.android.cts.deviceowner.NotificationListener");
         }
         mHasCreateAndManageUserFeature = mHasFeature && canCreateAdditionalUsers(1)
                 && hasDeviceFeature("android.software.managed_users");
@@ -709,23 +711,11 @@
             new DevicePolicyEventWrapper.Builder(EventId.RETRIEVE_PRE_REBOOT_SECURITY_LOGS_VALUE)
                     .setAdminPackageName(DEVICE_OWNER_PKG)
                     .build());
-
-        // Requesting a bug report (in AdminActionBookkeepingTest#testRequestBugreport) leaves a
-        // state where future bug report requests will fail
-        // TODO(b/130210665): replace this with use of NotificationListenerService to dismiss the
-        // bug report request
-        rebootAndWaitUntilReady();
-
         assertMetricsLogged(getDevice(), () -> {
             executeDeviceTestMethod(".AdminActionBookkeepingTest", "testRequestBugreport");
         }, new DevicePolicyEventWrapper.Builder(EventId.REQUEST_BUGREPORT_VALUE)
                 .setAdminPackageName(DEVICE_OWNER_PKG)
                 .build());
-        // Requesting a bug report (in AdminActionBookkeepingTest#testRequestBugreport) leaves a
-        // state where future bug report requests will fail
-        // TODO(b/130210665): replace this with use of NotificationListenerService to dismiss the
-        // bug report request
-        rebootAndWaitUntilReady();
     }
 
     public void testBluetoothRestriction() throws Exception {