Ensure that handles are hidden and stay hidden when not fully awake

Previously handles used either the wakefulness lifecycle or the doze
state to determine whether to hide the handles. This change makes them
use both and hides if either signal indicates that the phone is not
fully awake.

Test: atest google/perf/jank/SystemUI/UbSystemUIJankTests-Lock:android.platform.systemui.tests.jank.SystemUiJankTests#testUnlock -v
Test: atest AssistHandleLikeHomeBehaviorTest
BUG:141219006
FIX:141219006
Change-Id: Iedbf5421c369440a3170f1a0b1add66bd289ec7a
Merged-In: Iedbf5421c369440a3170f1a0b1add66bd289ec7a

3 files changed, 216 insertions, 39 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/assist/AssistHandleLikeHomeBehavior.java b/packages/SystemUI/src/com/android/systemui/assist/AssistHandleLikeHomeBehavior.java
index fa6ffe143022..31115bb940ad 100644
--- a/packages/SystemUI/src/com/android/systemui/assist/AssistHandleLikeHomeBehavior.java
+++ b/packages/SystemUI/src/com/android/systemui/assist/AssistHandleLikeHomeBehavior.java
@@ -22,6 +22,7 @@ import androidx.annotation.Nullable;
 
 import com.android.systemui.assist.AssistHandleBehaviorController.BehaviorController;
 import com.android.systemui.keyguard.WakefulnessLifecycle;
+import com.android.systemui.plugins.statusbar.StatusBarStateController;
 import com.android.systemui.recents.OverviewProxyService;
 import com.android.systemui.shared.system.QuickStepContract;
 
@@ -39,26 +40,33 @@ import dagger.Lazy;
 @Singleton
 final class AssistHandleLikeHomeBehavior implements BehaviorController {
 
+    private final StatusBarStateController.StateListener mStatusBarStateListener =
+            new StatusBarStateController.StateListener() {
+                @Override
+                public void onDozingChanged(boolean isDozing) {
+                    handleDozingChanged(isDozing);
+                }
+            };
     private final WakefulnessLifecycle.Observer mWakefulnessLifecycleObserver =
             new WakefulnessLifecycle.Observer() {
                 @Override
                 public void onStartedWakingUp() {
-                    handleDozingChanged(/* isDozing = */ true);
+                    handleWakefullnessChanged(/* isAwake = */ false);
                 }
 
                 @Override
                 public void onFinishedWakingUp() {
-                    handleDozingChanged(/* isDozing = */ false);
+                    handleWakefullnessChanged(/* isAwake = */ true);
                 }
 
                 @Override
                 public void onStartedGoingToSleep() {
-                    handleDozingChanged(/* isDozing = */ true);
+                    handleWakefullnessChanged(/* isAwake = */ false);
                 }
 
                 @Override
                 public void onFinishedGoingToSleep() {
-                    handleDozingChanged(/* isDozing = */ true);
+                    handleWakefullnessChanged(/* isAwake = */ false);
                 }
             };
     private final OverviewProxyService.OverviewProxyListener mOverviewProxyListener =
@@ -69,19 +77,22 @@ final class AssistHandleLikeHomeBehavior implements BehaviorController {
         }
     };
 
-
+    private final Lazy<StatusBarStateController> mStatusBarStateController;
     private final Lazy<WakefulnessLifecycle> mWakefulnessLifecycle;
     private final Lazy<OverviewProxyService> mOverviewProxyService;
 
     private boolean mIsDozing;
+    private boolean mIsAwake;
     private boolean mIsHomeHandleHiding;
 
     @Nullable private AssistHandleCallbacks mAssistHandleCallbacks;
 
     @Inject
     AssistHandleLikeHomeBehavior(
+            Lazy<StatusBarStateController> statusBarStateController,
             Lazy<WakefulnessLifecycle> wakefulnessLifecycle,
             Lazy<OverviewProxyService> overviewProxyService) {
+        mStatusBarStateController = statusBarStateController;
         mWakefulnessLifecycle = wakefulnessLifecycle;
         mOverviewProxyService = overviewProxyService;
     }
