Make DozeParameters responsible for setting PowerManager#setDozeAfterScreenOff.

Previously this was done in DozeUi, which is a component that is destroyed whenever we
unlock the phone (and re-created when we enter AOD).

Unfortunately, we relied on the fact that the keyguardUpdateMonitorCallback is not
unregistered in DozeUi#destroy, and therefore continued to receive events and
update setDozeAfterScreenOff until being garbage collected. That's obviously very
risky.

We rarely saw issues with this until status bar state was added as a criteria for
whether or not to control the screen off animation (we can't do unlocked screen off
if the shade is expanded). Once that was added, the status bar state at the time of
garbage collection became the permanent control screen off setting until the next doze.

Fixes: 203644049
Test: manual/atest SystemUITests
Change-Id: I85e88e2657363af7a131bd4c31d9a36715593181

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.

packages/SystemUI/src/com/android/systemui/doze/DozeUi.java

@@ -21,30 +21,21 @@ 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.unfold.FoldAodAnimationController;
-import com.android.systemui.unfold.FoldAodAnimationController.FoldAodAnimationStatus;
-import com.android.systemui.unfold.SysUIUnfoldComponent;
 import com.android.systemui.util.AlarmTimeout;
 import com.android.systemui.util.wakelock.WakeLock;
 
 import java.util.Calendar;
-import java.util.Optional;
 
 import javax.inject.Inject;
 
@@ -52,9 +43,7 @@ import javax.inject.Inject;
  * The policy controlling doze.
  */
 @DozeScope
-public class DozeUi implements DozeMachine.Part, TunerService.Tunable,
+public class DozeUi implements DozeMachine.Part {
-        ConfigurationController.ConfigurationListener, FoldAodAnimationStatus,
-        StatusBarStateController.StateListener {
     // 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
@@ -62,25 +51,14 @@ public class DozeUi implements DozeMachine.Part, TunerService.Tunable,
     private final DozeHost mHost;
     private final Handler mHandler;
     private final WakeLock mWakeLock;
-    private final FoldAodAnimationController mFoldAodAnimationController;
     private DozeMachine mMachine;
     private final AlarmTimeout mTimeTicker;
     private final boolean mCanAnimateTransition;
     private final DozeParameters mDozeParameters;
     private final DozeLog mDozeLog;
     private final 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.isDozing()) {
@@ -91,11 +69,6 @@ public class DozeUi implements DozeMachine.Part, TunerService.Tunable,
                         mHandler.post(mWakeLock.wrap(() -> {}));
                     }
                 }
-
-                @Override
-                public void onShadeExpandedChanged(boolean expanded) {
-                    updateAnimateScreenOff();
-                }
             };
 
     private long mLastTimeTickElapsed = 0;
@@ -104,10 +77,8 @@ 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,
             StatusBarStateController statusBarStateController,
-            Optional<SysUIUnfoldComponent> sysUiUnfoldComponent,
+            DozeLog dozeLog) {
-            ConfigurationController configurationController) {
         mContext = context;
         mWakeLock = wakeLock;
         mHost = host;
@@ -117,31 +88,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;
-        mStatusBarStateController.addCallback(this);
-
-        mTunerService.addTunable(this, Settings.Secure.DOZE_ALWAYS_ON);
-
-        mConfigurationController = configurationController;
-        mConfigurationController.addCallback(this);
-
-        mFoldAodAnimationController = sysUiUnfoldComponent
-                .map(SysUIUnfoldComponent::getFoldAodAnimationController).orElse(null);
-
-        if (mFoldAodAnimationController != null) {
-            mFoldAodAnimationController.addCallback(this);
-        }
-    }
-
-    @Override
-    public void destroy() {
-        mTunerService.removeTunable(this);
-        mConfigurationController.removeCallback(this);
-
-        if (mFoldAodAnimationController != null) {
-            mFoldAodAnimationController.removeCallback(this);
-        }
     }
 
     @Override
