Merge "Acquire WakeLocks for posting notifications" into udc-dev

5 files changed, 216 insertions, 48 deletions

diff --git a/core/java/com/android/internal/config/sysui/SystemUiSystemPropertiesFlags.java b/core/java/com/android/internal/config/sysui/SystemUiSystemPropertiesFlags.java
index 3d95dd341cc0..c9e76009136a 100644
--- a/core/java/com/android/internal/config/sysui/SystemUiSystemPropertiesFlags.java
+++ b/core/java/com/android/internal/config/sysui/SystemUiSystemPropertiesFlags.java
@@ -81,6 +81,10 @@ public class SystemUiSystemPropertiesFlags {
         /** Gating the logging of DND state change events. */
         public static final Flag LOG_DND_STATE_EVENTS =
                 releasedFlag("persist.sysui.notification.log_dnd_state_events");
+
+        /** Gating the holding of WakeLocks until NLSes are told about a new notification. */
+        public static final Flag WAKE_LOCK_FOR_POSTING_NOTIFICATION =
+                devFlag("persist.sysui.notification.wake_lock_for_posting_notification");
     }
 
     //// == End of flags.  Everything below this line is the implementation. == ////
diff --git a/services/core/java/com/android/server/notification/NotificationManagerService.java b/services/core/java/com/android/server/notification/NotificationManagerService.java
index aeff2b0f9af6..b39e8606e189 100644
--- a/services/core/java/com/android/server/notification/NotificationManagerService.java
+++ b/services/core/java/com/android/server/notification/NotificationManagerService.java
@@ -119,6 +119,7 @@ import static android.service.notification.NotificationListenerService.TRIM_LIGH
 import static android.view.WindowManager.LayoutParams.TYPE_TOAST;
 
 import static com.android.internal.config.sysui.SystemUiSystemPropertiesFlags.NotificationFlags.ALLOW_DISMISS_ONGOING;
+import static com.android.internal.config.sysui.SystemUiSystemPropertiesFlags.NotificationFlags.WAKE_LOCK_FOR_POSTING_NOTIFICATION;
 import static com.android.internal.util.FrameworkStatsLog.DND_MODE_RULE;
 import static com.android.internal.util.FrameworkStatsLog.PACKAGE_NOTIFICATION_CHANNEL_GROUP_PREFERENCES;
 import static com.android.internal.util.FrameworkStatsLog.PACKAGE_NOTIFICATION_CHANNEL_PREFERENCES;
@@ -223,6 +224,8 @@ import android.os.IInterface;
 import android.os.Looper;
 import android.os.Message;
 import android.os.ParcelFileDescriptor;
+import android.os.PowerManager;
+import android.os.PowerManager.WakeLock;
 import android.os.Process;
 import android.os.RemoteException;
 import android.os.ResultReceiver;
@@ -234,6 +237,7 @@ import android.os.Trace;
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.os.VibrationEffect;
+import android.os.WorkSource;
 import android.permission.PermissionManager;
 import android.provider.DeviceConfig;
 import android.provider.Settings;
@@ -559,6 +563,7 @@ public class NotificationManagerService extends SystemService {
     private PermissionHelper mPermissionHelper;
     private UsageStatsManagerInternal mUsageStatsManagerInternal;
     private TelecomManager mTelecomManager;
+    private PowerManager mPowerManager;
     private PostNotificationTrackerFactory mPostNotificationTrackerFactory;
 
     final IBinder mForegroundToken = new Binder();
@@ -923,7 +928,7 @@ public class NotificationManagerService extends SystemService {
         if (oldFlags != flags) {
             summary.getSbn().getNotification().flags = flags;
             mHandler.post(new EnqueueNotificationRunnable(userId, summary, isAppForeground,
-                    mPostNotificationTrackerFactory.newTracker()));
+                    mPostNotificationTrackerFactory.newTracker(null)));
         }
     }
 
@@ -1457,7 +1462,7 @@ public class NotificationManagerService extends SystemService {
                         // want to adjust the flag behaviour.
                         mHandler.post(new EnqueueNotificationRunnable(r.getUser().getIdentifier(),
                                 r, true /* isAppForeground*/,
-                                mPostNotificationTrackerFactory.newTracker()));
+                                mPostNotificationTrackerFactory.newTracker(null)));
                     }
                 }
             }
