diff --git a/packages/SystemUI/plugin/src/com/android/systemui/plugins/SensorManagerPlugin.java b/packages/SystemUI/plugin/src/com/android/systemui/plugins/SensorManagerPlugin.java
index 88b8dd8..fbd863d 100644
--- a/packages/SystemUI/plugin/src/com/android/systemui/plugins/SensorManagerPlugin.java
+++ b/packages/SystemUI/plugin/src/com/android/systemui/plugins/SensorManagerPlugin.java
@@ -11,13 +11,12 @@
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
- * limitations under the License
+ * limitations under the License.
  */
 
 package com.android.systemui.plugins;
 
-import android.hardware.Sensor;
-import android.hardware.TriggerEventListener;
+import android.hardware.SensorListener;
 
 import com.android.systemui.plugins.annotations.ProvidesInterface;
 
@@ -31,26 +30,30 @@
     int VERSION = 1;
 
     /**
-     * Registers for trigger events from the sensor. Trigger events are one-shot and need to
-     * re-registered in order for them to be fired again.
+     * Registers for sensor events. Events will be sent until the listener is unregistered.
      * @param sensor
      * @param listener
-     * @see android.hardware.SensorManager#requestTriggerSensor(
-     *     android.hardware.TriggerEventListener, android.hardware.Sensor)
+     * @see android.hardware.SensorManager#registerListener(SensorListener, int)
      */
-    void registerTriggerEvent(Sensor sensor, TriggerEventListener listener);
+    void registerListener(Sensor sensor, SensorEventListener listener);
 
     /**
-     * Unregisters trigger events from the sensor.
+     * Unregisters events from the sensor.
      * @param sensor
      * @param listener
      */
-    void unregisterTriggerEvent(Sensor sensor, TriggerEventListener listener);
+    void unregisterListener(Sensor sensor, SensorEventListener listener);
 
-    interface TriggerEventListener {
-        void onTrigger(TriggerEvent event);
+    /**
+     * Listener triggered whenever the Sensor has new data.
+     */
+    interface SensorEventListener {
+        void onSensorChanged(SensorEvent event);
     }
 
+    /**
+     * Sensor that can be defined in a plugin.
+     */
     class Sensor {
         public static final int TYPE_WAKE_LOCK_SCREEN = 1;
         public static final int TYPE_WAKE_DISPLAY = 2;
@@ -67,29 +70,32 @@
         }
     }
 