@@ -149,22 +96,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
-                    && mDozeParameters.shouldAnimateDozingChange());
-        }
-    }
-
     private void pulseWhileDozing(int reason) {
         mHost.pulseWhileDozing(
                 new DozeHost.PulseCallback() {
@@ -293,34 +224,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();
-    }
-
-    /**
-     * Called when StatusBar state changed, could affect unlocked screen off animation state
-     */
-    @Override
-    public void onStatePostChange() {
-        updateAnimateScreenOff();
-    }
-
-    @Override
-    public void onFoldToAodAnimationChanged() {
-        updateAnimateScreenOff();
-    }
 }

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;
@@ -36,13 +40,18 @@ import com.android.systemui.doze.DozeScreenState;
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.flags.Flags;
+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;
+import com.android.systemui.unfold.FoldAodAnimationController;
+import com.android.systemui.unfold.SysUIUnfoldComponent;
 
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
 import java.util.HashSet;
+import java.util.Optional;
 import java.util.Set;
 
 import javax.inject.Inject;
@@ -54,7 +63,8 @@ import javax.inject.Inject;
 public class DozeParameters implements
         TunerService.Tunable,
         com.android.systemui.plugins.statusbar.DozeParameters,
-        Dumpable {
+        Dumpable, ConfigurationController.ConfigurationListener,
+        StatusBarStateController.StateListener, FoldAodAnimationController.FoldAodAnimationStatus {
     private static final int MAX_DURATION = 60 * 1000;
     public static final boolean FORCE_NO_BLANKING =
             SystemProperties.getBoolean("debug.force_no_blanking", false);
@@ -69,12 +79,30 @@ public class DozeParameters implements
     private final BatteryController mBatteryController;
     private final FeatureFlags mFeatureFlags;
     private final ScreenOffAnimationController mScreenOffAnimationController;
+    private final FoldAodAnimationController mFoldAodAnimationController;
+    private final UnlockedScreenOffAnimationController mUnlockedScreenOffAnimationController;
 
     private final Set<Callback> mCallbacks = new HashSet<>();
 
     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,
@@ -85,7 +113,12 @@ public class DozeParameters implements
             TunerService tunerService,
             DumpManager dumpManager,
             FeatureFlags featureFlags,
-            ScreenOffAnimationController screenOffAnimationController) {
+            ScreenOffAnimationController screenOffAnimationController,
+            Optional<SysUIUnfoldComponent> sysUiUnfoldComponent,
+            UnlockedScreenOffAnimationController unlockedScreenOffAnimationController,
+            KeyguardUpdateMonitor keyguardUpdateMonitor,
+            ConfigurationController configurationController,
+            StatusBarStateController statusBarStateController) {
         mResources = resources;
         mAmbientDisplayConfiguration = ambientDisplayConfiguration;
         mAlwaysOnPolicy = alwaysOnDisplayPolicy;
@@ -97,11 +130,22 @@ public class DozeParameters implements
         mPowerManager.setDozeAfterScreenOff(!mControlScreenOffAnimation);
         mFeatureFlags = featureFlags;
         mScreenOffAnimationController = screenOffAnimationController;
+        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);
+
+        mFoldAodAnimationController = sysUiUnfoldComponent
+                .map(SysUIUnfoldComponent::getFoldAodAnimationController).orElse(null);
+
+        if (mFoldAodAnimationController != null) {
+            mFoldAodAnimationController.addCallback(this);
+        }
     }
 
     public boolean getDisplayStateSupported() {
@@ -222,13 +266,26 @@ public class DozeParameters implements
         mPowerManager.setDozeAfterScreenOff(!controlScreenOffAnimation);
     }
 
+    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 mScreenOffAnimationController.shouldControlUnlockedScreenOff();
+        return canControlUnlockedScreenOff()
+                && mUnlockedScreenOffAnimationController.shouldPlayUnlockedScreenOffAnimation();
     }
 
     public boolean shouldDelayKeyguardShow() {
@@ -325,12 +382,32 @@ 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();
         }
         mScreenOffAnimationController.onAlwaysOnChanged(getAlwaysOn());
     }
 