@@ -1488,7 +1493,7 @@ public class NotificationManagerService extends SystemService {
                         mHandler.post(
                                 new EnqueueNotificationRunnable(r.getUser().getIdentifier(), r,
                                         /* foreground= */ true,
-                                        mPostNotificationTrackerFactory.newTracker()));
+                                        mPostNotificationTrackerFactory.newTracker(null)));
                     }
                 }
             }
@@ -2233,7 +2238,7 @@ public class NotificationManagerService extends SystemService {
             UsageStatsManagerInternal usageStatsManagerInternal,
             TelecomManager telecomManager, NotificationChannelLogger channelLogger,
             SystemUiSystemPropertiesFlags.FlagResolver flagResolver,
-            PermissionManager permissionManager,
+            PermissionManager permissionManager, PowerManager powerManager,
             PostNotificationTrackerFactory postNotificationTrackerFactory) {
         mHandler = handler;
         Resources resources = getContext().getResources();
@@ -2265,6 +2270,7 @@ public class NotificationManagerService extends SystemService {
         mDpm = dpm;
         mUm = userManager;
         mTelecomManager = telecomManager;
+        mPowerManager = powerManager;
         mPostNotificationTrackerFactory = postNotificationTrackerFactory;
         mPlatformCompat = IPlatformCompat.Stub.asInterface(
                 ServiceManager.getService(Context.PLATFORM_COMPAT_SERVICE));
@@ -2568,6 +2574,7 @@ public class NotificationManagerService extends SystemService {
                 getContext().getSystemService(TelecomManager.class),
                 new NotificationChannelLoggerImpl(), SystemUiSystemPropertiesFlags.getResolver(),
                 getContext().getSystemService(PermissionManager.class),
+                getContext().getSystemService(PowerManager.class),
                 new PostNotificationTrackerFactory() {});
 
         publishBinderService(Context.NOTIFICATION_SERVICE, mService, /* allowIsolated= */ false,
@@ -2684,7 +2691,7 @@ public class NotificationManagerService extends SystemService {
                     final boolean isAppForeground =
                             mActivityManager.getPackageImportance(pkg) == IMPORTANCE_FOREGROUND;
                     mHandler.post(new EnqueueNotificationRunnable(userId, r, isAppForeground,
-                            mPostNotificationTrackerFactory.newTracker()));
+                            mPostNotificationTrackerFactory.newTracker(null)));
                 }
             }
 
