dpc: fix DisplayPowerControllerTest for cancelBlockScreenOn

Flag: com.android.server.display.feature.flags.offload_session_cancel_block_screen_on
Test: atest DisplayOffloadSessionImplTest
Test: atest DisplayPowerControllerTest
Bug: 346420374
Bug: 347415946
Bug: 349320166

Change-Id: I30bde106483a726cc0ba0cf9c8a08bf18c7116a4

2 files changed, 5 insertions, 6 deletions

diff --git a/services/core/java/com/android/server/display/DisplayPowerController.java b/services/core/java/com/android/server/display/DisplayPowerController.java
index c298bbf5a551..08fe6212301f 100644
--- a/services/core/java/com/android/server/display/DisplayPowerController.java
+++ b/services/core/java/com/android/server/display/DisplayPowerController.java
@@ -2096,7 +2096,7 @@ final class DisplayPowerController implements AutomaticBrightnessController.Call
     private void onDisplayOffloadUnblockScreenOn(DisplayOffloadSession displayOffloadSession) {
         Message msg = mHandler.obtainMessage(MSG_OFFLOADING_SCREEN_ON_UNBLOCKED,
                 displayOffloadSession);
-        mHandler.sendMessage(msg);
+        mHandler.sendMessageAtTime(msg, mClock.uptimeMillis());
     }
 
     private void unblockScreenOnByDisplayOffload() {
@@ -2114,7 +2114,7 @@ final class DisplayPowerController implements AutomaticBrightnessController.Call
         if (mDisplayOffloadSession == null) {
             return;
         }
-        if (mPendingScreenOnUnblockerByDisplayOffload != null) {
+        if (mPendingScreenOnUnblockerByDisplayOffload == null) {
             // Already unblocked.
             return;
         }
@@ -2132,7 +2132,6 @@ final class DisplayPowerController implements AutomaticBrightnessController.Call
 
         // If the screen is turning on, give displayoffload a chance to do something before the
         // screen actually turns on.
-        // TODO(b/316941732): add tests for this displayoffload screen-on blocker.
         if (isOn && changed && !mScreenTurningOnWasBlockedByDisplayOffload) {
             blockScreenOnByDisplayOffload(mDisplayOffloadSession);
         } else if (!isOn && mScreenTurningOnWasBlockedByDisplayOffload) {
diff --git a/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java b/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java
index bb774eec9d4e..7b8b712c1aee 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java
@@ -1707,7 +1707,8 @@ public final class DisplayPowerControllerTest {
         int initState = Display.STATE_OFF;
         mHolder = createDisplayPowerController(DISPLAY_ID, UNIQUE_ID);
         mHolder.dpc.setDisplayOffloadSession(mDisplayOffloadSession);
-        when(mDisplayOffloadSession.blockScreenOn(any())).thenReturn(true);
+        ArgumentCaptor<Runnable> argumentCaptor = ArgumentCaptor.forClass(Runnable.class);
+        when(mDisplayOffloadSession.blockScreenOn(argumentCaptor.capture())).thenReturn(true);
 
         // Start with OFF.
         when(mHolder.displayPowerState.getScreenState()).thenReturn(initState);
@@ -1721,8 +1722,7 @@ public final class DisplayPowerControllerTest {
         mHolder.dpc.requestPowerState(dpr, /* waitForNegativeProximity= */ false);
         advanceTime(1); // Run updatePowerState
 
-        ArgumentCaptor<Runnable> argumentCaptor = ArgumentCaptor.forClass(Runnable.class);
-        verify(mDisplayOffloadSession).blockScreenOn(argumentCaptor.capture());
+        verify(mDisplayOffloadSession).blockScreenOn(any());
 
         // Unblocked
         argumentCaptor.getValue().run();