+    @Override
+    public void onConfigChanged(Configuration newConfig) {
+        updateControlScreenOff();
+    }
+
+    @Override
+    public void onStatePostChange() {
+        updateControlScreenOff();
+    }
+
+    @Override
+    public void onFoldToAodAnimationChanged() {
+        updateControlScreenOff();
+    }
+
     @Override
     public void dump(@NonNull FileDescriptor fd, @NonNull PrintWriter pw, @NonNull String[] args) {
         pw.print("getAlwaysOn(): "); pw.println(getAlwaysOn());

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
@@ -356,11 +355,6 @@ public final class DozeServiceHost implements DozeHost {
         mAnimateWakeup = animateWakeup;
     }
 
-    @Override
-    public void setAnimateScreenOff(boolean animateScreenOff) {
-        mAnimateScreenOff = animateScreenOff;
-    }
-
     @Override
     public void onSlpiTap(float screenX, float screenY) {
         if (screenX > 0 && screenY > 0 && mAmbientIndicationContainer != null
@@ -440,10 +434,6 @@ public final class DozeServiceHost implements DozeHost {
         return mAnimateWakeup;
     }
 
-    boolean shouldAnimateScreenOff() {
-        return mAnimateScreenOff;
-    }
-
     boolean getIgnoreTouchWhilePulsing() {
         return mIgnoreTouchWhilePulsing;
     }

packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java

@@ -3162,7 +3162,7 @@ public class StatusBar extends CoreStartable 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);

packages/SystemUI/src/com/android/systemui/statusbar/phone/UnlockedScreenOffAnimationController.kt

@@ -229,10 +229,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") {

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;
 
@@ -43,10 +41,7 @@ import com.android.keyguard.KeyguardUpdateMonitor;
 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.unfold.FoldAodAnimationController;
-import com.android.systemui.unfold.SysUIUnfoldComponent;
 import com.android.systemui.util.wakelock.WakeLockFake;
 
 import org.junit.After;
@@ -56,8 +51,6 @@ import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import java.util.Optional;
-
 @RunWith(AndroidJUnit4.class)
 @SmallTest
 public class DozeUiTest extends SysuiTestCase {
@@ -82,12 +75,6 @@ public class DozeUiTest extends SysuiTestCase {
     private DozeUi mDozeUi;
     @Mock
     private StatusBarStateController mStatusBarStateController;
-    @Mock
-    private FoldAodAnimationController mFoldAodAnimationController;
-    @Mock
-    private SysUIUnfoldComponent mSysUIUnfoldComponent;
-    @Mock
-    private ConfigurationController mConfigurationController;
 
     @Before
     public void setUp() throws Exception {
@@ -98,13 +85,8 @@ public class DozeUiTest extends SysuiTestCase {
         mWakeLock = new WakeLockFake();
         mHandler = mHandlerThread.getThreadHandler();
 
-        when(mSysUIUnfoldComponent.getFoldAodAnimationController())
-                .thenReturn(mFoldAodAnimationController);
-
         mDozeUi = new DozeUi(mContext, mAlarmManager, mWakeLock, mHost, mHandler,
-                mDozeParameters, mKeyguardUpdateMonitor, mDozeLog, mTunerService,
+                mDozeParameters, mKeyguardUpdateMonitor, mStatusBarStateController, mDozeLog);
-                mStatusBarStateController, Optional.of(mSysUIUnfoldComponent),
-                mConfigurationController);
         mDozeUi.setDozeMachine(mMachine);
     }
 
@@ -116,7 +98,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);
@@ -129,60 +111,9 @@ public class DozeUiTest extends SysuiTestCase {
     }
 
     @Test
-    public void propagatesAnimateScreenOff_noAlwaysOn() {
+    public void transitionSetsAnimateWakeup_noAlwaysOn() {
-        reset(mHost);
+        mDozeUi.transitionTo(UNINITIALIZED, DOZE);
-        when(mDozeParameters.getAlwaysOn()).thenReturn(false);
+        verify(mHost).setAnimateWakeup(eq(false));
-        when(mDozeParameters.getDisplayNeedsBlanking()).thenReturn(false);
-        when(mDozeParameters.shouldAnimateDozingChange()).thenReturn(true);
-
-        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);
-        when(mDozeParameters.shouldAnimateDozingChange()).thenReturn(true);
-
-        // 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 propagatesAnimateScreenOff_alwaysOn_shouldAnimateDozingChangeIsFalse() {
-        reset(mHost);
-        when(mDozeParameters.getAlwaysOn()).thenReturn(true);
-        when(mDozeParameters.getDisplayNeedsBlanking()).thenReturn(false);
-        when(mDozeParameters.shouldAnimateDozingChange()).thenReturn(false);
-
-        // Take over when the keyguard is visible.
-        mDozeUi.getKeyguardCallback().onKeyguardVisibilityChanged(true);
-        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, Optional.of(mSysUIUnfoldComponent),
-                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));
     }
 
     @Test
@@ -192,49 +123,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);
-    }
 }

packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/DozeParametersTest.java

@@ -17,11 +17,12 @@
 package com.android.systemui.statusbar.phone;
 
 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.doAnswer;
-import static org.mockito.Mockito.reset;
-import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.content.res.Resources;
@@ -32,14 +33,19 @@ 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.flags.Flags;
+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 com.android.systemui.unfold.FoldAodAnimationController;
+import com.android.systemui.unfold.SysUIUnfoldComponent;
 
 import org.junit.Assert;
 import org.junit.Before;