@@ -6577,7 +6584,7 @@ public class NotificationManagerService extends SystemService {
     void enqueueNotificationInternal(final String pkg, final String opPkg, final int callingUid,
             final int callingPid, final String tag, final int id, final Notification notification,
             int incomingUserId, boolean postSilently) {
-        PostNotificationTracker tracker = mPostNotificationTrackerFactory.newTracker();
+        PostNotificationTracker tracker = acquireWakeLockForPost(pkg, callingUid);
         boolean enqueued = false;
         try {
             enqueued = enqueueNotificationInternal(pkg, opPkg, callingUid, callingPid, tag, id,
@@ -6589,6 +6596,22 @@ public class NotificationManagerService extends SystemService {
         }
     }
 
+    private PostNotificationTracker acquireWakeLockForPost(String pkg, int uid) {
+        if (mFlagResolver.isEnabled(WAKE_LOCK_FOR_POSTING_NOTIFICATION)) {
+            // The package probably doesn't have WAKE_LOCK permission and should not require it.
+            return Binder.withCleanCallingIdentity(() -> {
+                WakeLock wakeLock = mPowerManager.newWakeLock(PowerManager.PARTIAL_WAKE_LOCK,
+                        "NotificationManagerService:post:" + pkg);
+                wakeLock.setWorkSource(new WorkSource(uid, pkg));
+                // TODO(b/275044361): Adjust to a more reasonable number when we have the data.
+                wakeLock.acquire(30_000);
+                return mPostNotificationTrackerFactory.newTracker(wakeLock);
+            });
+        } else {
+            return mPostNotificationTrackerFactory.newTracker(null);
+        }
+    }
+
     /**
      * @return True if we successfully processed the notification and handed off the task of
      * enqueueing it to a background thread; false otherwise.
@@ -7106,7 +7129,7 @@ public class NotificationManagerService extends SystemService {
                             mHandler.post(
                                     new NotificationManagerService.EnqueueNotificationRunnable(
                                             r.getUser().getIdentifier(), r, isAppForeground,
-                                            mPostNotificationTrackerFactory.newTracker()));
+                                            mPostNotificationTrackerFactory.newTracker(null)));
                         }
                     }
                 }
@@ -12168,20 +12191,20 @@ public class NotificationManagerService extends SystemService {
     }
 
     interface PostNotificationTrackerFactory {
-        default PostNotificationTracker newTracker() {
-            return new PostNotificationTracker();
+        default PostNotificationTracker newTracker(@Nullable WakeLock optionalWakelock) {
+            return new PostNotificationTracker(optionalWakelock);
         }
     }
 
     static class PostNotificationTracker {
         @ElapsedRealtimeLong private final long mStartTime;
-        @Nullable private NotificationRecordLogger.NotificationReported mReport;
+        @Nullable private final WakeLock mWakeLock;
         private boolean mOngoing;
 
         @VisibleForTesting
-        PostNotificationTracker() {
-            // TODO(b/275044361): (Conditionally) receive a wakelock.
+        PostNotificationTracker(@Nullable WakeLock wakeLock) {
             mStartTime = SystemClock.elapsedRealtime();
+            mWakeLock = wakeLock;
             mOngoing = true;
             if (DBG) {
                 Slog.d(TAG, "PostNotification: Started");
@@ -12199,9 +12222,8 @@ public class NotificationManagerService extends SystemService {
         }
 
         /**
-         * Cancels the tracker (TODO(b/275044361): releasing the acquired WakeLock). Either
-         * {@link #finish} or {@link #cancel} (exclusively) should be called on this object before
-         * it's discarded.
+         * Cancels the tracker (releasing the acquired WakeLock). Either {@link #finish} or
+         * {@link #cancel} (exclusively) should be called on this object before it's discarded.
          */
         void cancel() {
             if (!isOngoing()) {
@@ -12209,9 +12231,9 @@ public class NotificationManagerService extends SystemService {
                 return;
             }
             mOngoing = false;
-
-            // TODO(b/275044361): Release wakelock.
-
+            if (mWakeLock != null) {
+                Binder.withCleanCallingIdentity(() -> mWakeLock.release());
+            }
             if (DBG) {
                 long elapsedTime = SystemClock.elapsedRealtime() - mStartTime;
                 Slog.d(TAG, TextUtils.formatSimple("PostNotification: Abandoned after %d ms",
@@ -12220,9 +12242,9 @@ public class NotificationManagerService extends SystemService {
         }
 
         /**
-         * Finishes the tracker (TODO(b/275044361): releasing the acquired WakeLock) and returns the
-         * time elapsed since the operation started, in milliseconds. Either {@link #finish} or
-         * {@link #cancel} (exclusively) should be called on this object before it's discarded.
+         * Finishes the tracker (releasing the acquired WakeLock) and returns the time elapsed since
+         * the operation started, in milliseconds. Either {@link #finish} or {@link #cancel}
+         * (exclusively) should be called on this object before it's discarded.
          */
         @DurationMillisLong
         long finish() {
@@ -12232,9 +12254,9 @@ public class NotificationManagerService extends SystemService {
                 return elapsedTime;
             }
             mOngoing = false;
-
-            // TODO(b/275044361): Release wakelock.
-
+            if (mWakeLock != null) {
+                Binder.withCleanCallingIdentity(() -> mWakeLock.release());
+            }
             if (DBG) {
                 Slog.d(TAG,
                         TextUtils.formatSimple("PostNotification: Finished in %d ms", elapsedTime));
diff --git a/services/tests/uiservicestests/AndroidManifest.xml b/services/tests/uiservicestests/AndroidManifest.xml
index f44c1d18614d..4315254f68a9 100644
--- a/services/tests/uiservicestests/AndroidManifest.xml
+++ b/services/tests/uiservicestests/AndroidManifest.xml
@@ -20,6 +20,7 @@
     <uses-permission android:name="android.permission.INTERACT_ACROSS_USERS" />
     <uses-permission android:name="android.permission.INTERACT_ACROSS_USERS_FULL" />
     <uses-permission android:name="android.permission.UPDATE_APP_OPS_STATS" />
+    <uses-permission android:name="android.permission.UPDATE_DEVICE_STATS" />
     <uses-permission android:name="android.permission.MANAGE_USERS" />
     <uses-permission android:name="android.permission.WRITE_SECURE_SETTINGS" />
     <uses-permission android:name="android.permission.ACCESS_NOTIFICATIONS" />
@@ -36,6 +37,7 @@
     <uses-permission android:name="android.permission.WRITE_ALLOWLISTED_DEVICE_CONFIG" />
     <uses-permission android:name="android.permission.READ_WRITE_SYNC_DISABLED_MODE_CONFIG" />
     <uses-permission android:name="android.permission.ACCESS_KEYGUARD_SECURE_STORAGE" />
+    <uses-permission android:name="android.permission.WAKE_LOCK" />
 
     <application android:debuggable="true">
         <uses-library android:name="android.test.runner" />
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java b/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
index 9166b3d75f60..2a0c745a0ffc 100755
--- a/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
@@ -61,6 +61,7 @@ import static android.content.pm.PackageManager.PERMISSION_DENIED;
 import static android.content.pm.PackageManager.PERMISSION_GRANTED;
 import static android.os.Build.VERSION_CODES.O_MR1;
 import static android.os.Build.VERSION_CODES.P;
+import static android.os.PowerManager.PARTIAL_WAKE_LOCK;
 import static android.os.UserHandle.USER_SYSTEM;
 import static android.os.UserManager.USER_TYPE_FULL_SECONDARY;
 import static android.os.UserManager.USER_TYPE_PROFILE_CLONE;
@@ -80,6 +81,7 @@ import static android.view.WindowManager.LayoutParams.TYPE_TOAST;
 import static com.android.internal.config.sysui.SystemUiSystemPropertiesFlags.NotificationFlags.ALLOW_DISMISS_ONGOING;
 import static com.android.internal.config.sysui.SystemUiSystemPropertiesFlags.NotificationFlags.FSI_FORCE_DEMOTE;
 import static com.android.internal.config.sysui.SystemUiSystemPropertiesFlags.NotificationFlags.SHOW_STICKY_HUN_FOR_DENIED_FSI;
+import static com.android.internal.config.sysui.SystemUiSystemPropertiesFlags.NotificationFlags.WAKE_LOCK_FOR_POSTING_NOTIFICATION;
 import static com.android.internal.widget.LockPatternUtils.StrongAuthTracker.STRONG_AUTH_REQUIRED_AFTER_USER_LOCKDOWN;
 import static com.android.server.notification.NotificationRecordLogger.NotificationReportedEvent.NOTIFICATION_ADJUSTED;
 import static com.android.server.notification.NotificationRecordLogger.NotificationReportedEvent.NOTIFICATION_POSTED;
@@ -119,12 +121,14 @@ import static org.mockito.Mockito.timeout;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.verifyZeroInteractions;
 import static org.mockito.Mockito.when;
 
 import static java.util.Collections.emptyList;
 import static java.util.Collections.singletonList;
 
 import android.Manifest;
+import android.annotation.Nullable;
 import android.annotation.SuppressLint;
 import android.annotation.UserIdInt;
 import android.app.ActivityManager;
@@ -181,11 +185,14 @@ import android.os.Bundle;
 import android.os.IBinder;
 import android.os.Looper;
 import android.os.Parcel;
+import android.os.PowerManager;
+import android.os.PowerManager.WakeLock;
 import android.os.Process;
 import android.os.RemoteException;
 import android.os.SystemClock;
 import android.os.UserHandle;
 import android.os.UserManager;
+import android.os.WorkSource;
 import android.permission.PermissionManager;
 import android.provider.DeviceConfig;
 import android.provider.MediaStore;
@@ -351,6 +358,9 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
     private PermissionManager mPermissionManager;
     @Mock
     private DevicePolicyManagerInternal mDevicePolicyManager;
+    @Mock
+    private PowerManager mPowerManager;
+    private final ArrayList<WakeLock> mAcquiredWakeLocks = new ArrayList<>();
     private final TestPostNotificationTrackerFactory mPostNotificationTrackerFactory =
             new TestPostNotificationTrackerFactory();
 
@@ -431,8 +441,9 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
         private final List<PostNotificationTracker> mCreatedTrackers = new ArrayList<>();
 
         @Override
-        public PostNotificationTracker newTracker() {
-            PostNotificationTracker tracker = PostNotificationTrackerFactory.super.newTracker();
+        public PostNotificationTracker newTracker(@Nullable WakeLock optionalWakeLock) {
+            PostNotificationTracker tracker = PostNotificationTrackerFactory.super.newTracker(
+                    optionalWakeLock);
             mCreatedTrackers.add(tracker);
             return tracker;
         }
@@ -563,6 +574,20 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
 
         when(mAssistants.isAdjustmentAllowed(anyString())).thenReturn(true);
 
+        // Use the real PowerManager to back up the mock w.r.t. creating WakeLocks.
+        // This is because 1) we need a mock to verify() calls and tracking the created WakeLocks,
+        // but 2) PowerManager and WakeLock perform their own checks (e.g. correct arguments, don't
+        // call release twice, etc) and we want the test to fail if such misuse happens, too.
+        PowerManager realPowerManager = mContext.getSystemService(PowerManager.class);
+        when(mPowerManager.newWakeLock(anyInt(), anyString())).then(
+                (Answer<WakeLock>) invocation -> {
+                    WakeLock wl = realPowerManager.newWakeLock(invocation.getArgument(0),
+                            invocation.getArgument(1));
+                    mAcquiredWakeLocks.add(wl);
+                    return wl;
+                });
+        mTestFlagResolver.setFlagOverride(WAKE_LOCK_FOR_POSTING_NOTIFICATION, true);
+
         // apps allowed as convos
         mService.setStringArrayResourceValue(PKG_O);
 
@@ -579,7 +604,7 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
                 mock(TelephonyManager.class),
                 mAmi, mToastRateLimiter, mPermissionHelper, mock(UsageStatsManagerInternal.class),
                 mTelecomManager, mLogger, mTestFlagResolver, mPermissionManager,
-                mPostNotificationTrackerFactory);
+                mPowerManager, mPostNotificationTrackerFactory);
         // Return first true for RoleObserver main-thread check
         when(mMainLooper.isCurrentThread()).thenReturn(true).thenReturn(false);
         mService.onBootPhase(SystemService.PHASE_SYSTEM_SERVICES_READY, mMainLooper);
@@ -686,6 +711,13 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
     }
 
     @After
+    public void assertAllWakeLocksReleased() {
+        for (WakeLock wakeLock : mAcquiredWakeLocks) {
+            assertThat(wakeLock.isHeld()).isFalse();
+        }
+    }
+
+    @After
     public void tearDown() throws Exception {
         if (mFile != null) mFile.delete();
         clearDeviceConfig();
@@ -1486,7 +1518,7 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
 
         NotificationManagerService.PostNotificationRunnable runnable =
                 mService.new PostNotificationRunnable(r.getKey(), r.getSbn().getPackageName(),
-                        r.getUid(), mPostNotificationTrackerFactory.newTracker());
+                        r.getUid(), mPostNotificationTrackerFactory.newTracker(null));
         runnable.run();
         waitForIdle();
 
@@ -1507,7 +1539,7 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
 
         NotificationManagerService.PostNotificationRunnable runnable =
                 mService.new PostNotificationRunnable(r.getKey(), r.getSbn().getPackageName(),
-                        r.getUid(), mPostNotificationTrackerFactory.newTracker());
+                        r.getUid(), mPostNotificationTrackerFactory.newTracker(null));
         runnable.run();
         waitForIdle();
 
@@ -1803,6 +1835,112 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
     }
 
     @Test
+    public void enqueueNotification_acquiresAndReleasesWakeLock() throws Exception {
+        mBinderService.enqueueNotificationWithTag(PKG, PKG,
+                "enqueueNotification_acquiresAndReleasesWakeLock", 0,
+                generateNotificationRecord(null).getNotification(), 0);
+
+        verify(mPowerManager).newWakeLock(eq(PARTIAL_WAKE_LOCK), anyString());
+        assertThat(mAcquiredWakeLocks).hasSize(1);
+        assertThat(mAcquiredWakeLocks.get(0).isHeld()).isTrue();
+
+        waitForIdle();
+
+        assertThat(mAcquiredWakeLocks).hasSize(1);
+        assertThat(mAcquiredWakeLocks.get(0).isHeld()).isFalse();
+    }
+
+    @Test
+    public void enqueueNotification_throws_acquiresAndReleasesWakeLock() throws Exception {
+        // Simulate not enqueued due to rejected inputs.
+        assertThrows(Exception.class,
+                () -> mBinderService.enqueueNotificationWithTag(PKG, PKG,
+                        "enqueueNotification_throws_acquiresAndReleasesWakeLock", 0,
+                        /* notification= */ null, 0));
+
+        verify(mPowerManager).newWakeLock(eq(PARTIAL_WAKE_LOCK), anyString());
+        assertThat(mAcquiredWakeLocks).hasSize(1);
+        assertThat(mAcquiredWakeLocks.get(0).isHeld()).isFalse();
+    }
+
+    @Test
+    public void enqueueNotification_notEnqueued_acquiresAndReleasesWakeLock() throws Exception {
+        // Simulate not enqueued due to snoozing inputs.
+        when(mSnoozeHelper.getSnoozeContextForUnpostedNotification(anyInt(), any(), any()))
+                .thenReturn("zzzzzzz");
+
+        mBinderService.enqueueNotificationWithTag(PKG, PKG,
+                "enqueueNotification_notEnqueued_acquiresAndReleasesWakeLock", 0,
+                generateNotificationRecord(null).getNotification(), 0);
+
+        verify(mPowerManager).newWakeLock(eq(PARTIAL_WAKE_LOCK), anyString());
+        assertThat(mAcquiredWakeLocks).hasSize(1);
+        assertThat(mAcquiredWakeLocks.get(0).isHeld()).isTrue();
+
+        waitForIdle();
+
+        assertThat(mAcquiredWakeLocks).hasSize(1);
+        assertThat(mAcquiredWakeLocks.get(0).isHeld()).isFalse();
+    }
+
+    @Test
+    public void enqueueNotification_notPosted_acquiresAndReleasesWakeLock() throws Exception {
+        // Simulate enqueued but not posted due to missing small icon.
+        Notification notif = new Notification.Builder(mContext, mTestNotificationChannel.getId())
+                .setContentTitle("foo")
+                .build();
+
+        mBinderService.enqueueNotificationWithTag(PKG, PKG,
+                "enqueueNotification_notPosted_acquiresAndReleasesWakeLock", 0,
+                notif, 0);
+
+        verify(mPowerManager).newWakeLock(eq(PARTIAL_WAKE_LOCK), anyString());
+        assertThat(mAcquiredWakeLocks).hasSize(1);
+        assertThat(mAcquiredWakeLocks.get(0).isHeld()).isTrue();
+
+        waitForIdle();
+
+        // NLSes were not called.
+        verify(mListeners, never()).prepareNotifyPostedLocked(any(), any(), anyBoolean());
+
+        assertThat(mAcquiredWakeLocks).hasSize(1);
+        assertThat(mAcquiredWakeLocks.get(0).isHeld()).isFalse();
+    }
+
+    @Test
+    public void enqueueNotification_setsWakeLockWorkSource() throws Exception {
+        // Use a "full" mock for the PowerManager (instead of the one that delegates to the real
+        // service) so we can return a mocked WakeLock that we can verify() on.
+        reset(mPowerManager);
+        WakeLock wakeLock = mock(WakeLock.class);
+        when(mPowerManager.newWakeLock(anyInt(), anyString())).thenReturn(wakeLock);
+
+        mBinderService.enqueueNotificationWithTag(PKG, PKG,
+                "enqueueNotification_setsWakeLockWorkSource", 0,
+                generateNotificationRecord(null).getNotification(), 0);
+        waitForIdle();
+
+        InOrder inOrder = inOrder(mPowerManager, wakeLock);
+        inOrder.verify(mPowerManager).newWakeLock(eq(PARTIAL_WAKE_LOCK), anyString());
+        inOrder.verify(wakeLock).setWorkSource(eq(new WorkSource(mUid, PKG)));
+        inOrder.verify(wakeLock).acquire(anyLong());
+        inOrder.verify(wakeLock).release();
+        inOrder.verifyNoMoreInteractions();
+    }
+
+    @Test
+    public void enqueueNotification_wakeLockFlagOff_noWakeLock() throws Exception {
+        mTestFlagResolver.setFlagOverride(WAKE_LOCK_FOR_POSTING_NOTIFICATION, false);
+
+        mBinderService.enqueueNotificationWithTag(PKG, PKG,
+                "enqueueNotification_setsWakeLockWorkSource", 0,
+                generateNotificationRecord(null).getNotification(), 0);
+        waitForIdle();
+
+        verifyZeroInteractions(mPowerManager);
+    }
+
+    @Test
     public void testCancelNonexistentNotification() throws Exception {
         mBinderService.cancelNotificationWithTag(PKG, PKG,
                 "testCancelNonexistentNotification", 0, 0);
@@ -4361,7 +4499,7 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
         mService.addEnqueuedNotification(r);
         NotificationManagerService.PostNotificationRunnable runnable =
                 mService.new PostNotificationRunnable(r.getKey(), r.getSbn().getPackageName(),
-                        r.getUid(), mPostNotificationTrackerFactory.newTracker());
+                        r.getUid(), mPostNotificationTrackerFactory.newTracker(null));
         runnable.run();
         waitForIdle();
 
@@ -4380,7 +4518,7 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
         NotificationManagerService.PostNotificationRunnable runnable =
                 mService.new PostNotificationRunnable(update.getKey(),
                         update.getSbn().getPackageName(), update.getUid(),
-                        mPostNotificationTrackerFactory.newTracker());
+                        mPostNotificationTrackerFactory.newTracker(null));
         runnable.run();
         waitForIdle();
 
@@ -4400,7 +4538,7 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
         NotificationManagerService.PostNotificationRunnable runnable =
                 mService.new PostNotificationRunnable(update.getKey(),
                         update.getSbn().getPackageName(), update.getUid(),
-                        mPostNotificationTrackerFactory.newTracker());
+                        mPostNotificationTrackerFactory.newTracker(null));
         runnable.run();
         waitForIdle();
 
@@ -4420,7 +4558,7 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
         NotificationManagerService.PostNotificationRunnable runnable =
                 mService.new PostNotificationRunnable(update.getKey(),
                         update.getSbn().getPackageName(),
-                        update.getUid(), mPostNotificationTrackerFactory.newTracker());
+                        update.getUid(), mPostNotificationTrackerFactory.newTracker(null));
         runnable.run();
         waitForIdle();
 
@@ -4434,13 +4572,13 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
         mService.addEnqueuedNotification(r);
         NotificationManagerService.PostNotificationRunnable runnable =
                 mService.new PostNotificationRunnable(r.getKey(), r.getSbn().getPackageName(),
-                        r.getUid(), mPostNotificationTrackerFactory.newTracker());
+                        r.getUid(), mPostNotificationTrackerFactory.newTracker(null));
         runnable.run();
 
         r = generateNotificationRecord(mTestNotificationChannel, 1, null, false);
         r.setCriticality(CriticalNotificationExtractor.CRITICAL);
         runnable = mService.new PostNotificationRunnable(r.getKey(), r.getSbn().getPackageName(),
-                r.getUid(), mPostNotificationTrackerFactory.newTracker());
+                r.getUid(), mPostNotificationTrackerFactory.newTracker(null));
         mService.addEnqueuedNotification(r);
 
         runnable.run();
@@ -5090,7 +5228,7 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
                 mService.new PostNotificationRunnable(original.getKey(),
                         original.getSbn().getPackageName(),
                         original.getUid(),
-                        mPostNotificationTrackerFactory.newTracker());
+                        mPostNotificationTrackerFactory.newTracker(null));
         runnable.run();
         waitForIdle();
 
