diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
index 59159bb..42c3630 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
@@ -4048,10 +4048,6 @@
         mRecentTasks.notifyTaskPersisterLocked(task, flush);
     }
 
-    boolean isKeyguardLocked(int displayId) {
-        return mKeyguardController.isKeyguardLocked(displayId);
-    }
-
     /**
      * Clears launch params for the given package.
      *
diff --git a/services/core/java/com/android/server/wm/KeyguardController.java b/services/core/java/com/android/server/wm/KeyguardController.java
index 8c1d8fa..1a319ad 100644
--- a/services/core/java/com/android/server/wm/KeyguardController.java
+++ b/services/core/java/com/android/server/wm/KeyguardController.java
@@ -423,7 +423,7 @@
 
         final TransitionController tc = mRootWindowContainer.mTransitionController;
 
-        final boolean occluded = isDisplayOccluded(displayId);
+        final boolean occluded = getDisplayState(displayId).mOccluded;
         final boolean performTransition = isKeyguardLocked(displayId);
         final boolean executeTransition = performTransition && !tc.isCollecting();
 
@@ -500,15 +500,6 @@
         }
     }
 
-    /**
-     * Returns {@code true} if the top activity on the display can occlude keyguard or the device
-     * is dreaming. Note that this method may return {@code true} even if the keyguard is disabled
-     * or not showing.
-     */
-    boolean isDisplayOccluded(int displayId) {
-        return getDisplayState(displayId).mOccluded;
-    }
-
     ActivityRecord getTopOccludingActivity(int displayId) {
         return getDisplayState(displayId).mTopOccludesActivity;
     }
@@ -601,6 +592,11 @@
         private boolean mAodShowing;
         private boolean mKeyguardGoingAway;
         private boolean mDismissalRequested;
+
+        /**
+         * True if the top activity on the display can occlude keyguard or the device is dreaming.
+         * Note that this can be true even if the keyguard is disabled or not showing.
+         */
         private boolean mOccluded;
         private boolean mShowingDream;
 
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index 6b4cc25..57f8268 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -2699,7 +2699,7 @@
             // transition exists, so this affects only when no lock screen is set. Otherwise
             // keyguard going away animation will be played.
             // See also AppTransitionController#getTransitCompatType for more details.
-            if ((!mTaskSupervisor.getKeyguardController().isDisplayOccluded(display.mDisplayId)
+            if ((!mTaskSupervisor.getKeyguardController().isKeyguardOccluded(display.mDisplayId)
                     && token.mTag.equals(KEYGUARD_SLEEP_TOKEN_TAG))
                     || token.mTag.equals(DISPLAY_OFF_SLEEP_TOKEN_TAG)) {
                 display.mSkipAppTransitionAnimation = true;
diff --git a/services/core/java/com/android/server/wm/Task.java b/services/core/java/com/android/server/wm/Task.java
index f9bbc68..387a876 100644
--- a/services/core/java/com/android/server/wm/Task.java
+++ b/services/core/java/com/android/server/wm/Task.java
@@ -3579,7 +3579,7 @@
                 && activity.info != info.taskInfo.topActivityInfo
                 ? activity.info : null;
         info.isKeyguardOccluded =
-            mAtmService.mKeyguardController.isDisplayOccluded(DEFAULT_DISPLAY);
+                mAtmService.mKeyguardController.isKeyguardOccluded(info.taskInfo.displayId);
 
         info.startingWindowTypeParameter = activity.mStartingData != null
                 ? activity.mStartingData.mTypeParams
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
index 302ad7f..31682bc 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
@@ -1541,7 +1541,8 @@
         // Make keyguard locked and set the top activity show-when-locked.
         KeyguardController keyguardController = activity.mTaskSupervisor.getKeyguardController();
         int displayId = activity.getDisplayId();
-        doReturn(true).when(keyguardController).isKeyguardLocked(eq(displayId));
+        keyguardController.setKeyguardShown(displayId, true /* keyguardShowing */,
+                false /* aodShowing */);
         final ActivityRecord topActivity = new ActivityBuilder(mAtm).setTask(task).build();
         topActivity.setVisibleRequested(true);
         topActivity.nowVisible = true;
@@ -1553,7 +1554,7 @@
 
         // Verify the stack-top activity is occluded keyguard.
         assertEquals(topActivity, task.topRunningActivity());
-        assertTrue(keyguardController.isDisplayOccluded(DEFAULT_DISPLAY));
+        assertTrue(keyguardController.isKeyguardOccluded(displayId));
 
         // Finish the top activity
         topActivity.setState(PAUSED, "true");
@@ -1562,7 +1563,7 @@
 
         // Verify new top activity does not occlude keyguard.
         assertEquals(activity, task.topRunningActivity());
-        assertFalse(keyguardController.isDisplayOccluded(DEFAULT_DISPLAY));
+        assertFalse(keyguardController.isKeyguardOccluded(displayId));
     }
 
     /**
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityTaskManagerServiceTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityTaskManagerServiceTests.java
index d169a58..5341588 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityTaskManagerServiceTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityTaskManagerServiceTests.java
@@ -24,7 +24,6 @@
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
 import static com.android.server.wm.ActivityInterceptorCallback.MAINLINE_FIRST_ORDERED_ID;
 import static com.android.server.wm.ActivityInterceptorCallback.SYSTEM_FIRST_ORDERED_ID;
@@ -41,11 +40,9 @@
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.anyBoolean;
-import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.atLeast;
 import static org.mockito.Mockito.doCallRealMethod;
-import static org.mockito.Mockito.when;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -60,7 +57,6 @@
 import android.content.res.Configuration;
 import android.graphics.Rect;
 import android.os.Binder;
-import android.os.IBinder;
 import android.os.LocaleList;
 import android.os.RemoteException;
 import android.platform.test.annotations.Presubmit;
@@ -76,7 +72,6 @@
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mockito;
-import org.mockito.MockitoSession;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -304,18 +299,11 @@
      */
     @Test
     public void testEnterPipModeWhenRecordParentChangesToNull() {
-        MockitoSession mockSession = mockitoSession()
-                .initMocks(this)
-                .mockStatic(ActivityRecord.class)
-                .startMocking();
-
-        ActivityRecord record = mock(ActivityRecord.class);
-        IBinder token = mock(IBinder.class);
+        final ActivityRecord record = new ActivityBuilder(mAtm).setCreateTask(true).build();
         PictureInPictureParams params = mock(PictureInPictureParams.class);
         record.pictureInPictureArgs = params;
 
         //mock operations in private method ensureValidPictureInPictureActivityParamsLocked()
-        when(ActivityRecord.forTokenLocked(token)).thenReturn(record);
         doReturn(true).when(record).supportsPictureInPicture();
         doReturn(false).when(params).hasSetAspectRatio();
 
@@ -323,15 +311,13 @@
         doReturn(true).when(record)
                 .checkEnterPictureInPictureState("enterPictureInPictureMode", false);
         doReturn(false).when(record).inPinnedWindowingMode();
-        doReturn(false).when(mAtm).isKeyguardLocked(anyInt());
+        doReturn(false).when(record).isKeyguardLocked();
 
         //to simulate NPE
         doReturn(null).when(record).getParent();
 
-        mAtm.mActivityClientController.enterPictureInPictureMode(token, params);
+        mAtm.mActivityClientController.enterPictureInPictureMode(record.token, params);
         //if record's null parent is not handled gracefully, test will fail with NPE
-
-        mockSession.finishMocking();
     }
 
     @Test