Merge "[MEDIA TTT] Apply wakelock for chipbar" into tm-qpr-dev

10 files changed, 188 insertions, 25 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/media/taptotransfer/receiver/MediaTttChipControllerReceiver.kt b/packages/SystemUI/src/com/android/systemui/media/taptotransfer/receiver/MediaTttChipControllerReceiver.kt
index 662d059aedcd..8bddffc842f5 100644
--- a/packages/SystemUI/src/com/android/systemui/media/taptotransfer/receiver/MediaTttChipControllerReceiver.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/taptotransfer/receiver/MediaTttChipControllerReceiver.kt
@@ -45,6 +45,7 @@ import com.android.systemui.temporarydisplay.TemporaryViewInfo
 import com.android.systemui.util.animation.AnimationUtil.Companion.frames
 import com.android.systemui.util.concurrency.DelayableExecutor
 import com.android.systemui.util.view.ViewUtil
+import com.android.systemui.util.wakelock.WakeLock
 import javax.inject.Inject
 
 /**
@@ -68,6 +69,7 @@ class MediaTttChipControllerReceiver @Inject constructor(
         private val mediaTttFlags: MediaTttFlags,
         private val uiEventLogger: MediaTttReceiverUiEventLogger,
         private val viewUtil: ViewUtil,
+        wakeLockBuilder: WakeLock.Builder,
 ) : TemporaryViewDisplayController<ChipReceiverInfo, MediaTttLogger>(
         context,
         logger,
@@ -77,6 +79,7 @@ class MediaTttChipControllerReceiver @Inject constructor(
         configurationController,
         powerManager,
         R.layout.media_ttt_chip_receiver,
+        wakeLockBuilder,
 ) {
     @SuppressLint("WrongConstant") // We're allowed to use LAYOUT_IN_DISPLAY_CUTOUT_MODE_ALWAYS
     override val windowLayoutParams = commonWindowLayoutParams.apply {
diff --git a/packages/SystemUI/src/com/android/systemui/temporarydisplay/TemporaryViewDisplayController.kt b/packages/SystemUI/src/com/android/systemui/temporarydisplay/TemporaryViewDisplayController.kt
index 637fac05f0b6..4cb41f3a977e 100644
--- a/packages/SystemUI/src/com/android/systemui/temporarydisplay/TemporaryViewDisplayController.kt
+++ b/packages/SystemUI/src/com/android/systemui/temporarydisplay/TemporaryViewDisplayController.kt
@@ -22,7 +22,6 @@ import android.graphics.PixelFormat
 import android.graphics.Rect
 import android.graphics.drawable.Drawable
 import android.os.PowerManager
-import android.os.SystemClock
 import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
@@ -35,6 +34,7 @@ import com.android.systemui.CoreStartable
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.statusbar.policy.ConfigurationController
 import com.android.systemui.util.concurrency.DelayableExecutor
+import com.android.systemui.util.wakelock.WakeLock
 
 /**
  * A generic controller that can temporarily display a new view in a new window.
@@ -54,6 +54,7 @@ abstract class TemporaryViewDisplayController<T : TemporaryViewInfo, U : Tempora
     private val configurationController: ConfigurationController,
     private val powerManager: PowerManager,
     @LayoutRes private val viewLayoutRes: Int,
+    private val wakeLockBuilder: WakeLock.Builder,
 ) : CoreStartable {
     /**
      * Window layout params that will be used as a starting point for the [windowLayoutParams] of
@@ -64,7 +65,8 @@ abstract class TemporaryViewDisplayController<T : TemporaryViewInfo, U : Tempora
         height = WindowManager.LayoutParams.WRAP_CONTENT
         type = WindowManager.LayoutParams.TYPE_SYSTEM_ERROR
         flags = WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE or
-            WindowManager.LayoutParams.FLAG_NOT_TOUCH_MODAL
+            WindowManager.LayoutParams.FLAG_NOT_TOUCH_MODAL or
+            WindowManager.LayoutParams.FLAG_KEEP_SCREEN_ON
         format = PixelFormat.TRANSLUCENT
         setTrustedOverlay()
     }
@@ -84,6 +86,15 @@ abstract class TemporaryViewDisplayController<T : TemporaryViewInfo, U : Tempora
     private var cancelViewTimeout: Runnable? = null
 
     /**
+     * A wakelock that is acquired when view is displayed and screen off,
+     * then released when view is removed.
+     */
+    private var wakeLock: WakeLock? = null
+
+    /** A string that keeps track of wakelock reason once it is acquired till it gets released */
+    private var wakeReasonAcquired: String? = null
+
+    /**
      * Displays the view with the provided [newInfo].
      *
      * This method handles inflating and attaching the view, then delegates to [updateView] to
@@ -113,11 +124,15 @@ abstract class TemporaryViewDisplayController<T : TemporaryViewInfo, U : Tempora
             // the view to show over the dream state, so we should only wake up if the screen is
             // completely off.)
             if (!powerManager.isScreenOn) {
-                powerManager.wakeUp(
-                    SystemClock.uptimeMillis(),
-                    PowerManager.WAKE_REASON_APPLICATION,
-                    "com.android.systemui:${newInfo.wakeReason}",
-                )
+                wakeLock = wakeLockBuilder
+                    .setTag(newInfo.windowTitle)
+                    .setLevelsAndFlags(
+                        PowerManager.FULL_WAKE_LOCK or
+                        PowerManager.ACQUIRE_CAUSES_WAKEUP
+                    )
+                    .build()
+                wakeLock?.acquire(newInfo.wakeReason)
+                wakeReasonAcquired = newInfo.wakeReason
             }
             logger.logViewAddition(newInfo.windowTitle)
             inflateAndUpdateView(newInfo)
@@ -155,6 +170,7 @@ abstract class TemporaryViewDisplayController<T : TemporaryViewInfo, U : Tempora
             it.copyFrom(windowLayoutParams)
             it.title = newInfo.windowTitle
         }
+        newView.keepScreenOn = true
         windowManager.addView(newView, paramsWithTitle)
         animateViewIn(newView)
     }
@@ -183,7 +199,10 @@ abstract class TemporaryViewDisplayController<T : TemporaryViewInfo, U : Tempora
         val currentDisplayInfo = displayInfo ?: return
 
         val currentView = currentDisplayInfo.view
-        animateViewOut(currentView) { windowManager.removeView(currentView) }
+        animateViewOut(currentView) {
+            windowManager.removeView(currentView)
+            wakeLock?.release(wakeReasonAcquired)
+        }
 
         logger.logViewRemoval(removalReason)
         configurationController.removeCallback(displayScaleListener)
diff --git a/packages/SystemUI/src/com/android/systemui/temporarydisplay/chipbar/ChipbarCoordinator.kt b/packages/SystemUI/src/com/android/systemui/temporarydisplay/chipbar/ChipbarCoordinator.kt
index 87b6e8d3af34..44e5ce865241 100644
--- a/packages/SystemUI/src/com/android/systemui/temporarydisplay/chipbar/ChipbarCoordinator.kt
+++ b/packages/SystemUI/src/com/android/systemui/temporarydisplay/chipbar/ChipbarCoordinator.kt
@@ -44,6 +44,7 @@ import com.android.systemui.statusbar.policy.ConfigurationController
 import com.android.systemui.temporarydisplay.TemporaryViewDisplayController
 import com.android.systemui.util.concurrency.DelayableExecutor
 import com.android.systemui.util.view.ViewUtil
+import com.android.systemui.util.wakelock.WakeLock
 import javax.inject.Inject
 
 /**
@@ -75,6 +76,7 @@ open class ChipbarCoordinator @Inject constructor(
         private val falsingCollector: FalsingCollector,
         private val viewUtil: ViewUtil,
         private val vibratorHelper: VibratorHelper,
+        wakeLockBuilder: WakeLock.Builder,
 ) : TemporaryViewDisplayController<ChipbarInfo, ChipbarLogger>(
         context,
         logger,
@@ -84,6 +86,7 @@ open class ChipbarCoordinator @Inject constructor(
         configurationController,
         powerManager,
         R.layout.chipbar,
+        wakeLockBuilder,
 ) {
 
     private lateinit var parent: ChipbarRootView
@@ -92,8 +95,6 @@ open class ChipbarCoordinator @Inject constructor(
         gravity = Gravity.TOP.or(Gravity.CENTER_HORIZONTAL)
     }
 
-    override fun start() {}
-
     override fun updateView(
         newInfo: ChipbarInfo,
         currentView: ViewGroup
@@ -192,6 +193,8 @@ open class ChipbarCoordinator @Inject constructor(
         )
     }
 
+    override fun start() {}
+
     override fun getTouchableRegion(view: View, outRect: Rect) {
         viewUtil.setRectToViewWindowLocation(view, outRect)
     }
diff --git a/packages/SystemUI/src/com/android/systemui/util/wakelock/WakeLock.java b/packages/SystemUI/src/com/android/systemui/util/wakelock/WakeLock.java
index 8d77c4a194a9..f320d071b54f 100644
--- a/packages/SystemUI/src/com/android/systemui/util/wakelock/WakeLock.java
+++ b/packages/SystemUI/src/com/android/systemui/util/wakelock/WakeLock.java
@@ -38,6 +38,11 @@ public interface WakeLock {
     long DEFAULT_MAX_TIMEOUT = 20000;
 
     /**
+     * Default wake-lock levels and flags.
+     */
+    int DEFAULT_LEVELS_AND_FLAGS = PowerManager.PARTIAL_WAKE_LOCK;
+
+    /**
      * @param why A tag that will be saved for sysui dumps.
      * @see android.os.PowerManager.WakeLock#acquire()
      **/
