Merge "Make DozeParameters responsible for setting PowerManager#setDozeAfterScreenOff." into sc-v2-dev

8 files changed, 220 insertions, 210 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeHost.java b/packages/SystemUI/src/com/android/systemui/doze/DozeHost.java
index 9c25b3596be6..2beed4c6a7e7 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeHost.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeHost.java
@@ -40,7 +40,6 @@ public interface DozeHost {
     void extendPulse(int reason);
 
     void setAnimateWakeup(boolean animateWakeup);
-    void setAnimateScreenOff(boolean animateScreenOff);
 
     /**
      * Reports that a tap event happend on the Sensors Low Power Island.
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeUi.java b/packages/SystemUI/src/com/android/systemui/doze/DozeUi.java
index 374bed31eb7d..5677072db968 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeUi.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeUi.java
@@ -21,22 +21,17 @@ import static com.android.systemui.doze.DozeMachine.State.DOZE_AOD_PAUSED;
 
 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;
 import android.text.format.Formatter;
 import android.util.Log;
 
-import com.android.internal.annotations.VisibleForTesting;
 import com.android.keyguard.KeyguardUpdateMonitor;
 import com.android.keyguard.KeyguardUpdateMonitorCallback;
 import com.android.systemui.dagger.qualifiers.Main;
 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;
 
@@ -50,8 +45,7 @@ import dagger.Lazy;
  * The policy controlling doze.
  */
 @DozeScope
-public class DozeUi implements DozeMachine.Part, TunerService.Tunable,
-        ConfigurationController.ConfigurationListener {
+public class DozeUi implements DozeMachine.Part {
     // 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
@@ -65,19 +59,9 @@ public class DozeUi implements DozeMachine.Part, TunerService.Tunable,
     private final DozeParameters mDozeParameters;
     private final DozeLog mDozeLog;
     private final Lazy<StatusBarStateController> mStatusBarStateController;
-    private final TunerService mTunerService;
-    private final ConfigurationController mConfigurationController;
-
-    private boolean mKeyguardShowing;
     private final KeyguardUpdateMonitorCallback mKeyguardVisibilityCallback =
             new KeyguardUpdateMonitorCallback() {
                 @Override
-                public void onKeyguardVisibilityChanged(boolean showing) {
-                    mKeyguardShowing = showing;
-                    updateAnimateScreenOff();
-                }
-
-                @Override
                 public void onTimeChanged() {
                     if (BURN_IN_TESTING_ENABLED && mStatusBarStateController != null
                             && mStatusBarStateController.get().isDozing()) {
@@ -88,11 +72,6 @@ public class DozeUi implements DozeMachine.Part, TunerService.Tunable,
                         mHandler.post(mWakeLock.wrap(() -> {}));
                     }
                 }
-
-                @Override
-                public void onShadeExpandedChanged(boolean expanded) {
-                    updateAnimateScreenOff();
-                }
             };
 
     private long mLastTimeTickElapsed = 0;
@@ -101,9 +80,7 @@ public class DozeUi implements DozeMachine.Part, TunerService.Tunable,
     public DozeUi(Context context, AlarmManager alarmManager,
             WakeLock wakeLock, DozeHost host, @Main Handler handler,
             DozeParameters params, KeyguardUpdateMonitor keyguardUpdateMonitor,
-            DozeLog dozeLog, TunerService tunerService,
-            Lazy<StatusBarStateController> statusBarStateController,
-            ConfigurationController configurationController) {
+            DozeLog dozeLog, Lazy<StatusBarStateController> statusBarStateController) {
         mContext = context;
         mWakeLock = wakeLock;
         mHost = host;
@@ -113,19 +90,7 @@ public class DozeUi implements DozeMachine.Part, TunerService.Tunable,
         mTimeTicker = new AlarmTimeout(alarmManager, this::onTimeTick, "doze_time_tick", handler);
         keyguardUpdateMonitor.registerCallback(mKeyguardVisibilityCallback);
         mDozeLog = dozeLog;
-        mTunerService = tunerService;
         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
@@ -133,21 +98,6 @@ public class DozeUi implements DozeMachine.Part, TunerService.Tunable,
         mMachine = dozeMachine;
     }
 
-    /**
-     * Decide if we're taking over the screen-off animation
-     * when the device was configured to skip doze after screen off.
-     */
-    private void updateAnimateScreenOff() {
-        if (mCanAnimateTransition) {
-            final boolean controlScreenOff =
-                    mDozeParameters.getAlwaysOn()
-                    && (mKeyguardShowing || mDozeParameters.shouldControlUnlockedScreenOff())
-                    && !mHost.isPowerSaveActive();
-            mDozeParameters.setControlScreenOffAnimation(controlScreenOff);
-            mHost.setAnimateScreenOff(controlScreenOff);
-        }
-    }
-
     private void pulseWhileDozing(int reason) {
         mHost.pulseWhileDozing(
                 new DozeHost.PulseCallback() {
@@ -276,21 +226,4 @@ public class DozeUi implements DozeMachine.Part, TunerService.Tunable,
 
         scheduleTimeTick();
     }
-
-    @VisibleForTesting
-    KeyguardUpdateMonitorCallback getKeyguardCallback() {
-        return mKeyguardVisibilityCallback;
-    }
-
-    @Override
-    public void onTuningChanged(String key, String newValue) {
-        if (key.equals(Settings.Secure.DOZE_ALWAYS_ON)) {
-            updateAnimateScreenOff();
-        }
-    }
-
-    @Override
-    public void onConfigChanged(Configuration newConfig) {
-        updateAnimateScreenOff();
-    }
 }
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 49e3fe7df2be..b5beca44c5e1 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeParameters.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeParameters.java
@@ -16,6 +16,7 @@
 
 package com.android.systemui.statusbar.phone;
 
+import android.content.res.Configuration;
 import android.content.res.Resources;
 import android.hardware.display.AmbientDisplayConfiguration;
 import android.os.PowerManager;
@@ -26,7 +27,10 @@ import android.util.Log;
 import android.util.MathUtils;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.VisibleForTesting;
 
+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;
@@ -35,7 +39,9 @@ import com.android.systemui.doze.AlwaysOnDisplayPolicy;
 import com.android.systemui.doze.DozeScreenState;
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.flags.FeatureFlags;
+import com.android.systemui.plugins.statusbar.StatusBarStateController;
 import com.android.systemui.statusbar.policy.BatteryController;
+import com.android.systemui.statusbar.policy.ConfigurationController;
 import com.android.systemui.statusbar.policy.DevicePostureController;
 import com.android.systemui.tuner.TunerService;
 
@@ -53,7 +59,8 @@ import javax.inject.Inject;
 public class DozeParameters implements
         TunerService.Tunable,
         com.android.systemui.plugins.statusbar.DozeParameters,
-        Dumpable {
+        Dumpable, ConfigurationController.ConfigurationListener,
+        StatusBarStateController.StateListener {
     private static final int MAX_DURATION = 60 * 1000;
     public static final boolean FORCE_NO_BLANKING =
             SystemProperties.getBoolean("debug.force_no_blanking", false);
@@ -74,6 +81,22 @@ public class DozeParameters implements
     private boolean mDozeAlwaysOn;
     private boolean mControlScreenOffAnimation;
 
+    private boolean mKeyguardShowing;
+    @VisibleForTesting
+    final KeyguardUpdateMonitorCallback mKeyguardVisibilityCallback =
+            new KeyguardUpdateMonitorCallback() {
+                @Override
+                public void onKeyguardVisibilityChanged(boolean showing) {
+                    mKeyguardShowing = showing;
+                    updateControlScreenOff();
+                }
+
+                @Override
+                public void onShadeExpandedChanged(boolean expanded) {
+                    updateControlScreenOff();
+                }
+            };
+
     @Inject
     protected DozeParameters(
             @Main Resources resources,
@@ -84,7 +107,10 @@ public class DozeParameters implements
             TunerService tunerService,
             DumpManager dumpManager,
             FeatureFlags featureFlags,
-            UnlockedScreenOffAnimationController unlockedScreenOffAnimationController) {
+            UnlockedScreenOffAnimationController unlockedScreenOffAnimationController,
+            KeyguardUpdateMonitor keyguardUpdateMonitor,
+            ConfigurationController configurationController,
+            StatusBarStateController statusBarStateController) {
         mResources = resources;
         mAmbientDisplayConfiguration = ambientDisplayConfiguration;
         mAlwaysOnPolicy = alwaysOnDisplayPolicy;
@@ -97,10 +123,13 @@ public class DozeParameters implements
         mFeatureFlags = featureFlags;
         mUnlockedScreenOffAnimationController = unlockedScreenOffAnimationController;
 
+        keyguardUpdateMonitor.registerCallback(mKeyguardVisibilityCallback);
         tunerService.addTunable(
                 this,
                 Settings.Secure.DOZE_ALWAYS_ON,
                 Settings.Secure.ACCESSIBILITY_DISPLAY_INVERSION_ENABLED);
+        configurationController.addCallback(this);
+        statusBarStateController.addCallback(this);
     }
 
     public boolean getDisplayStateSupported() {
@@ -222,12 +251,28 @@ public class DozeParameters implements
     }
 
     /**
+     *
+     */
+    public void updateControlScreenOff() {
+        if (!getDisplayNeedsBlanking()) {
+            final boolean controlScreenOff =
+                    getAlwaysOn() && (mKeyguardShowing || shouldControlUnlockedScreenOff());
+            setControlScreenOffAnimation(controlScreenOff);
+        }
+    }
+
+    /**
      * Whether we want to control the screen off animation when the device is unlocked. If we do,
      * we'll animate in AOD before turning off the screen, rather than simply fading to black and
      * then abruptly showing AOD.
+     *
+     * There are currently several reasons we might not want to control the screen off even if we
+     * are able to, such as the shade being expanded, being in landscape, or having animations
+     * disabled for a11y.
      */
     public boolean shouldControlUnlockedScreenOff() {
-        return mUnlockedScreenOffAnimationController.shouldPlayUnlockedScreenOffAnimation();
+        return canControlUnlockedScreenOff()
+                && mUnlockedScreenOffAnimationController.shouldPlayUnlockedScreenOffAnimation();
     }
 
     /**
@@ -308,12 +353,27 @@ public class DozeParameters implements
     @Override
     public void onTuningChanged(String key, String newValue) {
         mDozeAlwaysOn = mAmbientDisplayConfiguration.alwaysOnEnabled(UserHandle.USER_CURRENT);
+
+        if (key.equals(Settings.Secure.DOZE_ALWAYS_ON)) {
+            updateControlScreenOff();
+        }
+
         for (Callback callback : mCallbacks) {
             callback.onAlwaysOnChange();
         }
     }
 
     @Override
+    public void onConfigChanged(Configuration newConfig) {
+        updateControlScreenOff();
+    }
+
+    @Override
+    public void onStatePostChange() {
+        updateControlScreenOff();
+    }
+
+    @Override
     public void dump(@NonNull FileDescriptor fd, @NonNull PrintWriter pw, @NonNull String[] args) {
         pw.print("getAlwaysOn(): "); pw.println(getAlwaysOn());
         pw.print("getDisplayStateSupported(): "); pw.println(getDisplayStateSupported());
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeServiceHost.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeServiceHost.java
index 57b9c03ce576..a88a3b6392c2 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeServiceHost.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeServiceHost.java
@@ -63,7 +63,6 @@ public final class DozeServiceHost implements DozeHost {
     private final DozeLog mDozeLog;
     private final PowerManager mPowerManager;
     private boolean mAnimateWakeup;
-    private boolean mAnimateScreenOff;
     private boolean mIgnoreTouchWhilePulsing;
     private Runnable mPendingScreenOffCallback;
     @VisibleForTesting
@@ -357,11 +356,6 @@ public final class DozeServiceHost implements DozeHost {
     }
 
     @Override
-    public void setAnimateScreenOff(boolean animateScreenOff) {
-        mAnimateScreenOff = animateScreenOff;
-    }
-
-    @Override
     public void onSlpiTap(float screenX, float screenY) {
         if (screenX > 0 && screenY > 0 && mAmbientIndicationContainer != null
                 && mAmbientIndicationContainer.getVisibility() == View.VISIBLE) {
@@ -440,10 +434,6 @@ public final class DozeServiceHost implements DozeHost {
         return mAnimateWakeup;
     }
 
-    boolean shouldAnimateScreenOff() {
-        return mAnimateScreenOff;
-    }
-
     boolean getIgnoreTouchWhilePulsing() {
         return mIgnoreTouchWhilePulsing;
     }
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 b2be9948fea9..7a68b9616fa3 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
@@ -3190,7 +3190,7 @@ public class StatusBar extends SystemUI implements
         boolean wakeAndUnlock = mBiometricUnlockController.getMode()
                 == BiometricUnlockController.MODE_WAKE_AND_UNLOCK;
         boolean animate = (!mDozing && mDozeServiceHost.shouldAnimateWakeup() && !wakeAndUnlock)
-                || (mDozing && mDozeServiceHost.shouldAnimateScreenOff()
+                || (mDozing && mDozeParameters.shouldControlScreenOff()
                 && visibleNotOccludedOrWillBe);
 
         mNotificationPanelViewController.setDozing(mDozing, animate, mWakeUpTouchLocation);
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 cddde6487ba2..83ab8cb80b29 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/UnlockedScreenOffAnimationController.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/UnlockedScreenOffAnimationController.kt
@@ -226,10 +226,6 @@ class UnlockedScreenOffAnimationController @Inject constructor(
             return false
         }
 
-        if (!dozeParameters.get().canControlUnlockedScreenOff()) {
-            return false
-        }
-
         // If animations are disabled system-wide, don't play this one either.
         if (Settings.Global.getString(
                 context.contentResolver, Settings.Global.ANIMATOR_DURATION_SCALE) == "0") {
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 1d34aac3b1cb..519d7d71c07f 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeUiTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeUiTest.java
@@ -27,8 +27,6 @@ import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.clearInvocations;
-import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -91,8 +89,7 @@ public class DozeUiTest extends SysuiTestCase {
         mHandler = mHandlerThread.getThreadHandler();
 
         mDozeUi = new DozeUi(mContext, mAlarmManager, mWakeLock, mHost, mHandler,
-                mDozeParameters, mKeyguardUpdateMonitor, mDozeLog, mTunerService,
-                () -> mStatusBarStateController, mConfigurationController);
+                mDozeParameters, mKeyguardUpdateMonitor, mDozeLog, () -> mStatusBarStateController);
         mDozeUi.setDozeMachine(mMachine);
     }
 
@@ -104,7 +101,7 @@ public class DozeUiTest extends SysuiTestCase {
     }
 
     @Test
-    public void pausingAndUnpausingAod_registersTimeTickAfterUnpausing() throws Exception {
+    public void pausingAndUnpausingAod_registersTimeTickAfterUnpausing() {
         mDozeUi.transitionTo(UNINITIALIZED, INITIALIZED);
         mDozeUi.transitionTo(INITIALIZED, DOZE_AOD);
         mDozeUi.transitionTo(DOZE_AOD, DOZE_AOD_PAUSED);
@@ -117,45 +114,9 @@ public class DozeUiTest extends SysuiTestCase {
     }
 
     @Test
-    public void propagatesAnimateScreenOff_noAlwaysOn() {
-        reset(mHost);
-        when(mDozeParameters.getAlwaysOn()).thenReturn(false);
-        when(mDozeParameters.getDisplayNeedsBlanking()).thenReturn(false);
-
-        mDozeUi.getKeyguardCallback().onKeyguardVisibilityChanged(false);
-        verify(mHost).setAnimateScreenOff(eq(false));
-    }
-
-    @Test
-    public void propagatesAnimateScreenOff_alwaysOn() {
-        reset(mHost);
-        when(mDozeParameters.getAlwaysOn()).thenReturn(true);
-        when(mDozeParameters.getDisplayNeedsBlanking()).thenReturn(false);
-
-        // Take over when the keyguard is visible.
-        mDozeUi.getKeyguardCallback().onKeyguardVisibilityChanged(true);
-        verify(mHost).setAnimateScreenOff(eq(true));
-
-        // Do not animate screen-off when keyguard isn't visible - PowerManager will do it.
-        mDozeUi.getKeyguardCallback().onKeyguardVisibilityChanged(false);
-        verify(mHost).setAnimateScreenOff(eq(false));
-    }
-
-    @Test
-    public void neverAnimateScreenOff_whenNotSupported() {
-        // Re-initialize DozeParameters saying that the display requires blanking.
-        reset(mDozeParameters);
-        reset(mHost);
-        when(mDozeParameters.getDisplayNeedsBlanking()).thenReturn(true);
-        mDozeUi = new DozeUi(mContext, mAlarmManager, mWakeLock, mHost, mHandler,
-                mDozeParameters, mKeyguardUpdateMonitor, mDozeLog, mTunerService,
-                () -> mStatusBarStateController, mConfigurationController);
-        mDozeUi.setDozeMachine(mMachine);
-
-        // Never animate if display doesn't support it.
-        mDozeUi.getKeyguardCallback().onKeyguardVisibilityChanged(true);
-        mDozeUi.getKeyguardCallback().onKeyguardVisibilityChanged(false);
-        verify(mHost, never()).setAnimateScreenOff(eq(false));
+    public void transitionSetsAnimateWakeup_noAlwaysOn() {
+        mDozeUi.transitionTo(UNINITIALIZED, DOZE);
+        verify(mHost).setAnimateWakeup(eq(false));
     }
 
     @Test
@@ -165,49 +126,4 @@ public class DozeUiTest extends SysuiTestCase {
         mDozeUi.transitionTo(UNINITIALIZED, DOZE);
         verify(mHost).setAnimateWakeup(eq(true));
     }
-
-    @Test
-    public void keyguardVisibility_changesControlScreenOffAnimation() {
-        // Pre-condition
-        reset(mDozeParameters);
-        when(mDozeParameters.getAlwaysOn()).thenReturn(true);
-        when(mDozeParameters.getDisplayNeedsBlanking()).thenReturn(false);
-
-        mDozeUi.getKeyguardCallback().onKeyguardVisibilityChanged(false);
-        verify(mDozeParameters).setControlScreenOffAnimation(eq(false));
-        mDozeUi.getKeyguardCallback().onKeyguardVisibilityChanged(true);
-        verify(mDozeParameters).setControlScreenOffAnimation(eq(true));
-    }
-
-    @Test
-    public void transitionSetsAnimateWakeup_noAlwaysOn() {
-        mDozeUi.transitionTo(UNINITIALIZED, DOZE);
-        verify(mHost).setAnimateWakeup(eq(false));
-    }
-
-    @Test
-    public void controlScreenOffTrueWhenKeyguardNotShowingAndControlUnlockedScreenOff() {
-        when(mDozeParameters.getAlwaysOn()).thenReturn(true);
-        when(mDozeParameters.shouldControlUnlockedScreenOff()).thenReturn(true);
-
-        // Tell doze that keyguard is not visible.
-        mDozeUi.getKeyguardCallback().onKeyguardVisibilityChanged(false /* showing */);
-
-        // Since we're controlling the unlocked screen off animation, verify that we've asked to
-        // control the screen off animation despite being unlocked.
-        verify(mDozeParameters).setControlScreenOffAnimation(true);
-    }
-
-    @Test
-    public void controlScreenOffFalseWhenKeyguardNotShowingAndControlUnlockedScreenOffFalse() {
-        when(mDozeParameters.getAlwaysOn()).thenReturn(true);
-        when(mDozeParameters.shouldControlUnlockedScreenOff()).thenReturn(false);
-
-        // Tell doze that keyguard is not visible.
-        mDozeUi.getKeyguardCallback().onKeyguardVisibilityChanged(false /* showing */);
-
-        // Since we're not controlling the unlocked screen off animation, verify that we haven't
-        // asked to control the screen off animation since we're unlocked.
-        verify(mDozeParameters).setControlScreenOffAnimation(false);
-    }
 }
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 7a0b3669991c..f8035ade707f 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
@@ -20,10 +20,9 @@ import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.reset;
-import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.when;
 
 import android.content.res.Resources;
@@ -34,12 +33,15 @@ import android.test.suitebuilder.annotation.SmallTest;
 
 import androidx.test.runner.AndroidJUnit4;
 
+import com.android.keyguard.KeyguardUpdateMonitor;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.doze.AlwaysOnDisplayPolicy;
 import com.android.systemui.doze.DozeScreenState;
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.flags.FeatureFlags;
+import com.android.systemui.plugins.statusbar.StatusBarStateController;
 import com.android.systemui.statusbar.policy.BatteryController;
+import com.android.systemui.statusbar.policy.ConfigurationController;
 import com.android.systemui.tuner.TunerService;
 
 import org.junit.Assert;
@@ -52,7 +54,6 @@ import org.mockito.MockitoAnnotations;
 @SmallTest
 @RunWith(AndroidJUnit4.class)
 public class DozeParametersTest extends SysuiTestCase {
-
     private DozeParameters mDozeParameters;
 
     @Mock Resources mResources;
@@ -64,10 +65,31 @@ public class DozeParametersTest extends SysuiTestCase {
     @Mock private FeatureFlags mFeatureFlags;
     @Mock private DumpManager mDumpManager;
     @Mock private UnlockedScreenOffAnimationController mUnlockedScreenOffAnimationController;
+    @Mock private KeyguardUpdateMonitor mKeyguardUpdateMonitor;
+    @Mock private StatusBarStateController mStatusBarStateController;
+    @Mock private ConfigurationController mConfigurationController;
+
+    /**
+     * The current value of PowerManager's dozeAfterScreenOff property.
+     *
+     * This property controls whether System UI is controlling the screen off animation. If it's
+     * false (PowerManager should not doze after screen off) then System UI is controlling the
+     * animation. If true, we're not controlling it and PowerManager will doze immediately.
+     */
+    private boolean mPowerManagerDozeAfterScreenOff;
 
     @Before
     public void setup() {
         MockitoAnnotations.initMocks(this);
+
+        // Save the current value set for dozeAfterScreenOff so we can make assertions. This method
+        // is only called if the value changes, which makes it difficult to check that it was set
+        // correctly in tests.
+        doAnswer(invocation -> {
+            mPowerManagerDozeAfterScreenOff = invocation.getArgument(0);
+            return mPowerManagerDozeAfterScreenOff;
+        }).when(mPowerManager).setDozeAfterScreenOff(anyBoolean());
+
         mDozeParameters = new DozeParameters(
             mResources,
             mAmbientDisplayConfiguration,
@@ -77,23 +99,30 @@ public class DozeParametersTest extends SysuiTestCase {
             mTunerService,
             mDumpManager,
             mFeatureFlags,
-            mUnlockedScreenOffAnimationController
+            mUnlockedScreenOffAnimationController,
+            mKeyguardUpdateMonitor,
+            mConfigurationController,
+            mStatusBarStateController
         );
+
+        when(mFeatureFlags.useNewLockscreenAnimations()).thenReturn(true);
+
+        setAodEnabledForTest(true);
+        setShouldControlUnlockedScreenOffForTest(true);
+        setDisplayNeedsBlankingForTest(false);
     }
+
     @Test
-    public void testSetControlScreenOffAnimation_setsDozeAfterScreenOff_false() {
+    public void testSetControlScreenOffAnimation_setsDozeAfterScreenOff_correctly() {
+        // If we want to control screen off, we do NOT want PowerManager to doze after screen off.
+        // Obviously.
         mDozeParameters.setControlScreenOffAnimation(true);
-        reset(mPowerManager);
-        mDozeParameters.setControlScreenOffAnimation(false);
-        verify(mPowerManager).setDozeAfterScreenOff(eq(true));
-    }
+        assertFalse(mPowerManagerDozeAfterScreenOff);
 
-    @Test
-    public void testSetControlScreenOffAnimation_setsDozeAfterScreenOff_true() {
+        // If we don't want to control screen off, PowerManager is free to doze after screen off if
+        // that's what'll make it happy.
         mDozeParameters.setControlScreenOffAnimation(false);
-        reset(mPowerManager);
-        mDozeParameters.setControlScreenOffAnimation(true);
-        verify(mPowerManager).setDozeAfterScreenOff(eq(false));
+        assertTrue(mPowerManagerDozeAfterScreenOff);
     }
 
     @Test
@@ -122,37 +151,124 @@ public class DozeParametersTest extends SysuiTestCase {
         assertThat(mDozeParameters.getAlwaysOn()).isFalse();
     }
 
+    /**
+     * PowerManager.setDozeAfterScreenOff(true) means we are not controlling screen off, and calling
+     * it with false means we are. Confusing, but sure - make sure that we call PowerManager with
+     * the correct value depending on whether we want to control screen off.
+     */
     @Test
     public void testControlUnlockedScreenOffAnimation_dozeAfterScreenOff_false() {
-        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.
-        mDozeParameters.setControlScreenOffAnimation(mDozeParameters.shouldControlScreenOff());
+        // If AOD is disabled, we shouldn't want to control screen off. Also, let's double check
+        // that when that value is updated, we called through to PowerManager.
+        setAodEnabledForTest(false);
+        assertFalse(mDozeParameters.shouldControlScreenOff());
+        assertTrue(mPowerManagerDozeAfterScreenOff);
 
-        // We should have asked power manager not to doze after screen off no matter what, since
-        // we're animating and controlling screen off.
-        verify(mPowerManager).setDozeAfterScreenOff(eq(false));
+        // And vice versa...
+        setAodEnabledForTest(true);
+        assertTrue(mDozeParameters.shouldControlScreenOff());
+        assertFalse(mPowerManagerDozeAfterScreenOff);
     }
 
     @Test
     public void testControlUnlockedScreenOffAnimationDisabled_dozeAfterScreenOff() {
-        when(mAmbientDisplayConfiguration.alwaysOnEnabled(anyInt())).thenReturn(true);
-        mDozeParameters.onTuningChanged(Settings.Secure.DOZE_ALWAYS_ON, "1");
+        setShouldControlUnlockedScreenOffForTest(true);
         when(mFeatureFlags.useNewLockscreenAnimations()).thenReturn(false);
 
         assertFalse(mDozeParameters.shouldControlUnlockedScreenOff());
 
         // Trigger the setter for the current value.
         mDozeParameters.setControlScreenOffAnimation(mDozeParameters.shouldControlScreenOff());
+        assertFalse(mDozeParameters.shouldControlScreenOff());
+    }
+
+    @Test
+    public void propagatesAnimateScreenOff_noAlwaysOn() {
+        setAodEnabledForTest(false);
+        setDisplayNeedsBlankingForTest(false);
+
+        mDozeParameters.mKeyguardVisibilityCallback.onKeyguardVisibilityChanged(false);
+        assertFalse(mDozeParameters.shouldControlScreenOff());
+    }
+
+    @Test
+    public void propagatesAnimateScreenOff_alwaysOn() {
+        setAodEnabledForTest(true);
+        setDisplayNeedsBlankingForTest(false);
+        setShouldControlUnlockedScreenOffForTest(false);
+
+        // Take over when the keyguard is visible.
+        mDozeParameters.mKeyguardVisibilityCallback.onKeyguardVisibilityChanged(true);
+        assertTrue(mDozeParameters.shouldControlScreenOff());
 
-        // We should have asked power manager to doze only if we're not controlling screen off
-        // normally.
-        verify(mPowerManager).setDozeAfterScreenOff(
-                eq(!mDozeParameters.shouldControlScreenOff()));
+        // Do not animate screen-off when keyguard isn't visible.
+        mDozeParameters.mKeyguardVisibilityCallback.onKeyguardVisibilityChanged(false);
+        assertFalse(mDozeParameters.shouldControlScreenOff());
+    }
+
+
+    @Test
+    public void neverAnimateScreenOff_whenNotSupported() {
+        setDisplayNeedsBlankingForTest(true);
+
+        // Never animate if display doesn't support it.
+        mDozeParameters.mKeyguardVisibilityCallback.onKeyguardVisibilityChanged(true);
+        assertFalse(mDozeParameters.shouldControlScreenOff());
+        mDozeParameters.mKeyguardVisibilityCallback.onKeyguardVisibilityChanged(false);
+        assertFalse(mDozeParameters.shouldControlScreenOff());
+    }
+
+
+    @Test
+    public void controlScreenOffTrueWhenKeyguardNotShowingAndControlUnlockedScreenOff() {
+        setShouldControlUnlockedScreenOffForTest(true);
+
+        // Tell doze that keyguard is not visible.
+        mDozeParameters.mKeyguardVisibilityCallback.onKeyguardVisibilityChanged(
+                false /* showing */);
+
+        // Since we're controlling the unlocked screen off animation, verify that we've asked to
+        // control the screen off animation despite being unlocked.
+        assertTrue(mDozeParameters.shouldControlScreenOff());
+    }
+
+
+    @Test
+    public void keyguardVisibility_changesControlScreenOffAnimation() {
+        setShouldControlUnlockedScreenOffForTest(false);
+
+        mDozeParameters.mKeyguardVisibilityCallback.onKeyguardVisibilityChanged(false);
+        assertFalse(mDozeParameters.shouldControlScreenOff());
+        mDozeParameters.mKeyguardVisibilityCallback.onKeyguardVisibilityChanged(true);
+        assertTrue(mDozeParameters.shouldControlScreenOff());
+    }
+
+    @Test
+    public void keyguardVisibility_changesControlScreenOffAnimation_respectsUnlockedScreenOff() {
+        setShouldControlUnlockedScreenOffForTest(true);
+
+        // Even if the keyguard is gone, we should control screen off if we can control unlocked
+        // screen off.
+        mDozeParameters.mKeyguardVisibilityCallback.onKeyguardVisibilityChanged(false);
+        assertTrue(mDozeParameters.shouldControlScreenOff());
+
+        mDozeParameters.mKeyguardVisibilityCallback.onKeyguardVisibilityChanged(true);
+        assertTrue(mDozeParameters.shouldControlScreenOff());
+    }
+
+    private void setDisplayNeedsBlankingForTest(boolean needsBlanking) {
+        when(mResources.getBoolean(
+                com.android.internal.R.bool.config_displayBlanksAfterDoze)).thenReturn(
+                        needsBlanking);
+    }
+
+    private void setAodEnabledForTest(boolean enabled) {
+        when(mAmbientDisplayConfiguration.alwaysOnEnabled(anyInt())).thenReturn(enabled);
+        mDozeParameters.onTuningChanged(Settings.Secure.DOZE_ALWAYS_ON, "");
+    }
+
+    private void setShouldControlUnlockedScreenOffForTest(boolean shouldControl) {
+        when(mUnlockedScreenOffAnimationController.shouldPlayUnlockedScreenOffAnimation())
+                .thenReturn(shouldControl);
     }
 }