@@ -89,8 +100,10 @@ final class AssistHandleLikeHomeBehavior implements BehaviorController {
     @Override
     public void onModeActivated(Context context, AssistHandleCallbacks callbacks) {
         mAssistHandleCallbacks = callbacks;
-        mIsDozing = mWakefulnessLifecycle.get().getWakefulness()
-                != WakefulnessLifecycle.WAKEFULNESS_AWAKE;
+        mIsDozing = mStatusBarStateController.get().isDozing();
+        mStatusBarStateController.get().addCallback(mStatusBarStateListener);
+        mIsAwake = mWakefulnessLifecycle.get().getWakefulness()
+                == WakefulnessLifecycle.WAKEFULNESS_AWAKE;
         mWakefulnessLifecycle.get().addObserver(mWakefulnessLifecycleObserver);
         mOverviewProxyService.get().addCallback(mOverviewProxyListener);
         callbackForCurrentState();
@@ -99,6 +112,7 @@ final class AssistHandleLikeHomeBehavior implements BehaviorController {
     @Override
     public void onModeDeactivated() {
         mAssistHandleCallbacks = null;
+        mStatusBarStateController.get().removeCallback(mStatusBarStateListener);
         mWakefulnessLifecycle.get().removeObserver(mWakefulnessLifecycleObserver);
         mOverviewProxyService.get().removeCallback(mOverviewProxyListener);
     }
@@ -116,6 +130,15 @@ final class AssistHandleLikeHomeBehavior implements BehaviorController {
         callbackForCurrentState();
     }
 
+    private void handleWakefullnessChanged(boolean isAwake) {
+        if (mIsAwake == isAwake) {
+            return;
+        }
+
+        mIsAwake = isAwake;
+        callbackForCurrentState();
+    }
+
     private void handleSystemUiStateChange(int sysuiStateFlags) {
         boolean isHomeHandleHiding = isHomeHandleHiding(sysuiStateFlags);
         if (mIsHomeHandleHiding == isHomeHandleHiding) {
@@ -131,18 +154,23 @@ final class AssistHandleLikeHomeBehavior implements BehaviorController {
             return;
         }
 
-        if (mIsHomeHandleHiding || mIsDozing) {
+        if (mIsHomeHandleHiding || !isFullyAwake()) {
             mAssistHandleCallbacks.hide();
         } else {
             mAssistHandleCallbacks.showAndStay();
         }
     }
 
+    private boolean isFullyAwake() {
+        return mIsAwake && !mIsDozing;
+    }
+
     @Override
     public void dump(PrintWriter pw, String prefix) {
-        pw.println("Current AssistHandleLikeHomeBehavior State:");
+        pw.println(prefix + "Current AssistHandleLikeHomeBehavior State:");
 
         pw.println(prefix + "   mIsDozing=" + mIsDozing);
+        pw.println(prefix + "   mIsAwake=" + mIsAwake);
         pw.println(prefix + "   mIsHomeHandleHiding=" + mIsHomeHandleHiding);
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/assist/AssistHandleReminderExpBehavior.java b/packages/SystemUI/src/com/android/systemui/assist/AssistHandleReminderExpBehavior.java
index a567315f0296..039404800fbb 100644
--- a/packages/SystemUI/src/com/android/systemui/assist/AssistHandleReminderExpBehavior.java
+++ b/packages/SystemUI/src/com/android/systemui/assist/AssistHandleReminderExpBehavior.java
@@ -93,6 +93,11 @@ final class AssistHandleReminderExpBehavior implements BehaviorController {
                 public void onStateChanged(int newState) {
                     handleStatusBarStateChanged(newState);
                 }
+
+                @Override
+                public void onDozingChanged(boolean isDozing) {
+                    handleDozingChanged(isDozing);
+                }
             };
     private final TaskStackChangeListener mTaskStackChangeListener =
             new TaskStackChangeListener() {
@@ -121,15 +126,25 @@ final class AssistHandleReminderExpBehavior implements BehaviorController {
     private final WakefulnessLifecycle.Observer mWakefulnessLifecycleObserver =
             new WakefulnessLifecycle.Observer() {
                 @Override
+                public void onStartedWakingUp() {
+                    handleWakefullnessChanged(/* isAwake = */ false);
+                }
+
+                @Override
                 public void onFinishedWakingUp() {
-                    handleDozingChanged(false);
+                    handleWakefullnessChanged(/* isAwake = */ true);
                 }
 
                 @Override
                 public void onStartedGoingToSleep() {
-                    handleDozingChanged(true);
+                    handleWakefullnessChanged(/* isAwake = */ false);
                 }
-    };
+
+                @Override
+                public void onFinishedGoingToSleep() {
+                    handleWakefullnessChanged(/* isAwake = */ false);
+                }
+            };
     private final BroadcastReceiver mDefaultHomeBroadcastReceiver = new BroadcastReceiver() {
         @Override
         public void onReceive(Context context, Intent intent) {
@@ -150,6 +165,7 @@ final class AssistHandleReminderExpBehavior implements BehaviorController {
 
     private boolean mOnLockscreen;
     private boolean mIsDozing;
+    private boolean mIsAwake;
     private int mRunningTaskId;
     private boolean mIsNavBarHidden;
     private boolean mIsLauncherShowing;
@@ -200,14 +216,15 @@ final class AssistHandleReminderExpBehavior implements BehaviorController {
         mDefaultHome = getCurrentDefaultHome();
         context.registerReceiver(mDefaultHomeBroadcastReceiver, mDefaultHomeIntentFilter);
         mOnLockscreen = onLockscreen(mStatusBarStateController.get().getState());
+        mIsDozing = mStatusBarStateController.get().isDozing();
         mStatusBarStateController.get().addCallback(mStatusBarStateListener);
         ActivityManager.RunningTaskInfo runningTaskInfo =
                 mActivityManagerWrapper.get().getRunningTask();
         mRunningTaskId = runningTaskInfo == null ? 0 : runningTaskInfo.taskId;
         mActivityManagerWrapper.get().registerTaskStackListener(mTaskStackChangeListener);
         mOverviewProxyService.get().addCallback(mOverviewProxyListener);
-        mIsDozing = mWakefulnessLifecycle.get().getWakefulness()
-                != WakefulnessLifecycle.WAKEFULNESS_AWAKE;
+        mIsAwake = mWakefulnessLifecycle.get().getWakefulness()
+                == WakefulnessLifecycle.WAKEFULNESS_AWAKE;
         mWakefulnessLifecycle.get().addObserver(mWakefulnessLifecycleObserver);
 
         mLearningTimeElapsed = Settings.Secure.getLong(
@@ -249,7 +266,10 @@ final class AssistHandleReminderExpBehavior implements BehaviorController {
 
     @Override
     public void onAssistHandlesRequested() {
-        if (mAssistHandleCallbacks != null && !mIsDozing && !mIsNavBarHidden && !mOnLockscreen) {
+        if (mAssistHandleCallbacks != null
+                && isFullyAwake()
+                && !mIsNavBarHidden
+                && !mOnLockscreen) {
             mAssistHandleCallbacks.showAndGo();
         }
     }
@@ -296,6 +316,16 @@ final class AssistHandleReminderExpBehavior implements BehaviorController {
         callbackForCurrentState(/* justUnlocked = */ false);
     }
 
+    private void handleWakefullnessChanged(boolean isAwake) {
+        if (mIsAwake == isAwake) {
+            return;
+        }
+
+        resetConsecutiveTaskSwitches();
+        mIsAwake = isAwake;
+        callbackForCurrentState(/* justUnlocked = */ false);
+    }
+
     private void handleTaskStackTopChanged(int taskId, @Nullable ComponentName taskComponentName) {
         if (mRunningTaskId == taskId || taskComponentName == null) {
             return;
@@ -349,7 +379,7 @@ final class AssistHandleReminderExpBehavior implements BehaviorController {
             return;
         }
 
-        if (mIsDozing || mIsNavBarHidden || mOnLockscreen || !getShowWhenTaught()) {
+        if (!isFullyAwake() || mIsNavBarHidden || mOnLockscreen || !getShowWhenTaught()) {
             mAssistHandleCallbacks.hide();
         } else if (justUnlocked) {
             long currentEpochDay = LocalDate.now().toEpochDay();
@@ -371,7 +401,7 @@ final class AssistHandleReminderExpBehavior implements BehaviorController {
             return;
         }
 
-        if (mIsDozing || mIsNavBarHidden || isSuppressed()) {
+        if (!isFullyAwake() || mIsNavBarHidden || isSuppressed()) {
             mAssistHandleCallbacks.hide();
         } else if (mOnLockscreen) {
             mAssistHandleCallbacks.showAndStay();
@@ -422,6 +452,10 @@ final class AssistHandleReminderExpBehavior implements BehaviorController {
         mHandler.postDelayed(mResetConsecutiveTaskSwitches, getShowAndGoDelayResetTimeoutMs());
     }
 
+    private boolean isFullyAwake() {
+        return mIsAwake && !mIsDozing;
+    }
+
     private long getLearningTimeMs() {
         return mPhenotypeHelper.getLong(
                 SystemUiDeviceConfigFlags.ASSIST_HANDLES_LEARN_TIME_MS,
@@ -481,6 +515,7 @@ final class AssistHandleReminderExpBehavior implements BehaviorController {
         pw.println(prefix + "Current AssistHandleReminderExpBehavior State:");
         pw.println(prefix + "   mOnLockscreen=" + mOnLockscreen);
         pw.println(prefix + "   mIsDozing=" + mIsDozing);
+        pw.println(prefix + "   mIsAwake=" + mIsAwake);
         pw.println(prefix + "   mRunningTaskId=" + mRunningTaskId);
         pw.println(prefix + "   mDefaultHome=" + mDefaultHome);
         pw.println(prefix + "   mIsNavBarHidden=" + mIsNavBarHidden);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/assist/AssistHandleLikeHomeBehaviorTest.java b/packages/SystemUI/tests/src/com/android/systemui/assist/AssistHandleLikeHomeBehaviorTest.java
index 2aa866e2c179..21ef640d6bf6 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/assist/AssistHandleLikeHomeBehaviorTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/assist/AssistHandleLikeHomeBehaviorTest.java
@@ -30,6 +30,7 @@ import androidx.test.filters.SmallTest;
 
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.keyguard.WakefulnessLifecycle;
+import com.android.systemui.plugins.statusbar.StatusBarStateController;
 import com.android.systemui.recents.OverviewProxyService;
 import com.android.systemui.shared.system.QuickStepContract;
 
@@ -47,6 +48,7 @@ public class AssistHandleLikeHomeBehaviorTest extends SysuiTestCase {
 
     private AssistHandleLikeHomeBehavior mAssistHandleLikeHomeBehavior;
 
+    @Mock private StatusBarStateController mMockStatusBarStateController;
     @Mock private WakefulnessLifecycle mMockWakefulnessLifecycle;
     @Mock private OverviewProxyService mMockOverviewProxyService;
     @Mock private AssistHandleCallbacks mMockAssistHandleCallbacks;
@@ -55,7 +57,9 @@ public class AssistHandleLikeHomeBehaviorTest extends SysuiTestCase {
     public void setup() {
         MockitoAnnotations.initMocks(this);
         mAssistHandleLikeHomeBehavior = new AssistHandleLikeHomeBehavior(
-                () -> mMockWakefulnessLifecycle, () -> mMockOverviewProxyService);
+                () -> mMockStatusBarStateController,
+                () -> mMockWakefulnessLifecycle,
+                () -> mMockOverviewProxyService);
     }
 
     @Test
@@ -66,6 +70,9 @@ public class AssistHandleLikeHomeBehaviorTest extends SysuiTestCase {
         mAssistHandleLikeHomeBehavior.onModeActivated(mContext, mMockAssistHandleCallbacks);
 
         // Assert
+        verify(mMockStatusBarStateController).isDozing();
+        verify(mMockStatusBarStateController).addCallback(
+                any(StatusBarStateController.StateListener.class));
         verify(mMockWakefulnessLifecycle).getWakefulness();
         verify(mMockWakefulnessLifecycle).addObserver(any(WakefulnessLifecycle.Observer.class));
         verify(mMockOverviewProxyService).addCallback(any(
@@ -74,8 +81,9 @@ public class AssistHandleLikeHomeBehaviorTest extends SysuiTestCase {
     }
 
     @Test
-    public void onModeActivated_showsHandlesWhenAwake() {
+    public void onModeActivated_showsHandlesWhenFullyAwake() {
         // Arrange
+        when(mMockStatusBarStateController.isDozing()).thenReturn(false);
         when(mMockWakefulnessLifecycle.getWakefulness())
                 .thenReturn(WakefulnessLifecycle.WAKEFULNESS_AWAKE);
 
@@ -90,6 +98,7 @@ public class AssistHandleLikeHomeBehaviorTest extends SysuiTestCase {
     @Test
     public void onModeActivated_hidesHandlesWhenNotAwake() {
         // Arrange
+        when(mMockStatusBarStateController.isDozing()).thenReturn(true);
         when(mMockWakefulnessLifecycle.getWakefulness())
                 .thenReturn(WakefulnessLifecycle.WAKEFULNESS_ASLEEP);
 
@@ -102,72 +111,139 @@ public class AssistHandleLikeHomeBehaviorTest extends SysuiTestCase {
     }
 
     @Test
+    public void onModeActivated_hidesHandlesWhenDozing() {
+        // Arrange
+        when(mMockStatusBarStateController.isDozing()).thenReturn(true);
+        when(mMockWakefulnessLifecycle.getWakefulness())
+                .thenReturn(WakefulnessLifecycle.WAKEFULNESS_AWAKE);
+
+        // Act
+        mAssistHandleLikeHomeBehavior.onModeActivated(mContext, mMockAssistHandleCallbacks);
+
+        // Assert
+        verify(mMockAssistHandleCallbacks).hide();
+        verifyNoMoreInteractions(mMockAssistHandleCallbacks);
+    }
+
+    @Test
     public void onModeDeactivated_stopsObserving() {
         // Arrange
         mAssistHandleLikeHomeBehavior.onModeActivated(mContext, mMockAssistHandleCallbacks);
+        ArgumentCaptor<StatusBarStateController.StateListener> stateListener =
+                ArgumentCaptor.forClass(StatusBarStateController.StateListener.class);
         ArgumentCaptor<WakefulnessLifecycle.Observer> observer =
                 ArgumentCaptor.forClass(WakefulnessLifecycle.Observer.class);
         ArgumentCaptor<OverviewProxyService.OverviewProxyListener> overviewProxyListener =
                 ArgumentCaptor.forClass(OverviewProxyService.OverviewProxyListener.class);
+        verify(mMockStatusBarStateController).addCallback(stateListener.capture());
         verify(mMockWakefulnessLifecycle).addObserver(observer.capture());
         verify(mMockOverviewProxyService).addCallback(overviewProxyListener.capture());
-        reset(mMockWakefulnessLifecycle, mMockOverviewProxyService, mMockAssistHandleCallbacks);
+        reset(
+                mMockStatusBarStateController,
+                mMockWakefulnessLifecycle,
+                mMockOverviewProxyService,
+                mMockAssistHandleCallbacks);
 
         // Act
         mAssistHandleLikeHomeBehavior.onModeDeactivated();
 
         // Assert
+        verify(mMockStatusBarStateController).removeCallback(eq(stateListener.getValue()));
         verify(mMockWakefulnessLifecycle).removeObserver(eq(observer.getValue()));
         verify(mMockOverviewProxyService).removeCallback(eq(overviewProxyListener.getValue()));
         verifyNoMoreInteractions(
-                mMockWakefulnessLifecycle, mMockOverviewProxyService, mMockAssistHandleCallbacks);
+                mMockStatusBarStateController,
+                mMockWakefulnessLifecycle,
+                mMockOverviewProxyService,
+                mMockAssistHandleCallbacks);
     }
 
     @Test
     public void onAssistantGesturePerformed_doesNothing() {
         // Arrange
         mAssistHandleLikeHomeBehavior.onModeActivated(mContext, mMockAssistHandleCallbacks);
-        reset(mMockWakefulnessLifecycle, mMockOverviewProxyService, mMockAssistHandleCallbacks);
+        reset(
+                mMockStatusBarStateController,
+                mMockWakefulnessLifecycle,
+                mMockOverviewProxyService,
+                mMockAssistHandleCallbacks);
 
         // Act
         mAssistHandleLikeHomeBehavior.onAssistantGesturePerformed();
 
         // Assert
         verifyNoMoreInteractions(
-                mMockWakefulnessLifecycle, mMockOverviewProxyService, mMockAssistHandleCallbacks);
+                mMockStatusBarStateController,
+                mMockWakefulnessLifecycle,
+                mMockOverviewProxyService,
+                mMockAssistHandleCallbacks);
     }
 
     @Test
     public void onAssistHandlesRequested_doesNothing() {
         // Arrange
         mAssistHandleLikeHomeBehavior.onModeActivated(mContext, mMockAssistHandleCallbacks);
-        reset(mMockWakefulnessLifecycle, mMockOverviewProxyService, mMockAssistHandleCallbacks);
+        reset(
+                mMockStatusBarStateController,
+                mMockWakefulnessLifecycle,
+                mMockOverviewProxyService,
+                mMockAssistHandleCallbacks);
 
         // Act
         mAssistHandleLikeHomeBehavior.onAssistHandlesRequested();
 
         // Assert
         verifyNoMoreInteractions(
-                mMockWakefulnessLifecycle, mMockOverviewProxyService, mMockAssistHandleCallbacks);
+                mMockStatusBarStateController,
+                mMockWakefulnessLifecycle,
+                mMockOverviewProxyService,
+                mMockAssistHandleCallbacks);
     }
 
     @Test
-    public void onWake_handlesShow() {
+    public void onBothAwakeAndUnDoze_handlesShow() {
         // Arrange
+        when(mMockStatusBarStateController.isDozing()).thenReturn(true);
         when(mMockWakefulnessLifecycle.getWakefulness())
                 .thenReturn(WakefulnessLifecycle.WAKEFULNESS_ASLEEP);
+        ArgumentCaptor<StatusBarStateController.StateListener> stateListener =
+                ArgumentCaptor.forClass(StatusBarStateController.StateListener.class);
         ArgumentCaptor<WakefulnessLifecycle.Observer> observer =
                 ArgumentCaptor.forClass(WakefulnessLifecycle.Observer.class);
         mAssistHandleLikeHomeBehavior.onModeActivated(mContext, mMockAssistHandleCallbacks);
+        verify(mMockStatusBarStateController).addCallback(stateListener.capture());
         verify(mMockWakefulnessLifecycle).addObserver(observer.capture());
-        reset(mMockWakefulnessLifecycle, mMockOverviewProxyService, mMockAssistHandleCallbacks);
+        reset(
+                mMockStatusBarStateController,
+                mMockWakefulnessLifecycle,
+                mMockOverviewProxyService,
+                mMockAssistHandleCallbacks);
 
         // Act
-        observer.getValue().onStartedWakingUp();
+        observer.getValue().onFinishedWakingUp();
 
         // Assert
+        verify(mMockAssistHandleCallbacks).hide();
         verifyNoMoreInteractions(
-                mMockWakefulnessLifecycle, mMockOverviewProxyService, mMockAssistHandleCallbacks);
+                mMockStatusBarStateController,
+                mMockWakefulnessLifecycle,
+                mMockOverviewProxyService,
+                mMockAssistHandleCallbacks);
+
+        // Arrange
+        observer.getValue().onFinishedGoingToSleep();
+        reset(mMockAssistHandleCallbacks);
+
+        // Act
+        stateListener.getValue().onDozingChanged(false);
+
+        // Assert
+        verify(mMockAssistHandleCallbacks).hide();
+        verifyNoMoreInteractions(
+                mMockStatusBarStateController,
+                mMockWakefulnessLifecycle,
+                mMockOverviewProxyService,
+                mMockAssistHandleCallbacks);
 
         // Act
         observer.getValue().onFinishedWakingUp();
@@ -175,19 +251,30 @@ public class AssistHandleLikeHomeBehaviorTest extends SysuiTestCase {
         // Assert
         verify(mMockAssistHandleCallbacks).showAndStay();
         verifyNoMoreInteractions(
-                mMockWakefulnessLifecycle, mMockOverviewProxyService, mMockAssistHandleCallbacks);
+                mMockStatusBarStateController,
+                mMockWakefulnessLifecycle,
+                mMockOverviewProxyService,
+                mMockAssistHandleCallbacks);
     }
 
     @Test
-    public void onSleep_handlesHide() {
+    public void onSleepOrDoze_handlesHide() {
         // Arrange
+        when(mMockStatusBarStateController.isDozing()).thenReturn(false);
         when(mMockWakefulnessLifecycle.getWakefulness())
                 .thenReturn(WakefulnessLifecycle.WAKEFULNESS_AWAKE);
+        ArgumentCaptor<StatusBarStateController.StateListener> stateListener =
+                ArgumentCaptor.forClass(StatusBarStateController.StateListener.class);
         ArgumentCaptor<WakefulnessLifecycle.Observer> observer =
                 ArgumentCaptor.forClass(WakefulnessLifecycle.Observer.class);
         mAssistHandleLikeHomeBehavior.onModeActivated(mContext, mMockAssistHandleCallbacks);
+        verify(mMockStatusBarStateController).addCallback(stateListener.capture());
         verify(mMockWakefulnessLifecycle).addObserver(observer.capture());
-        reset(mMockWakefulnessLifecycle, mMockOverviewProxyService, mMockAssistHandleCallbacks);
+        reset(
+                mMockStatusBarStateController,
+                mMockWakefulnessLifecycle,
+                mMockOverviewProxyService,
+                mMockAssistHandleCallbacks);
 
         // Act
         observer.getValue().onStartedGoingToSleep();
@@ -195,26 +282,42 @@ public class AssistHandleLikeHomeBehaviorTest extends SysuiTestCase {
         // Assert
         verify(mMockAssistHandleCallbacks).hide();
         verifyNoMoreInteractions(
-                mMockWakefulnessLifecycle, mMockOverviewProxyService, mMockAssistHandleCallbacks);
+                mMockStatusBarStateController,
+                mMockWakefulnessLifecycle,
+                mMockOverviewProxyService,
+                mMockAssistHandleCallbacks);
+
+        // Arrange
+        observer.getValue().onFinishedWakingUp();
+        reset(mMockAssistHandleCallbacks);
 
         // Act
-        observer.getValue().onFinishedGoingToSleep();
+        stateListener.getValue().onDozingChanged(true);
 
         // Assert
+        verify(mMockAssistHandleCallbacks).hide();
         verifyNoMoreInteractions(
-                mMockWakefulnessLifecycle, mMockOverviewProxyService, mMockAssistHandleCallbacks);
+                mMockStatusBarStateController,
+                mMockWakefulnessLifecycle,
+                mMockOverviewProxyService,
+                mMockAssistHandleCallbacks);
     }
 
     @Test
     public void onHomeHandleHide_handlesHide() {
         // Arrange
+        when(mMockStatusBarStateController.isDozing()).thenReturn(false);
         when(mMockWakefulnessLifecycle.getWakefulness())
                 .thenReturn(WakefulnessLifecycle.WAKEFULNESS_AWAKE);
         ArgumentCaptor<OverviewProxyService.OverviewProxyListener> sysUiStateCallback =
                 ArgumentCaptor.forClass(OverviewProxyService.OverviewProxyListener.class);
         mAssistHandleLikeHomeBehavior.onModeActivated(mContext, mMockAssistHandleCallbacks);
         verify(mMockOverviewProxyService).addCallback(sysUiStateCallback.capture());
-        reset(mMockWakefulnessLifecycle, mMockOverviewProxyService, mMockAssistHandleCallbacks);
+        reset(
+                mMockStatusBarStateController,
+                mMockWakefulnessLifecycle,
+                mMockOverviewProxyService,
+                mMockAssistHandleCallbacks);
 
         // Act
         sysUiStateCallback.getValue().onSystemUiStateChanged(
@@ -223,12 +326,16 @@ public class AssistHandleLikeHomeBehaviorTest extends SysuiTestCase {
         // Assert
         verify(mMockAssistHandleCallbacks).hide();
         verifyNoMoreInteractions(
-                mMockWakefulnessLifecycle, mMockOverviewProxyService, mMockAssistHandleCallbacks);
+                mMockStatusBarStateController,
+                mMockWakefulnessLifecycle,
+                mMockOverviewProxyService,
+                mMockAssistHandleCallbacks);
     }
 
     @Test
     public void onHomeHandleUnhide_handlesShow() {
         // Arrange
+        when(mMockStatusBarStateController.isDozing()).thenReturn(false);
         when(mMockWakefulnessLifecycle.getWakefulness())
                 .thenReturn(WakefulnessLifecycle.WAKEFULNESS_AWAKE);
         ArgumentCaptor<OverviewProxyService.OverviewProxyListener> sysUiStateCallback =
@@ -237,7 +344,11 @@ public class AssistHandleLikeHomeBehaviorTest extends SysuiTestCase {
         verify(mMockOverviewProxyService).addCallback(sysUiStateCallback.capture());
         sysUiStateCallback.getValue().onSystemUiStateChanged(
                 QuickStepContract.SYSUI_STATE_NAV_BAR_HIDDEN);
-        reset(mMockWakefulnessLifecycle, mMockOverviewProxyService, mMockAssistHandleCallbacks);
+        reset(
+                mMockStatusBarStateController,
+                mMockWakefulnessLifecycle,
+                mMockOverviewProxyService,
+                mMockAssistHandleCallbacks);
 
         // Act
         sysUiStateCallback.getValue().onSystemUiStateChanged(
@@ -246,6 +357,9 @@ public class AssistHandleLikeHomeBehaviorTest extends SysuiTestCase {
         // Assert
         verify(mMockAssistHandleCallbacks).showAndStay();
         verifyNoMoreInteractions(
-                mMockWakefulnessLifecycle, mMockOverviewProxyService, mMockAssistHandleCallbacks);
+                mMockStatusBarStateController,
+                mMockWakefulnessLifecycle,
+                mMockOverviewProxyService,
+                mMockAssistHandleCallbacks);
     }
 }