Disable screen off animation in situations where it currently looks bad.

- Use DozeParameters as source of truth for whether we will control unlocked screen off.
- Update controlScreenOff when rotation/expansion change, not just when keyguard visibility changes.

Fixes: 190478850
Fixes: 189803757
Test: screen on/off with shade expanded/collapsed and in portrait/landscape
Change-Id: I89e3dc626090580a832af296bb6665a52980640c
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java b/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java
index f3a6d63..0e43b31 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java
@@ -249,6 +249,16 @@
         public void onStateChanged(int newState) {
             mStatusBarState = newState;
         }
+
+        @Override
+        public void onExpandedChanged(boolean isExpanded) {
+            for (int i = 0; i < mCallbacks.size(); i++) {
+                KeyguardUpdateMonitorCallback cb = mCallbacks.get(i).get();
+                if (cb != null) {
+                    cb.onShadeExpandedChanged(isExpanded);
+                }
+            }
+        }
     };
 
     HashMap<Integer, SimData> mSimDatas = new HashMap<>();
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitorCallback.java b/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitorCallback.java
index e561a5a..9849a7e 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitorCallback.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitorCallback.java
@@ -333,4 +333,8 @@
      */
     public void onRequireUnlockForNfc() { }
 
+    /**
+     * Called when the notification shade is expanded or collapsed.
+     */
+    public void onShadeExpandedChanged(boolean expanded) { }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeUi.java b/packages/SystemUI/src/com/android/systemui/doze/DozeUi.java
index ff3cb21..fbe06b0 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeUi.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeUi.java
@@ -21,6 +21,7 @@
 
 import android.app.AlarmManager;
 import android.content.Context;
+import android.content.res.Configuration;
 import android.os.Handler;
 import android.os.SystemClock;
 import android.provider.Settings;
@@ -34,6 +35,7 @@
 import com.android.systemui.doze.dagger.DozeScope;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
 import com.android.systemui.statusbar.phone.DozeParameters;
+import com.android.systemui.statusbar.policy.ConfigurationController;
 import com.android.systemui.tuner.TunerService;
 import com.android.systemui.util.AlarmTimeout;
 import com.android.systemui.util.wakelock.WakeLock;
@@ -48,7 +50,8 @@
  * The policy controlling doze.
  */
 @DozeScope
-public class DozeUi implements DozeMachine.Part, TunerService.Tunable {
+public class DozeUi implements DozeMachine.Part, TunerService.Tunable,
+        ConfigurationController.ConfigurationListener {
     // if enabled, calls dozeTimeTick() whenever the time changes:
     private static final boolean BURN_IN_TESTING_ENABLED = false;
     private static final long TIME_TICK_DEADLINE_MILLIS = 90 * 1000; // 1.5min
@@ -63,6 +66,7 @@
     private final DozeLog mDozeLog;
     private final Lazy<StatusBarStateController> mStatusBarStateController;
     private final TunerService mTunerService;
+    private final ConfigurationController mConfigurationController;
 
     private boolean mKeyguardShowing;
     private final KeyguardUpdateMonitorCallback mKeyguardVisibilityCallback =
@@ -84,6 +88,11 @@
                         mHandler.post(mWakeLock.wrap(() -> {}));
                     }
                 }
+
+                @Override
+                public void onShadeExpandedChanged(boolean expanded) {
+                    updateAnimateScreenOff();
+                }
             };
 
     private long mLastTimeTickElapsed = 0;
@@ -93,7 +102,8 @@
             WakeLock wakeLock, DozeHost host, @Main Handler handler,
             DozeParameters params, KeyguardUpdateMonitor keyguardUpdateMonitor,
             DozeLog dozeLog, TunerService tunerService,
-            Lazy<StatusBarStateController> statusBarStateController) {
+            Lazy<StatusBarStateController> statusBarStateController,
+            ConfigurationController configurationController) {
         mContext = context;
         mWakeLock = wakeLock;
         mHost = host;
@@ -107,11 +117,15 @@
         mStatusBarStateController = statusBarStateController;
 
         mTunerService.addTunable(this, Settings.Secure.DOZE_ALWAYS_ON);
+
+        mConfigurationController = configurationController;
+        mConfigurationController.addCallback(this);
     }
 
     @Override
     public void destroy() {
         mTunerService.removeTunable(this);
+        mConfigurationController.removeCallback(this);
     }
 
     @Override