-    class TriggerEvent {
+    /**
+     * Event sent by a {@link Sensor}.
+     */
+    class SensorEvent {
         Sensor mSensor;
         int mVendorType;
         float[] mValues;
 
         /**
-         * Creates a trigger event
+         * Creates a sensor event.
          * @param sensor The type of sensor, e.g. TYPE_WAKE_LOCK_SCREEN
          * @param vendorType The vendor type, which should be unique for each type of sensor,
          *                   e.g. SINGLE_TAP = 1, DOUBLE_TAP = 2, etc.
          */
-        public TriggerEvent(Sensor sensor, int vendorType) {
+        public SensorEvent(Sensor sensor, int vendorType) {
             this(sensor, vendorType, null);
         }
 
         /**
-         * Creates a trigger event
+         * Creates a sensor event.
          * @param sensor The type of sensor, e.g. TYPE_WAKE_LOCK_SCREEN
          * @param vendorType The vendor type, which should be unique for each type of sensor,
          *                   e.g. SINGLE_TAP = 1, DOUBLE_TAP = 2, etc.
          * @param values Values captured by the sensor.
          */
-        public TriggerEvent(Sensor sensor, int vendorType, float[] values) {
+        public SensorEvent(Sensor sensor, int vendorType, float[] values) {
             mSensor = sensor;
             mVendorType = vendorType;
             mValues = values;
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeLog.java b/packages/SystemUI/src/com/android/systemui/doze/DozeLog.java
index eda3c59..4fb1bc5 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeLog.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeLog.java
@@ -44,8 +44,7 @@
     public static final int PULSE_REASON_SENSOR_PICKUP = 3;
     public static final int PULSE_REASON_SENSOR_DOUBLE_TAP = 4;
     public static final int PULSE_REASON_SENSOR_LONG_PRESS = 5;
-    public static final int PULSE_REASON_SENSOR_WAKE_LOCK_SCREEN = 6;
-    public static final int REASON_SENSOR_WAKE_UP = 7;
+    public static final int REASON_SENSOR_WAKE_UP = 6;
 
     private static boolean sRegisterKeyguardCallback = true;
 
@@ -212,7 +211,6 @@
             case PULSE_REASON_SENSOR_PICKUP: return "pickup";
             case PULSE_REASON_SENSOR_DOUBLE_TAP: return "doubletap";
             case PULSE_REASON_SENSOR_LONG_PRESS: return "longpress";
-            case PULSE_REASON_SENSOR_WAKE_LOCK_SCREEN: return "wakeLockScreen";
             case REASON_SENSOR_WAKE_UP: return "wakeup";
             default: throw new IllegalArgumentException("bad reason: " + pulseReason);
         }
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeSensors.java b/packages/SystemUI/src/com/android/systemui/doze/DozeSensors.java
index 7e77843..c2676d0 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeSensors.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeSensors.java
@@ -17,7 +17,6 @@
 package com.android.systemui.doze;
 
 import static com.android.systemui.plugins.SensorManagerPlugin.Sensor.TYPE_WAKE_DISPLAY;
-import static com.android.systemui.plugins.SensorManagerPlugin.Sensor.TYPE_WAKE_LOCK_SCREEN;
 
 import android.annotation.AnyThread;
 import android.app.ActivityManager;
@@ -26,7 +25,6 @@
 import android.content.Context;
 import android.database.ContentObserver;
 import android.hardware.Sensor;
-import android.hardware.SensorEvent;
 import android.hardware.SensorEventListener;
 import android.hardware.SensorManager;
 import android.hardware.TriggerEvent;
@@ -114,14 +112,7 @@
                         DozeLog.PULSE_REASON_SENSOR_LONG_PRESS,
                         true /* reports touch coordinates */,
                         true /* touchscreen */),
-                new PluginTriggerSensor(
-                        new SensorManagerPlugin.Sensor(TYPE_WAKE_LOCK_SCREEN),
-                        Settings.Secure.DOZE_WAKE_LOCK_SCREEN_GESTURE,
-                        true /* configured */,
-                        DozeLog.PULSE_REASON_SENSOR_WAKE_LOCK_SCREEN,
-                        false /* reports touch coordinates */,
-                        false /* touchscreen */),
-                new PluginTriggerSensor(
+                new PluginSensor(
                         new SensorManagerPlugin.Sensor(TYPE_WAKE_DISPLAY),
                         Settings.Secure.DOZE_WAKE_SCREEN_GESTURE,
                         true /* configured */,
@@ -272,7 +263,7 @@
         }
 
         @Override
-        public void onSensorChanged(SensorEvent event) {
+        public void onSensorChanged(android.hardware.SensorEvent event) {
             if (DEBUG) Log.d(TAG, "onSensorChanged " + event);
 
             mCurrentlyFar = event.values[0] >= event.sensor.getMaximumRange();
@@ -417,7 +408,7 @@
 
         protected String triggerEventToString(TriggerEvent event) {
             if (event == null) return null;
-            final StringBuilder sb = new StringBuilder("TriggerEvent[")
+            final StringBuilder sb = new StringBuilder("SensorEvent[")
                     .append(event.timestamp).append(',')
                     .append(event.sensor.getName());
             if (event.values != null) {
@@ -432,23 +423,19 @@
     /**
      * A Sensor that is injected via plugin.
      */
-    private class PluginTriggerSensor extends TriggerSensor {
+    private class PluginSensor extends TriggerSensor {
 
         private final SensorManagerPlugin.Sensor mPluginSensor;
-        private final SensorManagerPlugin.TriggerEventListener mTriggerEventListener = (event) -> {
+        private final SensorManagerPlugin.SensorEventListener mTriggerEventListener = (event) -> {
             DozeLog.traceSensor(mContext, mPulseReason);
             mHandler.post(mWakeLock.wrap(() -> {
-                if (DEBUG) Log.d(TAG, "onTrigger: " + triggerEventToString(event));
-                mRegistered = false;
+                if (DEBUG) Log.d(TAG, "onSensorEvent: " + triggerEventToString(event));
                 mCallback.onSensorPulse(mPulseReason, true /* sensorPerformsProxCheck */, -1, -1,
                         event.getValues());
-                if (!mRegistered) {
-                    updateListener();  // reregister, this sensor only fires once
-                }
             }));
         };
 
-        PluginTriggerSensor(SensorManagerPlugin.Sensor sensor, String setting, boolean configured,
+        PluginSensor(SensorManagerPlugin.Sensor sensor, String setting, boolean configured,
                 int pulseReason, boolean reportsTouchCoordinates, boolean requiresTouchscreen) {
             super(null, setting, configured, pulseReason, reportsTouchCoordinates,
                     requiresTouchscreen);
@@ -460,13 +447,13 @@
             if (!mConfigured) return;
             AsyncSensorManager asyncSensorManager = (AsyncSensorManager) mSensorManager;
             if (mRequested && !mDisabled && enabledBySetting() && !mRegistered) {
-                asyncSensorManager.requestPluginTriggerSensor(mPluginSensor, mTriggerEventListener);
+                asyncSensorManager.registerPluginListener(mPluginSensor, mTriggerEventListener);
                 mRegistered = true;
-                if (DEBUG) Log.d(TAG, "requestPluginTriggerSensor");
+                if (DEBUG) Log.d(TAG, "registerPluginListener");
             } else if (mRegistered) {
-                asyncSensorManager.cancelPluginTriggerSensor(mPluginSensor, mTriggerEventListener);
+                asyncSensorManager.unregisterPluginListener(mPluginSensor, mTriggerEventListener);
                 mRegistered = false;
-                if (DEBUG) Log.d(TAG, "cancelPluginTriggerSensor");
+                if (DEBUG) Log.d(TAG, "unregisterPluginListener");
             }
         }
 
@@ -479,7 +466,7 @@
                     .append(", mSensor=").append(mPluginSensor).append("}").toString();
         }
 
-        private String triggerEventToString(SensorManagerPlugin.TriggerEvent event) {
+        private String triggerEventToString(SensorManagerPlugin.SensorEvent event) {
             if (event == null) return null;
             final StringBuilder sb = new StringBuilder("PluginTriggerEvent[")
                     .append(event.getSensor()).append(',')
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeTriggers.java b/packages/SystemUI/src/com/android/systemui/doze/DozeTriggers.java
index afe9a74..1da8976 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeTriggers.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeTriggers.java
@@ -53,6 +53,12 @@
     /** adb shell am broadcast -a com.android.systemui.doze.pulse com.android.systemui */
     private static final String PULSE_ACTION = "com.android.systemui.doze.pulse";
 
+    /**
+     * Last value sent by the wake-display sensor.
+     * Assuming that the screen should start on.
+     */
+    private static boolean sWakeDisplaySensorState = true;
+
     private final Context mContext;
     private final DozeMachine mMachine;
     private final DozeSensors mDozeSensors;
@@ -128,7 +134,6 @@
         boolean isDoubleTap = pulseReason == DozeLog.PULSE_REASON_SENSOR_DOUBLE_TAP;
         boolean isPickup = pulseReason == DozeLog.PULSE_REASON_SENSOR_PICKUP;
         boolean isLongPress = pulseReason == DozeLog.PULSE_REASON_SENSOR_LONG_PRESS;
-        boolean isWakeLockScreen = pulseReason == DozeLog.PULSE_REASON_SENSOR_WAKE_LOCK_SCREEN;
         boolean isWakeDisplay = pulseReason == DozeLog.REASON_SENSOR_WAKE_UP;
         boolean wakeEvent = rawValues != null && rawValues.length > 0 && rawValues[0] != 0;
 
@@ -145,14 +150,6 @@
                 if (isDoubleTap) {
                     mDozeHost.onDoubleTap(screenX, screenY);
                     mMachine.wakeUp();
-                } else if (isWakeLockScreen) {
-                    if (wakeEvent) {
-                        mDozeHost.setPassiveInterrupt(true);
-                        mMachine.wakeUp();
-                        DozeLog.traceLockScreenWakeUp(wakeEvent);
-                    } else {
-                        if (DEBUG) Log.d(TAG, "Unpulsing");
-                    }
                 } else if (isPickup) {
                     mDozeHost.setPassiveInterrupt(true);
                     mMachine.wakeUp();
@@ -199,6 +196,7 @@
         DozeMachine.State state = mMachine.getState();
         boolean paused = (state == DozeMachine.State.DOZE_AOD_PAUSED);
         boolean pausing = (state == DozeMachine.State.DOZE_AOD_PAUSING);
+        sWakeDisplaySensorState = wake;
 
         if (wake) {
             proximityCheckThenCall((result) -> {
@@ -234,6 +232,9 @@
                 }
                 mDozeSensors.setListening(true);
                 mDozeHost.setPassiveInterrupt(false);
+                if (newState == DozeMachine.State.DOZE_AOD && !sWakeDisplaySensorState) {
+                    onWakeScreen(false);
+                }
                 break;
             case DOZE_AOD_PAUSED:
             case DOZE_AOD_PAUSING:
diff --git a/packages/SystemUI/src/com/android/systemui/doze/LockScreenWakeUpController.java b/packages/SystemUI/src/com/android/systemui/doze/LockScreenWakeUpController.java
new file mode 100644
index 0000000..ebfafce
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/doze/LockScreenWakeUpController.java
@@ -0,0 +1,121 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.doze;
+
+import android.content.Context;
+import android.os.Handler;
+import android.os.PowerManager;
+import android.os.SystemClock;
+import android.os.UserHandle;
+import android.util.Log;
+
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.hardware.AmbientDisplayConfiguration;
+import com.android.systemui.Dependency;
+import com.android.systemui.plugins.SensorManagerPlugin;
+import com.android.systemui.statusbar.StatusBarState;
+import com.android.systemui.statusbar.StatusBarStateController;
+import com.android.systemui.util.AsyncSensorManager;
+
+/**
+ * Controller responsible for waking up or making the device sleep based on ambient sensors.
+ */
+public class LockScreenWakeUpController implements StatusBarStateController.StateListener,
+        SensorManagerPlugin.SensorEventListener {
+
+    private static final String TAG = LockScreenWakeUpController.class.getSimpleName();
+    private static final boolean DEBUG = Log.isLoggable(TAG, Log.DEBUG);
+
+    private final AsyncSensorManager mAsyncSensorManager;
+    private final SensorManagerPlugin.Sensor mSensor;
+    private final AmbientDisplayConfiguration mAmbientConfiguration;
+    private final PowerManager mPowerManager;
+    private final DozeHost mDozeHost;
+    private final Handler mHandler;
+    private boolean mRegistered;
+    private boolean mDozing;
+
+    public LockScreenWakeUpController(Context context, DozeHost dozeHost) {
+        this(Dependency.get(AsyncSensorManager.class),
+                new SensorManagerPlugin.Sensor(SensorManagerPlugin.Sensor.TYPE_WAKE_LOCK_SCREEN),
+                new AmbientDisplayConfiguration(context),
+                context.getSystemService(PowerManager.class),
+                dozeHost, Dependency.get(StatusBarStateController.class), new Handler());
+    }
+
+    @VisibleForTesting
+    public LockScreenWakeUpController(AsyncSensorManager asyncSensorManager,
+            SensorManagerPlugin.Sensor sensor, AmbientDisplayConfiguration ambientConfiguration,
+            PowerManager powerManager, DozeHost dozeHost,
+            StatusBarStateController statusBarStateController, Handler handler) {
+        mAsyncSensorManager = asyncSensorManager;
+        mSensor = sensor;
+        mAmbientConfiguration = ambientConfiguration;
+        mPowerManager = powerManager;
+        mDozeHost = dozeHost;
+        mHandler = handler;
+        statusBarStateController.addCallback(this);
+    }
+
+    @Override
+    public void onStateChanged(int newState) {
+        boolean isLockScreen = newState == StatusBarState.KEYGUARD
+                || newState == StatusBarState.SHADE_LOCKED;
+
+        if (!mAmbientConfiguration.wakeLockScreenGestureEnabled(UserHandle.USER_CURRENT)) {
+            if (mRegistered) {
+                mAsyncSensorManager.unregisterPluginListener(mSensor, this);
+                mRegistered = false;
+            }
+            return;
+        }
+
+        if (isLockScreen && !mRegistered) {
+            mAsyncSensorManager.registerPluginListener(mSensor, this);
+            mRegistered = true;
+        } else if (!isLockScreen && mRegistered) {
+            mAsyncSensorManager.unregisterPluginListener(mSensor, this);
+            mRegistered = false;
+        }
+    }
+
+    @Override
+    public void onDozingChanged(boolean isDozing) {
+        mDozing = isDozing;
+    }
+
+    @Override
+    public void onSensorChanged(SensorManagerPlugin.SensorEvent event) {
+        mHandler.post(()-> {
+            float[] rawValues = event.getValues();
+            boolean wakeEvent = rawValues != null && rawValues.length > 0 && rawValues[0] != 0;
+
+            DozeLog.traceLockScreenWakeUp(wakeEvent);
+            if (wakeEvent && mDozing) {
+                mDozeHost.setPassiveInterrupt(true);
+                if (DEBUG) Log.d(TAG, "Wake up.");
+                mPowerManager.wakeUp(SystemClock.uptimeMillis(), "com.android.systemui:NODOZE");
+            } else if (!wakeEvent && !mDozing) {
+                if (DEBUG) Log.d(TAG, "Nap time.");
+                mPowerManager.goToSleep(SystemClock.uptimeMillis(),
+                        PowerManager.GO_TO_SLEEP_REASON_SLEEP_BUTTON, 0);
+            } else if (DEBUG) {
+                Log.d(TAG, "Skip sensor event. Wake? " + wakeEvent + " dozing: " + mDozing);
+            }
+        });
+    }
+}
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 a508f1b..6caea1d 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
@@ -148,6 +148,7 @@
 import com.android.systemui.doze.DozeHost;
 import com.android.systemui.doze.DozeLog;
 import com.android.systemui.doze.DozeReceiver;
+import com.android.systemui.doze.LockScreenWakeUpController;
 import com.android.systemui.fragments.ExtensionFragmentListener;
 import com.android.systemui.fragments.FragmentHostManager;
 import com.android.systemui.keyguard.KeyguardViewMediator;
@@ -581,6 +582,7 @@
     protected NotificationPresenter mPresenter;
     private NotificationActivityStarter mNotificationActivityStarter;
     private boolean mPulsing;
+    private LockScreenWakeUpController mLockScreenWakeUpController;
 
     @Override
     public void onActiveStateChanged(int code, int uid, String packageName, boolean active) {
@@ -984,6 +986,7 @@
         for (int i = 0; i < pattern.length; i++) {
             mCameraLaunchGestureVibePattern[i] = pattern[i];
         }
+        mLockScreenWakeUpController = new LockScreenWakeUpController(mContext, mDozeServiceHost);
 
         // receive broadcasts
         IntentFilter filter = new IntentFilter();
diff --git a/packages/SystemUI/src/com/android/systemui/util/AsyncSensorManager.java b/packages/SystemUI/src/com/android/systemui/util/AsyncSensorManager.java
index 0dd8937..88cbbb5 100644
--- a/packages/SystemUI/src/com/android/systemui/util/AsyncSensorManager.java
+++ b/packages/SystemUI/src/com/android/systemui/util/AsyncSensorManager.java
@@ -76,8 +76,9 @@
     }
 
     @Override
-    protected boolean registerListenerImpl(SensorEventListener listener, Sensor sensor, int delayUs,
-            Handler handler, int maxReportLatencyUs, int reservedFlags) {
+    protected boolean registerListenerImpl(SensorEventListener listener,
+            Sensor sensor, int delayUs, Handler handler, int maxReportLatencyUs,
+            int reservedFlags) {
         mHandler.post(() -> {
             if (!mInner.registerListener(listener, sensor, delayUs, maxReportLatencyUs, handler)) {
                 Log.e(TAG, "Registering " + listener + " for " + sensor + " failed.");
@@ -146,23 +147,28 @@
      * @param sensor
      * @param listener
      */
-    public void requestPluginTriggerSensor(SensorManagerPlugin.Sensor sensor,
-            SensorManagerPlugin.TriggerEventListener listener) {
+    public void registerPluginListener(SensorManagerPlugin.Sensor sensor,
+            SensorManagerPlugin.SensorEventListener listener) {
         if (mPlugins.isEmpty()) {
             Log.w(TAG, "No plugins registered");
         }
         mHandler.post(() -> {
             for (int i = 0; i < mPlugins.size(); i++) {
-                mPlugins.get(i).registerTriggerEvent(sensor, listener);
+                mPlugins.get(i).registerListener(sensor, listener);
             }
         });
     }
 
-    public void cancelPluginTriggerSensor(SensorManagerPlugin.Sensor sensor,
-            SensorManagerPlugin.TriggerEventListener listener) {
+    /**
+     * Unregisters all sensors that match the give type for all plugins.
+     * @param sensor
+     * @param listener
+     */
+    public void unregisterPluginListener(SensorManagerPlugin.Sensor sensor,
+            SensorManagerPlugin.SensorEventListener listener) {
         mHandler.post(() -> {
             for (int i = 0; i < mPlugins.size(); i++) {
-                mPlugins.get(i).unregisterTriggerEvent(sensor, listener);
+                mPlugins.get(i).unregisterListener(sensor, listener);
             }
         });
     }
@@ -185,7 +191,8 @@
     }
 
     @Override
-    protected void unregisterListenerImpl(SensorEventListener listener, Sensor sensor) {
+    protected void unregisterListenerImpl(SensorEventListener listener,
+            Sensor sensor) {
         mHandler.post(() -> {
             if (sensor == null) {
                 mInner.unregisterListener(listener);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/doze/LockScreenWakeUpControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/doze/LockScreenWakeUpControllerTest.java
new file mode 100644
index 0000000..8963b59
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/doze/LockScreenWakeUpControllerTest.java
@@ -0,0 +1,142 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.doze;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import android.os.Handler;
+import android.os.PowerManager;
+import android.support.test.filters.SmallTest;
+
+import com.android.internal.hardware.AmbientDisplayConfiguration;
+import com.android.systemui.SysuiTestCase;
+import com.android.systemui.plugins.SensorManagerPlugin;
+import com.android.systemui.statusbar.StatusBarState;
+import com.android.systemui.statusbar.StatusBarStateController;
+import com.android.systemui.util.AsyncSensorManager;
+
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+@RunWith(JUnit4.class)
+@SmallTest
+public class LockScreenWakeUpControllerTest extends SysuiTestCase {
+
+    @Mock
+    private AsyncSensorManager mAsyncSensorManager;
+    @Mock
+    private SensorManagerPlugin.Sensor mSensor;
+    @Mock
+    private AmbientDisplayConfiguration mAmbientDisplayConfiguration;
+    @Mock
+    private PowerManager mPowerManager;
+    @Mock
+    private DozeHost mDozeHost;
+    @Mock
+    private StatusBarStateController mStatusBarStateController;
+    @Mock
+    private Handler mHandler;
+
+    private LockScreenWakeUpController mLockScreenWakeUpController;
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+        doAnswer(invocation -> {
+            ((Runnable) invocation.getArgument(0)).run();
+            return null;
+        }).when(mHandler).post(any());
+
+        mLockScreenWakeUpController = new LockScreenWakeUpController(mAsyncSensorManager, mSensor,
+                mAmbientDisplayConfiguration, mPowerManager, mDozeHost, mStatusBarStateController,
+                mHandler);
+    }
+
+    @Test
+    public void testOnStateChanged_registersUnregistersListener() {
+        when(mAmbientDisplayConfiguration.wakeLockScreenGestureEnabled(anyInt())).thenReturn(true);
+        mLockScreenWakeUpController.onStateChanged(StatusBarState.KEYGUARD);
+        mLockScreenWakeUpController.onStateChanged(StatusBarState.SHADE);
+
+        verify(mAsyncSensorManager, times(1)).registerPluginListener(eq(mSensor),
+                eq(mLockScreenWakeUpController));
+
+        mLockScreenWakeUpController.onStateChanged(StatusBarState.SHADE);
+        verify(mAsyncSensorManager).unregisterPluginListener(eq(mSensor),
+                eq(mLockScreenWakeUpController));
+    }
+
+    @Test
+    public void testOnStateChanged_disabledSensor() {
+        when(mAmbientDisplayConfiguration.wakeLockScreenGestureEnabled(anyInt()))
+                .thenReturn(false);
+        mLockScreenWakeUpController.onStateChanged(StatusBarState.KEYGUARD);
+        mLockScreenWakeUpController.onStateChanged(StatusBarState.SHADE);
+
+        verify(mAsyncSensorManager, never()).registerPluginListener(eq(mSensor),
+                eq(mLockScreenWakeUpController));
+    }
+
+    @Test
+    public void testOnSensorChanged_postsToMainThread() {
+        SensorManagerPlugin.SensorEvent event = new SensorManagerPlugin.SensorEvent(mSensor, 0);
+        mLockScreenWakeUpController.onSensorChanged(event);
+
+        verify(mHandler).post(any());
+    }
+
+    @Test
+    public void testOnSensorChanged_wakeUpWhenDozing() {
+        SensorManagerPlugin.SensorEvent event =
+                new SensorManagerPlugin.SensorEvent(mSensor, 0, new float[] {1});
+        mLockScreenWakeUpController.onSensorChanged(event);
+        verify(mPowerManager, never()).wakeUp(anyLong(), any());
+
+        mLockScreenWakeUpController.onDozingChanged(true);
+        mLockScreenWakeUpController.onSensorChanged(event);
+        verify(mPowerManager).wakeUp(anyLong(), any());
+    }
+
+    @Test
+    public void testOnSensorChanged_sleepsWhenAwake() {
+        boolean[] goToSleep = new boolean[] {false};
+        doAnswer(invocation -> goToSleep[0] = true)
+                .when(mPowerManager).goToSleep(anyLong(), anyInt(), anyInt());
+        SensorManagerPlugin.SensorEvent event =
+                new SensorManagerPlugin.SensorEvent(mSensor, 0, new float[] {0});
+        mLockScreenWakeUpController.onDozingChanged(true);
+        mLockScreenWakeUpController.onSensorChanged(event);
+        Assert.assertFalse("goToSleep should have never been called.", goToSleep[0]);
+
+        mLockScreenWakeUpController.onDozingChanged(false);
+        mLockScreenWakeUpController.onSensorChanged(event);
+        Assert.assertTrue("goToSleep should have been called.", goToSleep[0]);
+    }
+}