@@ -5114,7 +5252,7 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
                 mService.new PostNotificationRunnable(update.getKey(),
                         update.getSbn().getPackageName(),
                         update.getUid(),
-                        mPostNotificationTrackerFactory.newTracker());
+                        mPostNotificationTrackerFactory.newTracker(null));
         runnable.run();
         waitForIdle();
 
@@ -7536,7 +7674,7 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
 
         NotificationManagerService.PostNotificationRunnable runnable =
                 mService.new PostNotificationRunnable(update.getKey(), r.getSbn().getPackageName(),
-                        r.getUid(), mPostNotificationTrackerFactory.newTracker());
+                        r.getUid(), mPostNotificationTrackerFactory.newTracker(null));
         runnable.run();
         waitForIdle();
 
@@ -10234,7 +10372,7 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
         mService.addEnqueuedNotification(r);
         NotificationManagerService.PostNotificationRunnable runnable =
                 mService.new PostNotificationRunnable(r.getKey(), r.getSbn().getPackageName(),
-                        r.getUid(), mPostNotificationTrackerFactory.newTracker());
+                        r.getUid(), mPostNotificationTrackerFactory.newTracker(null));
         runnable.run();
         waitForIdle();
 
@@ -10251,7 +10389,7 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
 
         mService.addEnqueuedNotification(r);
         runnable = mService.new PostNotificationRunnable(r.getKey(), r.getSbn().getPackageName(),
-                r.getUid(), mPostNotificationTrackerFactory.newTracker());
+                r.getUid(), mPostNotificationTrackerFactory.newTracker(null));
         runnable.run();
         waitForIdle();
 
