do: fix flagging logic for cancelBlockScreenOn

Flag: EXEMPT test flag fix
Test: atest DisplayPowerControllerTest
Bug: 350124222
Change-Id: I7d705745f36d51a04620cf77068f69dec2231f64

1 files changed, 6 insertions, 2 deletions

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 7b8b712c1aee..d070aaa93b01 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java
@@ -1700,10 +1700,12 @@ public final class DisplayPowerControllerTest {
         verify(mDisplayOffloadSession, never()).cancelBlockScreenOn();
     }
 
-    @RequiresFlagsEnabled(Flags.FLAG_OFFLOAD_SESSION_CANCEL_BLOCK_SCREEN_ON)
     @Test
     public void testOffloadBlocker_turnON_thenOFF_cancelBlockScreenOnNotCalledIfUnblocked() {
         // Set up.
+        when(mDisplayManagerFlagsMock.isDisplayOffloadEnabled()).thenReturn(true);
+        when(mDisplayManagerFlagsMock.isOffloadSessionCancelBlockScreenOnEnabled())
+                .thenReturn(true);
         int initState = Display.STATE_OFF;
         mHolder = createDisplayPowerController(DISPLAY_ID, UNIQUE_ID);
         mHolder.dpc.setDisplayOffloadSession(mDisplayOffloadSession);
@@ -1737,10 +1739,12 @@ public final class DisplayPowerControllerTest {
         verify(mDisplayOffloadSession, never()).cancelBlockScreenOn();
     }
 
-    @RequiresFlagsEnabled(Flags.FLAG_OFFLOAD_SESSION_CANCEL_BLOCK_SCREEN_ON)
     @Test
     public void testOffloadBlocker_turnON_thenOFF_cancelBlockScreenOn() {
         // Set up.
+        when(mDisplayManagerFlagsMock.isDisplayOffloadEnabled()).thenReturn(true);
+        when(mDisplayManagerFlagsMock.isOffloadSessionCancelBlockScreenOnEnabled())
+                .thenReturn(true);
         int initState = Display.STATE_OFF;
         mHolder = createDisplayPowerController(DISPLAY_ID, UNIQUE_ID);
         mHolder.dpc.setDisplayOffloadSession(mDisplayOffloadSession);