2 files changed, 32 insertions, 4 deletions

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 12fbf2d979f4..a6a9d74a894c 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
@@ -342,7 +342,8 @@ public class StatusBar extends SystemUI implements DemoMode,
     private int mStatusBarWindowState = WINDOW_STATE_SHOWING;
     protected StatusBarWindowController mStatusBarWindowController;
     protected UnlockMethodCache mUnlockMethodCache;
-    private DozeServiceHost mDozeServiceHost = new DozeServiceHost();
+    @VisibleForTesting
+    DozeServiceHost mDozeServiceHost = new DozeServiceHost();
     private boolean mWakeUpComingFromTouch;
     private PointF mWakeUpTouchLocation;
 
@@ -479,7 +480,7 @@ public class StatusBar extends SystemUI implements DemoMode,
     private boolean mLaunchCameraOnScreenTurningOn;
     private boolean mLaunchCameraOnFinishedGoingToSleep;
     private int mLastCameraLaunchSource;
-    private PowerManager.WakeLock mGestureWakeLock;
+    protected PowerManager.WakeLock mGestureWakeLock;
     private Vibrator mVibrator;
     private long[] mCameraLaunchGestureVibePattern;
 
@@ -3609,6 +3610,7 @@ public class StatusBar extends SystemUI implements DemoMode,
         }
     }
 
+    @VisibleForTesting
     final WakefulnessLifecycle.Observer mWakefulnessObserver = new WakefulnessLifecycle.Observer() {
         @Override
         public void onFinishedGoingToSleep() {
@@ -3650,6 +3652,7 @@ public class StatusBar extends SystemUI implements DemoMode,
             mNotificationPanel.setTouchAndAnimationDisabled(false);
             updateVisibleToUser();
             updateIsKeyguard();
+            mDozeServiceHost.stopDozing();
         }
     };
 
@@ -3856,7 +3859,8 @@ public class StatusBar extends SystemUI implements DemoMode,
         return mStatusBarKeyguardViewManager.isShowing();
     }
 
-    private final class DozeServiceHost implements DozeHost {
+    @VisibleForTesting
+    final class DozeServiceHost implements DozeHost {
         private final ArrayList<Callback> mCallbacks = new ArrayList<>();
         private boolean mAnimateWakeup;
         private boolean mAnimateScreenOff;
@@ -3944,7 +3948,6 @@ public class StatusBar extends SystemUI implements DemoMode,
                 mDozingRequested = false;
                 DozeLog.traceDozing(mContext, mDozing);
                 updateDozing();
-                mWakefulnessLifecycle.dispatchStartedWakingUp();
             }
         }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarTest.java
index 882f2615f443..d442de20db1e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarTest.java
@@ -593,6 +593,30 @@ public class StatusBarTest extends SysuiTestCase {
         verify(mStatusBarStateController).setState(eq(StatusBarState.FULLSCREEN_USER_SWITCHER));
     }
 
+    @Test
+    public void testStartStopDozing() {
+        mStatusBar.setBarStateForTest(StatusBarState.KEYGUARD);
+        when(mStatusBarStateController.isKeyguardRequested()).thenReturn(true);
+
+        mStatusBar.mDozeServiceHost.startDozing();
+        verify(mStatusBarStateController).setIsDozing(eq(true));
+
+        mStatusBar.mDozeServiceHost.stopDozing();
+        verify(mStatusBarStateController).setIsDozing(eq(false));
+    }
+
+    @Test
+    public void testOnStartedWakingUp_isNotDozing() {
+        mStatusBar.setBarStateForTest(StatusBarState.KEYGUARD);
+        when(mStatusBarStateController.isKeyguardRequested()).thenReturn(true);
+
+        mStatusBar.mDozeServiceHost.startDozing();
+        verify(mStatusBarStateController).setIsDozing(eq(true));
+
+        mStatusBar.mWakefulnessObserver.onStartedWakingUp();
+        verify(mStatusBarStateController).setIsDozing(eq(false));
+    }
+
     static class TestableStatusBar extends StatusBar {
         public TestableStatusBar(StatusBarKeyguardViewManager man,
                 UnlockMethodCache unlock, KeyguardIndicationController key,
@@ -642,6 +666,7 @@ public class StatusBarTest extends SysuiTestCase {
             mLockscreenUserManager = notificationLockscreenUserManager;
             mCommandQueue = commandQueue;
             mPresenter = notificationPresenter;
+            mGestureWakeLock = mock(PowerManager.WakeLock.class);
         }
 
         private WakefulnessLifecycle createAwakeWakefulnessLifecycle() {