@@ -48,10 +54,11 @@ import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+import java.util.Optional;
+
 @SmallTest
 @RunWith(AndroidJUnit4.class)
 public class DozeParametersTest extends SysuiTestCase {
-
     private DozeParameters mDozeParameters;
 
     @Mock Resources mResources;
@@ -63,10 +70,37 @@ public class DozeParametersTest extends SysuiTestCase {
     @Mock private FeatureFlags mFeatureFlags;
     @Mock private DumpManager mDumpManager;
     @Mock private ScreenOffAnimationController mScreenOffAnimationController;
+    @Mock private FoldAodAnimationController mFoldAodAnimationController;
+    @Mock private SysUIUnfoldComponent mSysUIUnfoldComponent;
+    @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());
+
+        when(mSysUIUnfoldComponent.getFoldAodAnimationController())
+                .thenReturn(mFoldAodAnimationController);
+
         mDozeParameters = new DozeParameters(
             mResources,
             mAmbientDisplayConfiguration,
@@ -76,23 +110,31 @@ public class DozeParametersTest extends SysuiTestCase {
             mTunerService,
             mDumpManager,
             mFeatureFlags,
-            mScreenOffAnimationController
+            mScreenOffAnimationController,
+            Optional.of(mSysUIUnfoldComponent),
+            mUnlockedScreenOffAnimationController,
+            mKeyguardUpdateMonitor,
+            mConfigurationController,
+            mStatusBarStateController
         );
-    }
+
-    @Test
+        when(mFeatureFlags.isEnabled(Flags.LOCKSCREEN_ANIMATIONS)).thenReturn(true);
-    public void testSetControlScreenOffAnimation_setsDozeAfterScreenOff_false() {
+        setAodEnabledForTest(true);
-        mDozeParameters.setControlScreenOffAnimation(true);
+        setShouldControlUnlockedScreenOffForTest(true);
-        reset(mPowerManager);
+        setDisplayNeedsBlankingForTest(false);
-        mDozeParameters.setControlScreenOffAnimation(false);
-        verify(mPowerManager).setDozeAfterScreenOff(eq(true));
     }
 
     @Test
-    public void testSetControlScreenOffAnimation_setsDozeAfterScreenOff_true() {
+    public void testSetControlScreenOffAnimation_setsDozeAfterScreenOff_correctly() {
-        mDozeParameters.setControlScreenOffAnimation(false);
+        // If we want to control screen off, we do NOT want PowerManager to doze after screen off.
-        reset(mPowerManager);
+        // Obviously.
         mDozeParameters.setControlScreenOffAnimation(true);
-        verify(mPowerManager).setDozeAfterScreenOff(eq(false));
+        assertFalse(mPowerManagerDozeAfterScreenOff);
+
+        // 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);
+        assertTrue(mPowerManagerDozeAfterScreenOff);
     }
 
     @Test
@@ -121,35 +163,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);
+        // If AOD is disabled, we shouldn't want to control screen off. Also, let's double check
-        mDozeParameters.onTuningChanged(Settings.Secure.DOZE_ALWAYS_ON, "1");
+        // that when that value is updated, we called through to PowerManager.
-        when(mFeatureFlags.isEnabled(Flags.LOCKSCREEN_ANIMATIONS)).thenReturn(true);
+        setAodEnabledForTest(false);
-        when(mDozeParameters.shouldControlUnlockedScreenOff()).thenReturn(true);
+        assertFalse(mDozeParameters.shouldControlScreenOff());
+        assertTrue(mPowerManagerDozeAfterScreenOff);
 
-        // Trigger the setter for the current value.
+        // And vice versa...
-        mDozeParameters.setControlScreenOffAnimation(mDozeParameters.shouldControlScreenOff());
+        setAodEnabledForTest(true);
-
+        assertTrue(mDozeParameters.shouldControlScreenOff());
-        // We should have asked power manager not to doze after screen off no matter what, since
+        assertFalse(mPowerManagerDozeAfterScreenOff);
-        // we're animating and controlling screen off.
-        verify(mPowerManager).setDozeAfterScreenOff(eq(false));
     }
 
     @Test
     public void testControlUnlockedScreenOffAnimationDisabled_dozeAfterScreenOff() {
-        when(mAmbientDisplayConfiguration.alwaysOnEnabled(anyInt())).thenReturn(true);
+        setShouldControlUnlockedScreenOffForTest(true);
-        mDozeParameters.onTuningChanged(Settings.Secure.DOZE_ALWAYS_ON, "1");
         when(mFeatureFlags.isEnabled(Flags.LOCKSCREEN_ANIMATIONS)).thenReturn(false);
 
         assertFalse(mDozeParameters.shouldControlUnlockedScreenOff());
 
         // Trigger the setter for the current value.
         mDozeParameters.setControlScreenOffAnimation(mDozeParameters.shouldControlScreenOff());
+        assertFalse(mDozeParameters.shouldControlScreenOff());
+    }
 
-        // We should have asked power manager to doze only if we're not controlling screen off
+    @Test
-        // normally.
+    public void propagatesAnimateScreenOff_noAlwaysOn() {
-        verify(mPowerManager).setDozeAfterScreenOff(
+        setAodEnabledForTest(false);
-                eq(!mDozeParameters.shouldControlScreenOff()));
+        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());
+
+        // 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);
     }
 }