@@ -60,13 +65,21 @@ public interface WakeLock {
      * Creates a {@link WakeLock} that has a default release timeout.
      * @see android.os.PowerManager.WakeLock#acquire(long) */
     static WakeLock createPartial(Context context, String tag, long maxTimeout) {
-        return wrap(createPartialInner(context, tag), maxTimeout);
+        return wrap(createWakeLockInner(context, tag, DEFAULT_LEVELS_AND_FLAGS), maxTimeout);
+    }
+
+    /**
+     * Creates a {@link WakeLock} that has a default release timeout and flags.
+     */
+    static WakeLock createWakeLock(Context context, String tag, int flags, long maxTimeout) {
+        return wrap(createWakeLockInner(context, tag, flags), maxTimeout);
     }
 
     @VisibleForTesting
-    static PowerManager.WakeLock createPartialInner(Context context, String tag) {
+    static PowerManager.WakeLock createWakeLockInner(
+            Context context, String tag, int levelsAndFlags) {
         return context.getSystemService(PowerManager.class)
-                    .newWakeLock(PowerManager.PARTIAL_WAKE_LOCK, tag);
+                    .newWakeLock(levelsAndFlags, tag);
     }
 
     static Runnable wrapImpl(WakeLock w, Runnable r) {
@@ -131,6 +144,7 @@ public interface WakeLock {
     class Builder {
         private final Context mContext;
         private String mTag;
+        private int mLevelsAndFlags = DEFAULT_LEVELS_AND_FLAGS;
         private long mMaxTimeout = DEFAULT_MAX_TIMEOUT;
 
         @Inject
@@ -143,13 +157,18 @@ public interface WakeLock {
             return this;
         }
 
+        public Builder setLevelsAndFlags(int levelsAndFlags) {
+            this.mLevelsAndFlags = levelsAndFlags;
+            return this;
+        }
+
         public Builder setMaxTimeout(long maxTimeout) {
             this.mMaxTimeout = maxTimeout;
             return this;
         }
 
         public WakeLock build() {
-            return WakeLock.createPartial(mContext, mTag, mMaxTimeout);
+            return WakeLock.createWakeLock(mContext, mTag, mLevelsAndFlags, mMaxTimeout);
         }
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/taptotransfer/receiver/MediaTttChipControllerReceiverTest.kt b/packages/SystemUI/tests/src/com/android/systemui/media/taptotransfer/receiver/MediaTttChipControllerReceiverTest.kt
index 8c3ae3d01f1d..68a5f47c5e0b 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/media/taptotransfer/receiver/MediaTttChipControllerReceiverTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/taptotransfer/receiver/MediaTttChipControllerReceiverTest.kt
@@ -43,6 +43,7 @@ import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.eq
 import com.android.systemui.util.time.FakeSystemClock
 import com.android.systemui.util.view.ViewUtil
+import com.android.systemui.util.wakelock.WakeLockFake
 import com.google.common.truth.Truth.assertThat
 import org.junit.Before
 import org.junit.Test
@@ -85,6 +86,10 @@ class MediaTttChipControllerReceiverTest : SysuiTestCase() {
     private lateinit var fakeAppIconDrawable: Drawable
     private lateinit var uiEventLoggerFake: UiEventLoggerFake
     private lateinit var receiverUiEventLogger: MediaTttReceiverUiEventLogger
+    private lateinit var fakeClock: FakeSystemClock
+    private lateinit var fakeExecutor: FakeExecutor
+    private lateinit var fakeWakeLockBuilder: WakeLockFake.Builder
+    private lateinit var fakeWakeLock: WakeLockFake
 
     @Before
     fun setUp() {
@@ -99,15 +104,22 @@ class MediaTttChipControllerReceiverTest : SysuiTestCase() {
         )).thenReturn(applicationInfo)
         context.setMockPackageManager(packageManager)
 
+        fakeClock = FakeSystemClock()
+        fakeExecutor = FakeExecutor(fakeClock)
+
         uiEventLoggerFake = UiEventLoggerFake()
         receiverUiEventLogger = MediaTttReceiverUiEventLogger(uiEventLoggerFake)
 
+        fakeWakeLock = WakeLockFake()
+        fakeWakeLockBuilder = WakeLockFake.Builder(context)
+        fakeWakeLockBuilder.setWakeLock(fakeWakeLock)
+
         controllerReceiver = MediaTttChipControllerReceiver(
             commandQueue,
             context,
             logger,
             windowManager,
-            FakeExecutor(FakeSystemClock()),
+            fakeExecutor,
             accessibilityManager,
             configurationController,
             powerManager,
@@ -115,6 +127,7 @@ class MediaTttChipControllerReceiverTest : SysuiTestCase() {
             mediaTttFlags,
             receiverUiEventLogger,
             viewUtil,
+            fakeWakeLockBuilder,
         )
         controllerReceiver.start()
 
@@ -141,6 +154,7 @@ class MediaTttChipControllerReceiverTest : SysuiTestCase() {
             mediaTttFlags,
             receiverUiEventLogger,
             viewUtil,
+            fakeWakeLockBuilder,
         )
         controllerReceiver.start()
 
@@ -200,6 +214,39 @@ class MediaTttChipControllerReceiverTest : SysuiTestCase() {
     }
 
     @Test
+    fun commandQueueCallback_closeThenFar_wakeLockAcquiredThenReleased() {
+        commandQueueCallback.updateMediaTapToTransferReceiverDisplay(
+                StatusBarManager.MEDIA_TRANSFER_RECEIVER_STATE_CLOSE_TO_SENDER,
+                routeInfo,
+                null,
+                null
+        )
+
+        assertThat(fakeWakeLock.isHeld).isTrue()
+
+        commandQueueCallback.updateMediaTapToTransferReceiverDisplay(
+                StatusBarManager.MEDIA_TRANSFER_RECEIVER_STATE_FAR_FROM_SENDER,
+                routeInfo,
+                null,
+                null
+        )
+
+        assertThat(fakeWakeLock.isHeld).isFalse()
+    }
+
+    @Test
+    fun commandQueueCallback_closeThenFar_wakeLockNeverAcquired() {
+        commandQueueCallback.updateMediaTapToTransferReceiverDisplay(
+                StatusBarManager.MEDIA_TRANSFER_RECEIVER_STATE_FAR_FROM_SENDER,
+                routeInfo,
+                null,
+                null
+        )
+
+        assertThat(fakeWakeLock.isHeld).isFalse()
+    }
+
+    @Test
     fun receivesNewStateFromCommandQueue_isLogged() {
         commandQueueCallback.updateMediaTapToTransferReceiverDisplay(
             StatusBarManager.MEDIA_TRANSFER_RECEIVER_STATE_CLOSE_TO_SENDER,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/taptotransfer/sender/MediaTttSenderCoordinatorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/media/taptotransfer/sender/MediaTttSenderCoordinatorTest.kt
index ad19bc2a80e0..4437394da943 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/media/taptotransfer/sender/MediaTttSenderCoordinatorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/taptotransfer/sender/MediaTttSenderCoordinatorTest.kt
@@ -52,6 +52,7 @@ import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.eq
 import com.android.systemui.util.time.FakeSystemClock
 import com.android.systemui.util.view.ViewUtil
+import com.android.systemui.util.wakelock.WakeLockFake
 import com.google.common.truth.Truth.assertThat
 import org.junit.Before
 import org.junit.Test
@@ -89,6 +90,8 @@ class MediaTttSenderCoordinatorTest : SysuiTestCase() {
     @Mock private lateinit var viewUtil: ViewUtil
     @Mock private lateinit var windowManager: WindowManager
     @Mock private lateinit var vibratorHelper: VibratorHelper
+    private lateinit var fakeWakeLockBuilder: WakeLockFake.Builder
+    private lateinit var fakeWakeLock: WakeLockFake
     private lateinit var chipbarCoordinator: ChipbarCoordinator
     private lateinit var commandQueueCallback: CommandQueue.Callbacks
     private lateinit var fakeAppIconDrawable: Drawable
@@ -118,6 +121,10 @@ class MediaTttSenderCoordinatorTest : SysuiTestCase() {
         fakeClock = FakeSystemClock()
         fakeExecutor = FakeExecutor(fakeClock)
 
+        fakeWakeLock = WakeLockFake()
+        fakeWakeLockBuilder = WakeLockFake.Builder(context)
+        fakeWakeLockBuilder.setWakeLock(fakeWakeLock)
+
         uiEventLoggerFake = UiEventLoggerFake()
         uiEventLogger = MediaTttSenderUiEventLogger(uiEventLoggerFake)
 
@@ -134,6 +141,7 @@ class MediaTttSenderCoordinatorTest : SysuiTestCase() {
                 falsingCollector,
                 viewUtil,
                 vibratorHelper,
+                fakeWakeLockBuilder,
             )
         chipbarCoordinator.start()
 
@@ -472,6 +480,36 @@ class MediaTttSenderCoordinatorTest : SysuiTestCase() {
     }
 
     @Test
+    fun commandQueueCallback_almostCloseThenFarFromReceiver_wakeLockAcquiredThenReleased() {
+        commandQueueCallback.updateMediaTapToTransferSenderDisplay(
+            StatusBarManager.MEDIA_TRANSFER_SENDER_STATE_ALMOST_CLOSE_TO_START_CAST,
+            routeInfo,
+            null
+        )
+
+        assertThat(fakeWakeLock.isHeld).isTrue()
+
+        commandQueueCallback.updateMediaTapToTransferSenderDisplay(
+            StatusBarManager.MEDIA_TRANSFER_SENDER_STATE_FAR_FROM_RECEIVER,
+            routeInfo,
+            null
+        )
+
+        assertThat(fakeWakeLock.isHeld).isFalse()
+    }
+
+    @Test
+    fun commandQueueCallback_FarFromReceiver_wakeLockNeverReleased() {
+        commandQueueCallback.updateMediaTapToTransferSenderDisplay(
+            StatusBarManager.MEDIA_TRANSFER_SENDER_STATE_FAR_FROM_RECEIVER,
+            routeInfo,
+            null
+        )
+
+        assertThat(fakeWakeLock.isHeld).isFalse()
+    }
+
+    @Test
     fun commandQueueCallback_invalidStateParam_noChipShown() {
         commandQueueCallback.updateMediaTapToTransferSenderDisplay(100, routeInfo, null)
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/temporarydisplay/TemporaryViewDisplayControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/temporarydisplay/TemporaryViewDisplayControllerTest.kt
index 91b5c35d9661..9dea48e3b47c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/temporarydisplay/TemporaryViewDisplayControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/temporarydisplay/TemporaryViewDisplayControllerTest.kt
@@ -35,6 +35,8 @@ import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.argumentCaptor
 import com.android.systemui.util.mockito.capture
 import com.android.systemui.util.time.FakeSystemClock
+import com.android.systemui.util.wakelock.WakeLock
+import com.android.systemui.util.wakelock.WakeLockFake
 import com.google.common.truth.Truth.assertThat
 import org.junit.Before
 import org.junit.Test
@@ -53,6 +55,9 @@ class TemporaryViewDisplayControllerTest : SysuiTestCase() {
     private lateinit var fakeClock: FakeSystemClock
     private lateinit var fakeExecutor: FakeExecutor
 
+    private lateinit var fakeWakeLockBuilder: WakeLockFake.Builder
+    private lateinit var fakeWakeLock: WakeLockFake
+
     @Mock
     private lateinit var logger: TemporaryViewLogger
     @Mock
@@ -74,6 +79,10 @@ class TemporaryViewDisplayControllerTest : SysuiTestCase() {
         fakeClock = FakeSystemClock()
         fakeExecutor = FakeExecutor(fakeClock)
 
+        fakeWakeLock = WakeLockFake()
+        fakeWakeLockBuilder = WakeLockFake.Builder(context)
+        fakeWakeLockBuilder.setWakeLock(fakeWakeLock)
+
         underTest = TestController(
                 context,
                 logger,
@@ -82,7 +91,9 @@ class TemporaryViewDisplayControllerTest : SysuiTestCase() {
                 accessibilityManager,
                 configurationController,
                 powerManager,
+                fakeWakeLockBuilder,
         )
+        underTest.start()
     }
 
     @Test
@@ -112,25 +123,33 @@ class TemporaryViewDisplayControllerTest : SysuiTestCase() {
     }
 
     @Test
-    fun displayView_screenOff_screenWakes() {
-        whenever(powerManager.isScreenOn).thenReturn(false)
-
+    fun displayView_screenOff_wakeLockAcquired() {
         underTest.displayView(getState())
 
-        verify(powerManager).wakeUp(any(), any(), any())
+        assertThat(fakeWakeLock.isHeld).isTrue()
     }
 
     @Test
-    fun displayView_screenAlreadyOn_screenNotWoken() {
+    fun displayView_screenAlreadyOn_wakeLockNotAcquired() {
         whenever(powerManager.isScreenOn).thenReturn(true)
 
         underTest.displayView(getState())
 
-        verify(powerManager, never()).wakeUp(any(), any(), any())
+        assertThat(fakeWakeLock.isHeld).isFalse()
+    }
+
+    @Test
+    fun displayView_screenOff_wakeLockCanBeReleasedAfterTimeOut() {
+        underTest.displayView(getState())
+        assertThat(fakeWakeLock.isHeld).isTrue()
+
+        fakeClock.advanceTime(TIMEOUT_MS + 1)
+
+        assertThat(fakeWakeLock.isHeld).isFalse()
     }
 
     @Test
-    fun displayView_twiceWithSameWindowTitle_viewNotAddedTwice() {
+    fun displayView_twice_viewNotAddedTwice() {
         underTest.displayView(getState())
         reset(windowManager)
 
@@ -269,6 +288,7 @@ class TemporaryViewDisplayControllerTest : SysuiTestCase() {
         accessibilityManager: AccessibilityManager,
         configurationController: ConfigurationController,
         powerManager: PowerManager,
+        wakeLockBuilder: WakeLock.Builder,
     ) : TemporaryViewDisplayController<ViewInfo, TemporaryViewLogger>(
         context,
         logger,
@@ -278,13 +298,12 @@ class TemporaryViewDisplayControllerTest : SysuiTestCase() {
         configurationController,
         powerManager,
         R.layout.chipbar,
+        wakeLockBuilder,
     ) {
         var mostRecentViewInfo: ViewInfo? = null
 
         override val windowLayoutParams = commonWindowLayoutParams
 
-        override fun start() {}
-
         override fun updateView(newInfo: ViewInfo, currentView: ViewGroup) {
             mostRecentViewInfo = newInfo
         }
@@ -292,6 +311,8 @@ class TemporaryViewDisplayControllerTest : SysuiTestCase() {
         override fun getTouchableRegion(view: View, outRect: Rect) {
             outRect.setEmpty()
         }
+
+        override fun start() {}
     }
 
     inner class ViewInfo(
diff --git a/packages/SystemUI/tests/src/com/android/systemui/temporarydisplay/chipbar/ChipbarCoordinatorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/temporarydisplay/chipbar/ChipbarCoordinatorTest.kt
index f64397325867..8e37aa292240 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/temporarydisplay/chipbar/ChipbarCoordinatorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/temporarydisplay/chipbar/ChipbarCoordinatorTest.kt
@@ -43,6 +43,7 @@ import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.eq
 import com.android.systemui.util.time.FakeSystemClock
 import com.android.systemui.util.view.ViewUtil
+import com.android.systemui.util.wakelock.WakeLockFake
 import com.google.common.truth.Truth.assertThat
 import org.junit.Before
 import org.junit.Test
@@ -69,6 +70,8 @@ class ChipbarCoordinatorTest : SysuiTestCase() {
     @Mock private lateinit var falsingCollector: FalsingCollector
     @Mock private lateinit var viewUtil: ViewUtil
     @Mock private lateinit var vibratorHelper: VibratorHelper
+    private lateinit var fakeWakeLockBuilder: WakeLockFake.Builder
+    private lateinit var fakeWakeLock: WakeLockFake
     private lateinit var fakeClock: FakeSystemClock
     private lateinit var fakeExecutor: FakeExecutor
     private lateinit var uiEventLoggerFake: UiEventLoggerFake
@@ -81,6 +84,10 @@ class ChipbarCoordinatorTest : SysuiTestCase() {
         fakeClock = FakeSystemClock()
         fakeExecutor = FakeExecutor(fakeClock)
 
+        fakeWakeLock = WakeLockFake()
+        fakeWakeLockBuilder = WakeLockFake.Builder(context)
+        fakeWakeLockBuilder.setWakeLock(fakeWakeLock)
+
         uiEventLoggerFake = UiEventLoggerFake()
 
         underTest =
@@ -96,6 +103,7 @@ class ChipbarCoordinatorTest : SysuiTestCase() {
                 falsingCollector,
                 viewUtil,
                 vibratorHelper,
+                fakeWakeLockBuilder,
             )
         underTest.start()
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/temporarydisplay/chipbar/FakeChipbarCoordinator.kt b/packages/SystemUI/tests/src/com/android/systemui/temporarydisplay/chipbar/FakeChipbarCoordinator.kt
index 574f70e7fddc..beedf9f337bc 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/temporarydisplay/chipbar/FakeChipbarCoordinator.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/temporarydisplay/chipbar/FakeChipbarCoordinator.kt
@@ -27,6 +27,7 @@ import com.android.systemui.statusbar.VibratorHelper
 import com.android.systemui.statusbar.policy.ConfigurationController
 import com.android.systemui.util.concurrency.DelayableExecutor
 import com.android.systemui.util.view.ViewUtil
+import com.android.systemui.util.wakelock.WakeLock
 
 /** A fake implementation of [ChipbarCoordinator] for testing. */
 class FakeChipbarCoordinator(
@@ -41,6 +42,7 @@ class FakeChipbarCoordinator(
     falsingCollector: FalsingCollector,
     viewUtil: ViewUtil,
     vibratorHelper: VibratorHelper,
+    wakeLockBuilder: WakeLock.Builder,
 ) :
     ChipbarCoordinator(
         context,
@@ -54,6 +56,7 @@ class FakeChipbarCoordinator(
         falsingCollector,
         viewUtil,
         vibratorHelper,
+        wakeLockBuilder,
     ) {
     override fun animateViewOut(view: ViewGroup, onAnimationEnd: Runnable) {
         // Just bypass the animation in tests
diff --git a/packages/SystemUI/tests/src/com/android/systemui/util/wakelock/WakeLockTest.java b/packages/SystemUI/tests/src/com/android/systemui/util/wakelock/WakeLockTest.java
index fe01f841aa16..6e109eafd748 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/util/wakelock/WakeLockTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/util/wakelock/WakeLockTest.java
@@ -42,7 +42,9 @@ public class WakeLockTest extends SysuiTestCase {
 
     @Before
     public void setUp() {
-        mInner = WakeLock.createPartialInner(mContext, WakeLockTest.class.getName());
+        mInner = WakeLock.createWakeLockInner(mContext,
+                WakeLockTest.class.getName(),
+                PowerManager.PARTIAL_WAKE_LOCK);
         mWakeLock = WakeLock.wrap(mInner, 20000);
     }