@@ -10268,7 +10406,7 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
 
         mService.addEnqueuedNotification(r);
         runnable = mService.new PostNotificationRunnable(r.getKey(), r.getSbn().getPackageName(),
-                r.getUid(), mPostNotificationTrackerFactory.newTracker());
+                r.getUid(), mPostNotificationTrackerFactory.newTracker(null));
         runnable.run();
         waitForIdle();
 
@@ -10361,7 +10499,7 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
         // normal blocked notifications - blocked
         mService.addEnqueuedNotification(r);
         mService.new PostNotificationRunnable(r.getKey(), r.getSbn().getPackageName(), r.getUid(),
-                mPostNotificationTrackerFactory.newTracker()).run();
+                mPostNotificationTrackerFactory.newTracker(null)).run();
         waitForIdle();
 
         verify(mUsageStats).registerBlocked(any());
@@ -10379,7 +10517,7 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
 
         mService.addEnqueuedNotification(r);
         mService.new PostNotificationRunnable(r.getKey(), r.getSbn().getPackageName(), r.getUid(),
-                mPostNotificationTrackerFactory.newTracker()).run();
+                mPostNotificationTrackerFactory.newTracker(null)).run();
         waitForIdle();
 
         verify(mUsageStats).registerBlocked(any());
@@ -10392,7 +10530,7 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
 
         mService.addEnqueuedNotification(r);
         mService.new PostNotificationRunnable(r.getKey(), r.getSbn().getPackageName(), r.getUid(),
-                mPostNotificationTrackerFactory.newTracker()).run();
+                mPostNotificationTrackerFactory.newTracker(null)).run();
         waitForIdle();
 
         verify(mUsageStats, never()).registerBlocked(any());