@@ -274,4 +288,9 @@
             updateAnimateScreenOff();
         }
     }
+
+    @Override
+    public void onConfigChanged(Configuration newConfig) {
+        updateAnimateScreenOff();
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationWakeUpCoordinator.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationWakeUpCoordinator.kt
index 760bee2..b0a7767 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationWakeUpCoordinator.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationWakeUpCoordinator.kt
@@ -264,7 +264,7 @@
     }
 
     override fun onStateChanged(newState: Int) {
-        if (unlockedScreenOffAnimationController.shouldPlayScreenOffAnimation()) {
+        if (dozeParameters.shouldControlUnlockedScreenOff()) {
             if (unlockedScreenOffAnimationController.isScreenOffAnimationPlaying() &&
                     state == StatusBarState.KEYGUARD &&
                     newState == StatusBarState.SHADE) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeParameters.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeParameters.java
index 01d489f..c4d1abc 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeParameters.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeParameters.java
@@ -63,6 +63,7 @@
     private final Resources mResources;
     private final BatteryController mBatteryController;
     private final FeatureFlags mFeatureFlags;
+    private final UnlockedScreenOffAnimationController mUnlockedScreenOffAnimationController;
 
     private final Set<Callback> mCallbacks = new HashSet<>();
 
@@ -78,7 +79,8 @@
             BatteryController batteryController,
             TunerService tunerService,
             DumpManager dumpManager,
-            FeatureFlags featureFlags) {
+            FeatureFlags featureFlags,
+            UnlockedScreenOffAnimationController unlockedScreenOffAnimationController) {
         mResources = resources;
         mAmbientDisplayConfiguration = ambientDisplayConfiguration;
         mAlwaysOnPolicy = alwaysOnDisplayPolicy;
@@ -89,6 +91,7 @@
         mPowerManager = powerManager;
         mPowerManager.setDozeAfterScreenOff(!mControlScreenOffAnimation);
         mFeatureFlags = featureFlags;
+        mUnlockedScreenOffAnimationController = unlockedScreenOffAnimationController;
 
         tunerService.addTunable(
                 this,
@@ -220,7 +223,8 @@
      * then abruptly showing AOD.
      */
     public boolean shouldControlUnlockedScreenOff() {
-        return getAlwaysOn() && mFeatureFlags.useNewLockscreenAnimations();
+        return getAlwaysOn() && mFeatureFlags.useNewLockscreenAnimations()
+                && mUnlockedScreenOffAnimationController.shouldPlayUnlockedScreenOffAnimation();
     }
 
     private boolean getBoolean(String propName, int resId) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationPanelViewController.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationPanelViewController.java
index 528827f..57d0bf1 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationPanelViewController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationPanelViewController.java
@@ -4222,7 +4222,7 @@
             int oldState = mBarState;
             boolean keyguardShowing = statusBarState == KEYGUARD;
 
-            if (mUnlockedScreenOffAnimationController.shouldPlayScreenOffAnimation()
+            if (mDozeParameters.shouldControlUnlockedScreenOff()
                     && oldState == StatusBarState.SHADE
                     && statusBarState == KEYGUARD) {
                 // This means we're doing the screen off animation - position the keyguard status
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationShadeWindowControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationShadeWindowControllerImpl.java
index 52f9aca..c958796 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationShadeWindowControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationShadeWindowControllerImpl.java
@@ -26,11 +26,9 @@
 import android.app.IActivityManager;
 import android.content.Context;
 import android.content.pm.ActivityInfo;
-import android.content.res.Resources;
 import android.graphics.PixelFormat;
 import android.os.Binder;
 import android.os.RemoteException;
-import android.os.SystemProperties;
 import android.os.Trace;
 import android.util.Log;
 import android.view.Display;
@@ -53,6 +51,7 @@
 import com.android.systemui.statusbar.SysuiStatusBarStateController;
 import com.android.systemui.statusbar.policy.ConfigurationController;
 import com.android.systemui.statusbar.policy.ConfigurationController.ConfigurationListener;
+import com.android.systemui.statusbar.policy.KeyguardStateController;
 
 import com.google.android.collect.Lists;
 
@@ -108,12 +107,14 @@
             StatusBarStateController statusBarStateController,
             ConfigurationController configurationController,
             KeyguardViewMediator keyguardViewMediator,
-            KeyguardBypassController keyguardBypassController, SysuiColorExtractor colorExtractor,
-            DumpManager dumpManager) {
+            KeyguardBypassController keyguardBypassController,
+            SysuiColorExtractor colorExtractor,
+            DumpManager dumpManager,
+            KeyguardStateController keyguardStateController) {
         mContext = context;
         mWindowManager = windowManager;
         mActivityManager = activityManager;
-        mKeyguardScreenRotation = shouldEnableKeyguardScreenRotation();
+        mKeyguardScreenRotation = keyguardStateController.isKeyguardScreenRotationAllowed();
         mDozeParameters = dozeParameters;
         mScreenBrightnessDoze = mDozeParameters.getScreenBrightnessDoze();
         mLpChanged = new LayoutParams();
@@ -173,12 +174,6 @@
         }
     }
 
-    private boolean shouldEnableKeyguardScreenRotation() {
-        Resources res = mContext.getResources();
-        return SystemProperties.getBoolean("lockscreen.rot_override", false)
-                || res.getBoolean(R.bool.config_enableLockScreenRotation);
-    }
-
     /**
      * Adds the notification shade view to the window manager.
      */
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
index 5ee5e48..0271d5d 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
@@ -4060,7 +4060,7 @@
 
             // The screen off animation uses our LightRevealScrim - we need to be expanded for it to
             // be visible.
-            if (mUnlockedScreenOffAnimationController.shouldPlayScreenOffAnimation()) {
+            if (mDozeParameters.shouldControlUnlockedScreenOff()) {
                 makeExpandedVisible(true);
             }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/UnlockedScreenOffAnimationController.kt b/packages/SystemUI/src/com/android/systemui/statusbar/phone/UnlockedScreenOffAnimationController.kt
index e135cc5..f8e8ee0 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/UnlockedScreenOffAnimationController.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/UnlockedScreenOffAnimationController.kt
@@ -3,6 +3,8 @@
 import android.animation.Animator
 import android.animation.AnimatorListenerAdapter
 import android.animation.ValueAnimator
+import android.content.Context
+import android.content.res.Configuration
 import android.os.Handler
 import android.view.View
 import com.android.systemui.animation.Interpolators
@@ -16,6 +18,7 @@
 import com.android.systemui.statusbar.notification.PropertyAnimator
 import com.android.systemui.statusbar.notification.stack.AnimationProperties
 import com.android.systemui.statusbar.notification.stack.StackStateAnimator
+import com.android.systemui.statusbar.policy.KeyguardStateController
 import javax.inject.Inject
 
 /**
@@ -38,10 +41,11 @@
  */
 @SysUISingleton
 class UnlockedScreenOffAnimationController @Inject constructor(
+    private val context: Context,
     private val wakefulnessLifecycle: WakefulnessLifecycle,
     private val statusBarStateControllerImpl: StatusBarStateControllerImpl,
     private val keyguardViewMediatorLazy: dagger.Lazy<KeyguardViewMediator>,
-    private val dozeParameters: DozeParameters
+    private val keyguardStateController: KeyguardStateController
 ) : WakefulnessLifecycle.Observer {
     private val handler = Handler()
 
@@ -137,7 +141,7 @@
     }
 
     override fun onStartedGoingToSleep() {
-        if (shouldPlayScreenOffAnimation()) {
+        if (shouldPlayUnlockedScreenOffAnimation()) {
             lightRevealAnimationPlaying = true
             lightRevealAnimator.start()
 
@@ -151,13 +155,31 @@
     }
 
     /**
-     * Whether we should play the screen off animation when the phone starts going to sleep. We can
-     * do that if dozeParameters says we can control the unlocked screen off animation and we are in
-     * the SHADE state. If we're in KEYGUARD or SHADE_LOCKED, the regular
+     * Whether we want to play the screen off animation when the phone starts going to sleep, based
+     * on the current state of the device.
      */
-    fun shouldPlayScreenOffAnimation(): Boolean {
-        return dozeParameters.shouldControlUnlockedScreenOff() &&
-                statusBarStateControllerImpl.state == StatusBarState.SHADE
+    fun shouldPlayUnlockedScreenOffAnimation(): Boolean {
+        // We only play the unlocked screen off animation if we are... unlocked.
+        if (statusBarStateControllerImpl.state != StatusBarState.SHADE) {
+            return false
+        }
+
+        // We currently draw both the light reveal scrim, and the AOD UI, in the shade. If it's
+        // already expanded and showing notifications/QS, the animation looks really messy. For now,
+        // disable it if the notification panel is expanded.
+        if (statusBar.notificationPanelViewController.isFullyExpanded) {
+            return false
+        }
+
+        // If we're not allowed to rotate the keyguard, then only do the screen off animation if
+        // we're in portrait. Otherwise, AOD will animate in sideways, which looks weird.
+        if (!keyguardStateController.isKeyguardScreenRotationAllowed &&
+                context.resources.configuration.orientation != Configuration.ORIENTATION_PORTRAIT) {
+            return false
+        }
+
+        // Otherwise, good to go.
+        return true
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/KeyguardStateController.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/KeyguardStateController.java
index af7bf95..fcfc967 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/KeyguardStateController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/KeyguardStateController.java
@@ -57,6 +57,11 @@
     boolean canPerformSmartSpaceTransition();
 
     /**
+     * Whether the keyguard is allowed to rotate, or needs to be locked to the default orientation.
+     */
+    boolean isKeyguardScreenRotationAllowed();
+
+    /**
      * If the device has PIN/pattern/password or a lock screen at all.
      */
     boolean isMethodSecure();
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/KeyguardStateControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/KeyguardStateControllerImpl.java
index 0945a3f..64750bd 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/KeyguardStateControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/KeyguardStateControllerImpl.java
@@ -23,6 +23,7 @@
 import android.content.IntentFilter;
 import android.hardware.biometrics.BiometricSourceType;
 import android.os.Build;
+import android.os.SystemProperties;
 import android.os.Trace;
 
 import androidx.annotation.VisibleForTesting;
@@ -31,6 +32,7 @@
 import com.android.keyguard.KeyguardUpdateMonitor;
 import com.android.keyguard.KeyguardUpdateMonitorCallback;
 import com.android.systemui.Dumpable;
+import com.android.systemui.R;
 import com.android.systemui.dagger.SysUISingleton;
 import com.android.systemui.shared.system.smartspace.SmartspaceTransitionController;
 
@@ -50,6 +52,7 @@
     private static final String AUTH_BROADCAST_KEY = "debug_trigger_auth";
 
     private final ArrayList<Callback> mCallbacks = new ArrayList<>();
+    private final Context mContext;
     private final KeyguardUpdateMonitor mKeyguardUpdateMonitor;
     private final LockPatternUtils mLockPatternUtils;
     private final KeyguardUpdateMonitorCallback mKeyguardUpdateMonitorCallback =
@@ -100,6 +103,7 @@
     public KeyguardStateControllerImpl(Context context,
             KeyguardUpdateMonitor keyguardUpdateMonitor, LockPatternUtils lockPatternUtils,
             SmartspaceTransitionController smartspaceTransitionController) {
+        mContext = context;
         mKeyguardUpdateMonitor = keyguardUpdateMonitor;
         mLockPatternUtils = lockPatternUtils;
         mKeyguardUpdateMonitor.registerCallback(mKeyguardUpdateMonitorCallback);
@@ -243,6 +247,12 @@
     }
 
     @Override
+    public boolean isKeyguardScreenRotationAllowed() {
+        return SystemProperties.getBoolean("lockscreen.rot_override", false)
+                || mContext.getResources().getBoolean(R.bool.config_enableLockScreenRotation);
+    }
+
+    @Override
     public boolean isFaceAuthEnabled() {
         return mFaceAuthEnabled;
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeUiTest.java b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeUiTest.java
index afe5c0b..1d34aac 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeUiTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeUiTest.java
@@ -43,6 +43,7 @@
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
 import com.android.systemui.statusbar.phone.DozeParameters;
+import com.android.systemui.statusbar.policy.ConfigurationController;
 import com.android.systemui.tuner.TunerService;
 import com.android.systemui.util.wakelock.WakeLockFake;
 
@@ -77,6 +78,8 @@
     private DozeUi mDozeUi;
     @Mock
     private StatusBarStateController mStatusBarStateController;
+    @Mock
+    private ConfigurationController mConfigurationController;
 
     @Before
     public void setUp() throws Exception {
@@ -89,7 +92,7 @@
 
         mDozeUi = new DozeUi(mContext, mAlarmManager, mWakeLock, mHost, mHandler,
                 mDozeParameters, mKeyguardUpdateMonitor, mDozeLog, mTunerService,
-                () -> mStatusBarStateController);
+                () -> mStatusBarStateController, mConfigurationController);
         mDozeUi.setDozeMachine(mMachine);
     }
 
@@ -146,7 +149,7 @@
         when(mDozeParameters.getDisplayNeedsBlanking()).thenReturn(true);
         mDozeUi = new DozeUi(mContext, mAlarmManager, mWakeLock, mHost, mHandler,
                 mDozeParameters, mKeyguardUpdateMonitor, mDozeLog, mTunerService,
-                () -> mStatusBarStateController);
+                () -> mStatusBarStateController, mConfigurationController);
         mDozeUi.setDozeMachine(mMachine);
 
         // Never animate if display doesn't support it.
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/DozeParametersTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/DozeParametersTest.java
index 9ac600a..5bf1bb3 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/DozeParametersTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/DozeParametersTest.java
@@ -63,6 +63,7 @@
     @Mock private BatteryController mBatteryController;
     @Mock private FeatureFlags mFeatureFlags;
     @Mock private DumpManager mDumpManager;
+    @Mock private UnlockedScreenOffAnimationController mUnlockedScreenOffAnimationController;
 
     @Before
     public void setup() {
@@ -75,7 +76,8 @@
             mBatteryController,
             mTunerService,
             mDumpManager,
-            mFeatureFlags
+            mFeatureFlags,
+            mUnlockedScreenOffAnimationController
         );
     }
     @Test
@@ -125,7 +127,8 @@
         when(mAmbientDisplayConfiguration.alwaysOnEnabled(anyInt())).thenReturn(true);
         mDozeParameters.onTuningChanged(Settings.Secure.DOZE_ALWAYS_ON, "1");
         when(mFeatureFlags.useNewLockscreenAnimations()).thenReturn(true);
-
+        when(mUnlockedScreenOffAnimationController.shouldPlayUnlockedScreenOffAnimation())
+                .thenReturn(true);
         assertTrue(mDozeParameters.shouldControlUnlockedScreenOff());
 
         // Trigger the setter for the current value.
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/NotificationShadeWindowControllerImplTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/NotificationShadeWindowControllerImplTest.java
index 3238430..9fe47ec 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/NotificationShadeWindowControllerImplTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/NotificationShadeWindowControllerImplTest.java
@@ -45,6 +45,7 @@
 import com.android.systemui.keyguard.KeyguardViewMediator;
 import com.android.systemui.statusbar.SysuiStatusBarStateController;
 import com.android.systemui.statusbar.policy.ConfigurationController;
+import com.android.systemui.statusbar.policy.KeyguardStateController;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -70,6 +71,7 @@
     @Mock private SysuiColorExtractor mColorExtractor;
     @Mock ColorExtractor.GradientColors mGradientColors;
     @Mock private DumpManager mDumpManager;
+    @Mock private KeyguardStateController mKeyguardStateController;
     @Captor private ArgumentCaptor<WindowManager.LayoutParams> mLayoutParameters;
 
     private NotificationShadeWindowControllerImpl mNotificationShadeWindowController;
@@ -83,7 +85,7 @@
         mNotificationShadeWindowController = new NotificationShadeWindowControllerImpl(mContext,
                 mWindowManager, mActivityManager, mDozeParameters, mStatusBarStateController,
                 mConfigurationController, mKeyguardViewMediator, mKeyguardBypassController,
-                mColorExtractor, mDumpManager);
+                mColorExtractor, mDumpManager, mKeyguardStateController);
         mNotificationShadeWindowController.setNotificationShadeView(mNotificationShadeWindowView);
 
         mNotificationShadeWindowController.attach();
diff --git a/packages/SystemUI/tests/src/com/android/systemui/utils/leaks/FakeKeyguardStateController.java b/packages/SystemUI/tests/src/com/android/systemui/utils/leaks/FakeKeyguardStateController.java
index 1aebf1c..e136d00 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/utils/leaks/FakeKeyguardStateController.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/utils/leaks/FakeKeyguardStateController.java
@@ -132,4 +132,9 @@
     public boolean canPerformSmartSpaceTransition() {
         return false;
     }
+
+    @Override
+    public boolean isKeyguardScreenRotationAllowed() {
+        return false;
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java b/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
index 6e2e4cb..cf2e3d0 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
@@ -98,6 +98,7 @@
 import com.android.systemui.statusbar.policy.BatteryController;
 import com.android.systemui.statusbar.policy.ConfigurationController;
 import com.android.systemui.statusbar.policy.HeadsUpManager;
+import com.android.systemui.statusbar.policy.KeyguardStateController;
 import com.android.systemui.statusbar.policy.ZenModeController;
 import com.android.wm.shell.R;
 import com.android.wm.shell.ShellTaskOrganizer;
@@ -227,6 +228,8 @@
     private TaskStackListenerImpl mTaskStackListener;
     @Mock
     private ShellTaskOrganizer mShellTaskOrganizer;
+    @Mock
+    private KeyguardStateController mKeyguardStateController;
 
     private TestableBubblePositioner mPositioner;
 
@@ -249,7 +252,7 @@
         mNotificationShadeWindowController = new NotificationShadeWindowControllerImpl(mContext,
                 mWindowManager, mActivityManager, mDozeParameters, mStatusBarStateController,
                 mConfigurationController, mKeyguardViewMediator, mKeyguardBypassController,
-                mColorExtractor, mDumpManager);
+                mColorExtractor, mDumpManager, mKeyguardStateController);
         mNotificationShadeWindowController.setNotificationShadeView(mNotificationShadeWindowView);
         mNotificationShadeWindowController.attach();
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/wmshell/NewNotifPipelineBubblesTest.java b/packages/SystemUI/tests/src/com/android/systemui/wmshell/NewNotifPipelineBubblesTest.java
index 9339f81..7d1979d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/wmshell/NewNotifPipelineBubblesTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/wmshell/NewNotifPipelineBubblesTest.java
@@ -82,6 +82,7 @@
 import com.android.systemui.statusbar.policy.BatteryController;
 import com.android.systemui.statusbar.policy.ConfigurationController;
 import com.android.systemui.statusbar.policy.HeadsUpManager;
+import com.android.systemui.statusbar.policy.KeyguardStateController;
 import com.android.systemui.statusbar.policy.ZenModeController;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.WindowManagerShellWrapper;
@@ -192,6 +193,8 @@
     private TaskStackListenerImpl mTaskStackListener;
     @Mock
     private ShellTaskOrganizer mShellTaskOrganizer;
+    @Mock
+    private KeyguardStateController mKeyguardStateController;
 
     private TestableBubblePositioner mPositioner;
 
@@ -213,7 +216,7 @@
         mNotificationShadeWindowController = new NotificationShadeWindowControllerImpl(mContext,
                 mWindowManager, mActivityManager, mDozeParameters, mStatusBarStateController,
                 mConfigurationController, mKeyguardViewMediator, mKeyguardBypassController,
-                mColorExtractor, mDumpManager);
+                mColorExtractor, mDumpManager, mKeyguardStateController);
         mNotificationShadeWindowController.setNotificationShadeView(mNotificationShadeWindowView);
         mNotificationShadeWindowController.attach();