@@ -10406,7 +10544,7 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
 
         mService.addEnqueuedNotification(r);
         mService.new PostNotificationRunnable(r.getKey(), r.getSbn().getPackageName(), r.getUid(),
-                mPostNotificationTrackerFactory.newTracker()).run();
+                mPostNotificationTrackerFactory.newTracker(null)).run();
         waitForIdle();
 
         verify(mUsageStats, never()).registerBlocked(any());
@@ -10420,7 +10558,7 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
 
         mService.addEnqueuedNotification(r);
         mService.new PostNotificationRunnable(r.getKey(), r.getSbn().getPackageName(), r.getUid(),
-                mPostNotificationTrackerFactory.newTracker()).run();
+                mPostNotificationTrackerFactory.newTracker(null)).run();
         waitForIdle();
 
         verify(mUsageStats).registerBlocked(any());
@@ -10435,7 +10573,7 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
 
         mService.addEnqueuedNotification(r);
         mService.new PostNotificationRunnable(r.getKey(), r.getSbn().getPackageName(), r.getUid(),
-                mPostNotificationTrackerFactory.newTracker()).run();
+                mPostNotificationTrackerFactory.newTracker(null)).run();
         waitForIdle();
 
         verify(mUsageStats).registerBlocked(any());
@@ -10448,7 +10586,7 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
 
         mService.addEnqueuedNotification(r);
         mService.new PostNotificationRunnable(r.getKey(), r.getSbn().getPackageName(), r.getUid(),
-                mPostNotificationTrackerFactory.newTracker()).run();
+                mPostNotificationTrackerFactory.newTracker(null)).run();
         waitForIdle();
 
         verify(mUsageStats).registerBlocked(any());
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/RoleObserverTest.java b/services/tests/uiservicestests/src/com/android/server/notification/RoleObserverTest.java
index 66c1e35754c5..81c573d8fb1e 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/RoleObserverTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/RoleObserverTest.java
@@ -48,6 +48,7 @@ import android.content.Context;
 import android.content.pm.IPackageManager;
 import android.content.pm.PackageManager;
 import android.os.Looper;
+import android.os.PowerManager;
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.permission.PermissionManager;
@@ -169,6 +170,7 @@ public class RoleObserverTest extends UiServiceTestCase {
                     mock(UsageStatsManagerInternal.class), mock(TelecomManager.class),
                     mock(NotificationChannelLogger.class), new TestableFlagResolver(),
                     mock(PermissionManager.class),
+                    mock(PowerManager.class),
                     new NotificationManagerService.PostNotificationTrackerFactory() {});
         } catch (SecurityException e) {
             if (!e.getMessage().contains("Permission Denial: not allowed to send broadcast")) {