Merge "Reorienting wakelockController to have generic exposed functions"
diff --git a/services/core/java/com/android/server/display/DisplayPowerController2.java b/services/core/java/com/android/server/display/DisplayPowerController2.java
index 23c020e..3c1bf0b 100644
--- a/services/core/java/com/android/server/display/DisplayPowerController2.java
+++ b/services/core/java/com/android/server/display/DisplayPowerController2.java
@@ -1138,10 +1138,7 @@
         mHandler.removeCallbacksAndMessages(null);
 
         // Release any outstanding wakelocks we're still holding because of pending messages.
-        mWakelockController.releaseUnfinishedBusinessSuspendBlocker();
-        mWakelockController.releaseStateChangedSuspendBlocker();
-        mWakelockController.releaseProxPositiveSuspendBlocker();
-        mWakelockController.releaseProxNegativeSuspendBlocker();
+        mWakelockController.releaseAll();
 
         final float brightness = mPowerState != null
                 ? mPowerState.getScreenBrightness()
@@ -1729,7 +1726,7 @@
 
         // Grab a wake lock if we have unfinished business.
         if (!finished) {
-            mWakelockController.acquireUnfinishedBusinessSuspendBlocker();
+            mWakelockController.acquireWakelock(WakelockController.WAKE_LOCK_UNFINISHED_BUSINESS);
         }
 
         // Notify the power manager when ready.
@@ -1749,7 +1746,7 @@
 
         // Release the wake lock when we have no unfinished business.
         if (finished) {
-            mWakelockController.releaseUnfinishedBusinessSuspendBlocker();
+            mWakelockController.releaseWakelock(WakelockController.WAKE_LOCK_UNFINISHED_BUSINESS);
         }
 
         // Record if dozing for future comparison.
@@ -2248,7 +2245,8 @@
                 mSensorManager.unregisterListener(mProximitySensorListener);
                 // release wake lock(must be last)
                 boolean proxDebounceSuspendBlockerReleased =
-                        mWakelockController.releaseProxDebounceSuspendBlocker();
+                        mWakelockController.releaseWakelock(
+                                WakelockController.WAKE_LOCK_PROXIMITY_DEBOUNCE);
                 if (proxDebounceSuspendBlockerReleased) {
                     mPendingProximityDebounceTime = -1;
                 }
@@ -2272,11 +2270,13 @@
             if (positive) {
                 mPendingProximity = PROXIMITY_POSITIVE;
                 mPendingProximityDebounceTime = time + PROXIMITY_SENSOR_POSITIVE_DEBOUNCE_DELAY;
-                mWakelockController.acquireProxDebounceSuspendBlocker(); // acquire wake lock
+                mWakelockController.acquireWakelock(
+                        WakelockController.WAKE_LOCK_PROXIMITY_DEBOUNCE); // acquire wake lock
             } else {
                 mPendingProximity = PROXIMITY_NEGATIVE;
                 mPendingProximityDebounceTime = time + PROXIMITY_SENSOR_NEGATIVE_DEBOUNCE_DELAY;
-                mWakelockController.acquireProxDebounceSuspendBlocker(); // acquire wake lock
+                mWakelockController.acquireWakelock(
+                        WakelockController.WAKE_LOCK_PROXIMITY_DEBOUNCE); // acquire wake lock
             }
 
             // Debounce the new sensor reading.
@@ -2300,7 +2300,8 @@
                 updatePowerState();
                 // (must be last)
                 boolean proxDebounceSuspendBlockerReleased =
-                        mWakelockController.releaseProxDebounceSuspendBlocker();
+                        mWakelockController.releaseWakelock(
+                                WakelockController.WAKE_LOCK_PROXIMITY_DEBOUNCE);
                 if (proxDebounceSuspendBlockerReleased) {
                     mPendingProximityDebounceTime = -1;
                 }
@@ -2315,7 +2316,8 @@
     }
 
     private void sendOnStateChangedWithWakelock() {
-        boolean wakeLockAcquired = mWakelockController.acquireStateChangedSuspendBlocker();
+        boolean wakeLockAcquired = mWakelockController.acquireWakelock(
+                WakelockController.WAKE_LOCK_STATE_CHANGED);
         if (wakeLockAcquired) {
             mHandler.post(mWakelockController.getOnStateChangedRunnable());
         }
@@ -2479,13 +2481,13 @@
     }
 
     private void sendOnProximityPositiveWithWakelock() {
-        mWakelockController.acquireProxPositiveSuspendBlocker();
+        mWakelockController.acquireWakelock(WakelockController.WAKE_LOCK_PROXIMITY_POSITIVE);
         mHandler.post(mWakelockController.getOnProximityPositiveRunnable());
     }
 
 
     private void sendOnProximityNegativeWithWakelock() {
-        mWakelockController.acquireProxNegativeSuspendBlocker();
+        mWakelockController.acquireWakelock(WakelockController.WAKE_LOCK_PROXIMITY_NEGATIVE);
         mHandler.post(mWakelockController.getOnProximityNegativeRunnable());
     }
 
diff --git a/services/core/java/com/android/server/display/WakelockController.java b/services/core/java/com/android/server/display/WakelockController.java
index cbf559f..6511f4f 100644
--- a/services/core/java/com/android/server/display/WakelockController.java
+++ b/services/core/java/com/android/server/display/WakelockController.java
@@ -16,12 +16,15 @@
 
 package com.android.server.display;
 
+import android.annotation.IntDef;
 import android.hardware.display.DisplayManagerInternal;
 import android.util.Slog;
 
 import com.android.internal.annotations.VisibleForTesting;
 
 import java.io.PrintWriter;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
 
 /**
  * A utility class to acquire/release suspend blockers and manage appropriate states around it.
@@ -29,8 +32,26 @@
  * display states as needed.
  */
 public final class WakelockController {
+    public static final int WAKE_LOCK_PROXIMITY_POSITIVE = 1;
+    public static final int WAKE_LOCK_PROXIMITY_NEGATIVE = 2;
+    public static final int WAKE_LOCK_PROXIMITY_DEBOUNCE = 3;
+    public static final int WAKE_LOCK_STATE_CHANGED = 4;
+    public static final int WAKE_LOCK_UNFINISHED_BUSINESS = 5;
+
+    private static final int WAKE_LOCK_MAX = WAKE_LOCK_UNFINISHED_BUSINESS;
     private static final boolean DEBUG = false;
 
+    @IntDef(flag = true, prefix = "WAKE_LOCK_", value = {
+            WAKE_LOCK_PROXIMITY_POSITIVE,
+            WAKE_LOCK_PROXIMITY_NEGATIVE,
+            WAKE_LOCK_PROXIMITY_DEBOUNCE,
+            WAKE_LOCK_STATE_CHANGED,
+            WAKE_LOCK_UNFINISHED_BUSINESS
+    })
+    @Retention(RetentionPolicy.SOURCE)
+    public @interface WAKE_LOCK_TYPE {
+    }
+
     // Asynchronous callbacks into the power manager service.
     // Only invoked from the handler thread while no locks are held.
     private final DisplayManagerInternal.DisplayPowerCallbacks mDisplayPowerCallbacks;
@@ -58,17 +79,17 @@
     // (i.e. DisplayPowerController2).
     private boolean mOnStateChangedPending;
 
-    // Count of positive proximity messages currently held. Used to keep track of how many
-    // suspend blocker acquisitions are pending when shutting down the DisplayPowerController2.
-    // Should only be accessed on the Handler thread of the class managing the Display states
-    // (i.e. DisplayPowerController2).
-    private int mOnProximityPositiveMessages;
+    // When true, it implies that a positive proximity wakelock is currently held. Used to keep
+    // track if suspend blocker acquisitions is pending when shutting down the
+    // DisplayPowerController2. Should only be accessed on the Handler thread of the class
+    // managing the Display states (i.e. DisplayPowerController2).
+    private boolean mIsProximityPositiveAcquired;
 
-    // Count of negative proximity messages currently held. Used to keep track of how many
-    // suspend blocker acquisitions are pending when shutting down the DisplayPowerController2.
-    // Should only be accessed on the Handler thread of the class managing the Display states
-    // (i.e. DisplayPowerController2).
-    private int mOnProximityNegativeMessages;
+    // When true, it implies that a negative proximity wakelock is currently held. Used to keep
+    // track if suspend blocker acquisitions is pending when shutting down the
+    // DisplayPowerController2. Should only be accessed on the Handler thread of the class
+    // managing the Display states (i.e. DisplayPowerController2).
+    private boolean mIsProximityNegativeAcquired;
 
     /**
      * The constructor of WakelockController. Manages the initialization of all the local entities
@@ -87,9 +108,86 @@
     }
 
     /**
+     * A utility to acquire a wakelock
+     *
+     * @param wakelock The type of Wakelock to be acquired
+     * @return True of the wakelock is successfully acquired. False if it is already acquired
+     */
+    public boolean acquireWakelock(@WAKE_LOCK_TYPE int wakelock) {
+        return acquireWakelockInternal(wakelock);
+    }
+
+    /**
+     * A utility to release a wakelock
+     *
+     * @param wakelock The type of Wakelock to be released
+     * @return True of an acquired wakelock is successfully released. False if it is already
+     * acquired
+     */
+    public boolean releaseWakelock(@WAKE_LOCK_TYPE int wakelock) {
+        return releaseWakelockInternal(wakelock);
+    }
+
+    /**
+     * A utility to release all the wakelock acquired by the system
+     */
+    public void releaseAll() {
+        for (int i = WAKE_LOCK_PROXIMITY_POSITIVE; i < WAKE_LOCK_MAX; i++) {
+            releaseWakelockInternal(i);
+        }
+    }
+
+    private boolean acquireWakelockInternal(@WAKE_LOCK_TYPE int wakelock) {
+        switch (wakelock) {
+            case WAKE_LOCK_PROXIMITY_POSITIVE:
+                return acquireProxPositiveSuspendBlocker();
+            case WAKE_LOCK_PROXIMITY_NEGATIVE:
+                return acquireProxNegativeSuspendBlocker();
+            case WAKE_LOCK_PROXIMITY_DEBOUNCE:
+                return acquireProxDebounceSuspendBlocker();
+            case WAKE_LOCK_STATE_CHANGED:
+                return acquireStateChangedSuspendBlocker();
+            case WAKE_LOCK_UNFINISHED_BUSINESS:
+                return acquireUnfinishedBusinessSuspendBlocker();
+            default:
+                throw new RuntimeException("Invalid wakelock attempted to be acquired");
+        }
+    }
+
+    private boolean releaseWakelockInternal(@WAKE_LOCK_TYPE int wakelock) {
+        switch (wakelock) {
+            case WAKE_LOCK_PROXIMITY_POSITIVE:
+                return releaseProxPositiveSuspendBlocker();
+            case WAKE_LOCK_PROXIMITY_NEGATIVE:
+                return releaseProxNegativeSuspendBlocker();
+            case WAKE_LOCK_PROXIMITY_DEBOUNCE:
+                return releaseProxDebounceSuspendBlocker();
+            case WAKE_LOCK_STATE_CHANGED:
+                return releaseStateChangedSuspendBlocker();
+            case WAKE_LOCK_UNFINISHED_BUSINESS:
+                return releaseUnfinishedBusinessSuspendBlocker();
+            default:
+                throw new RuntimeException("Invalid wakelock attempted to be released");
+        }
+    }
+
+    /**
+     * Acquires the proximity positive wakelock and notifies the PowerManagerService about the
+     * changes.
+     */
+    private boolean acquireProxPositiveSuspendBlocker() {
+        if (!mIsProximityPositiveAcquired) {
+            mDisplayPowerCallbacks.acquireSuspendBlocker(mSuspendBlockerIdProxPositive);
+            mIsProximityPositiveAcquired = true;
+            return true;
+        }
+        return false;
+    }
+
+    /**
      * Acquires the state change wakelock and notifies the PowerManagerService about the changes.
      */
-    public boolean acquireStateChangedSuspendBlocker() {
+    private boolean acquireStateChangedSuspendBlocker() {
         // Grab a wake lock if we have change of the display state
         if (!mOnStateChangedPending) {
             if (DEBUG) {
@@ -105,18 +203,20 @@
     /**
      * Releases the state change wakelock and notifies the PowerManagerService about the changes.
      */
-    public void releaseStateChangedSuspendBlocker() {
+    private boolean releaseStateChangedSuspendBlocker() {
         if (mOnStateChangedPending) {
             mDisplayPowerCallbacks.releaseSuspendBlocker(mSuspendBlockerIdOnStateChanged);
             mOnStateChangedPending = false;
+            return true;
         }
+        return false;
     }
 
     /**
      * Acquires the unfinished business wakelock and notifies the PowerManagerService about the
      * changes.
      */
-    public void acquireUnfinishedBusinessSuspendBlocker() {
+    private boolean acquireUnfinishedBusinessSuspendBlocker() {
         // Grab a wake lock if we have unfinished business.
         if (!mUnfinishedBusiness) {
             if (DEBUG) {
@@ -124,79 +224,84 @@
             }
             mDisplayPowerCallbacks.acquireSuspendBlocker(mSuspendBlockerIdUnfinishedBusiness);
             mUnfinishedBusiness = true;
+            return true;
         }
+        return false;
     }
 
     /**
      * Releases the unfinished business wakelock and notifies the PowerManagerService about the
      * changes.
      */
-    public void releaseUnfinishedBusinessSuspendBlocker() {
+    private boolean releaseUnfinishedBusinessSuspendBlocker() {
         if (mUnfinishedBusiness) {
             if (DEBUG) {
                 Slog.d(mTag, "Finished business...");
             }
             mDisplayPowerCallbacks.releaseSuspendBlocker(mSuspendBlockerIdUnfinishedBusiness);
             mUnfinishedBusiness = false;
+            return true;
         }
-    }
-
-    /**
-     * Acquires the proximity positive wakelock and notifies the PowerManagerService about the
-     * changes.
-     */
-    public void acquireProxPositiveSuspendBlocker() {
-        mDisplayPowerCallbacks.acquireSuspendBlocker(mSuspendBlockerIdProxPositive);
-        mOnProximityPositiveMessages++;
+        return false;
     }
 
     /**
      * Releases the proximity positive wakelock and notifies the PowerManagerService about the
      * changes.
      */
-    public void releaseProxPositiveSuspendBlocker() {
-        for (int i = 0; i < mOnProximityPositiveMessages; i++) {
+    private boolean releaseProxPositiveSuspendBlocker() {
+        if (mIsProximityPositiveAcquired) {
             mDisplayPowerCallbacks.releaseSuspendBlocker(mSuspendBlockerIdProxPositive);
+            mIsProximityPositiveAcquired = false;
+            return true;
         }
-        mOnProximityPositiveMessages = 0;
+        return false;
     }
 
     /**
      * Acquires the proximity negative wakelock and notifies the PowerManagerService about the
      * changes.
      */
-    public void acquireProxNegativeSuspendBlocker() {
-        mOnProximityNegativeMessages++;
-        mDisplayPowerCallbacks.acquireSuspendBlocker(mSuspendBlockerIdProxNegative);
+    private boolean acquireProxNegativeSuspendBlocker() {
+        if (!mIsProximityNegativeAcquired) {
+            mDisplayPowerCallbacks.acquireSuspendBlocker(mSuspendBlockerIdProxNegative);
+            mIsProximityNegativeAcquired = true;
+            return true;
+        }
+        return false;
     }
 
     /**
      * Releases the proximity negative wakelock and notifies the PowerManagerService about the
      * changes.
      */
-    public void releaseProxNegativeSuspendBlocker() {
-        for (int i = 0; i < mOnProximityNegativeMessages; i++) {
+    private boolean releaseProxNegativeSuspendBlocker() {
+        if (mIsProximityNegativeAcquired) {
             mDisplayPowerCallbacks.releaseSuspendBlocker(mSuspendBlockerIdProxNegative);
+            mIsProximityNegativeAcquired = false;
+            return true;
         }
-        mOnProximityNegativeMessages = 0;
+        return false;
     }
 
     /**
      * Acquires the proximity debounce wakelock and notifies the PowerManagerService about the
      * changes.
      */
-    public void acquireProxDebounceSuspendBlocker() {
+    private boolean acquireProxDebounceSuspendBlocker() {
         if (!mHasProximityDebounced) {
             mDisplayPowerCallbacks.acquireSuspendBlocker(mSuspendBlockerIdProxDebounce);
+            mHasProximityDebounced = true;
+            return true;
         }
-        mHasProximityDebounced = true;
+        return false;
     }
 
     /**
      * Releases the proximity debounce wakelock and notifies the PowerManagerService about the
      * changes.
      */
-    public boolean releaseProxDebounceSuspendBlocker() {
+    private boolean releaseProxDebounceSuspendBlocker() {
         if (mHasProximityDebounced) {
             mDisplayPowerCallbacks.releaseSuspendBlocker(mSuspendBlockerIdProxDebounce);
             mHasProximityDebounced = false;
@@ -210,9 +315,11 @@
      */
     public Runnable getOnProximityPositiveRunnable() {
         return () -> {
-            mOnProximityPositiveMessages--;
-            mDisplayPowerCallbacks.onProximityPositive();
-            mDisplayPowerCallbacks.releaseSuspendBlocker(mSuspendBlockerIdProxPositive);
+            if (mIsProximityPositiveAcquired) {
+                mIsProximityPositiveAcquired = false;
+                mDisplayPowerCallbacks.onProximityPositive();
+                mDisplayPowerCallbacks.releaseSuspendBlocker(mSuspendBlockerIdProxPositive);
+            }
         };
     }
 
@@ -221,9 +328,11 @@
      */
     public Runnable getOnStateChangedRunnable() {
         return () -> {
-            mOnStateChangedPending = false;
-            mDisplayPowerCallbacks.onStateChanged();
-            mDisplayPowerCallbacks.releaseSuspendBlocker(mSuspendBlockerIdOnStateChanged);
+            if (mOnStateChangedPending) {
+                mOnStateChangedPending = false;
+                mDisplayPowerCallbacks.onStateChanged();
+                mDisplayPowerCallbacks.releaseSuspendBlocker(mSuspendBlockerIdOnStateChanged);
+            }
         };
     }
 
@@ -232,9 +341,11 @@
      */
     public Runnable getOnProximityNegativeRunnable() {
         return () -> {
-            mOnProximityNegativeMessages--;
-            mDisplayPowerCallbacks.onProximityNegative();
-            mDisplayPowerCallbacks.releaseSuspendBlocker(mSuspendBlockerIdProxNegative);
+            if (mIsProximityNegativeAcquired) {
+                mIsProximityNegativeAcquired = false;
+                mDisplayPowerCallbacks.onProximityNegative();
+                mDisplayPowerCallbacks.releaseSuspendBlocker(mSuspendBlockerIdProxNegative);
+            }
         };
     }
 
@@ -246,8 +357,8 @@
         pw.println("  mDisplayId=" + mDisplayId);
         pw.println("  mUnfinishedBusiness=" + hasUnfinishedBusiness());
         pw.println("  mOnStateChangePending=" + isOnStateChangedPending());
-        pw.println("  mOnProximityPositiveMessages=" + getOnProximityPositiveMessages());
-        pw.println("  mOnProximityNegativeMessages=" + getOnProximityNegativeMessages());
+        pw.println("  mOnProximityPositiveMessages=" + isProximityPositiveAcquired());
+        pw.println("  mOnProximityNegativeMessages=" + isProximityNegativeAcquired());
     }
 
     @VisibleForTesting
@@ -286,13 +397,13 @@
     }
 
     @VisibleForTesting
-    int getOnProximityPositiveMessages() {
-        return mOnProximityPositiveMessages;
+    boolean isProximityPositiveAcquired() {
+        return mIsProximityPositiveAcquired;
     }
 
     @VisibleForTesting
-    int getOnProximityNegativeMessages() {
-        return mOnProximityNegativeMessages;
+    boolean isProximityNegativeAcquired() {
+        return mIsProximityNegativeAcquired;
     }
 
     @VisibleForTesting
diff --git a/services/tests/mockingservicestests/src/com/android/server/display/DisplayPowerController2Test.java b/services/tests/mockingservicestests/src/com/android/server/display/DisplayPowerController2Test.java
index 2df6823a..20af02e 100644
--- a/services/tests/mockingservicestests/src/com/android/server/display/DisplayPowerController2Test.java
+++ b/services/tests/mockingservicestests/src/com/android/server/display/DisplayPowerController2Test.java
@@ -177,15 +177,19 @@
         advanceTime(1);
 
         // two times, one for unfinished business and one for proximity
-        verify(mWakelockController).acquireUnfinishedBusinessSuspendBlocker();
-        verify(mWakelockController).acquireProxDebounceSuspendBlocker();
+        verify(mWakelockController).acquireWakelock(
+                WakelockController.WAKE_LOCK_UNFINISHED_BUSINESS);
+        verify(mWakelockController).acquireWakelock(
+                WakelockController.WAKE_LOCK_PROXIMITY_DEBOUNCE);
 
 
         dpc.stop();
         advanceTime(1);
         // two times, one for unfinished business and one for proximity
-        verify(mWakelockController).acquireUnfinishedBusinessSuspendBlocker();
-        verify(mWakelockController).acquireProxDebounceSuspendBlocker();
+        verify(mWakelockController).acquireWakelock(
+                WakelockController.WAKE_LOCK_UNFINISHED_BUSINESS);
+        verify(mWakelockController).acquireWakelock(
+                WakelockController.WAKE_LOCK_PROXIMITY_DEBOUNCE);
     }
 
     /**
diff --git a/services/tests/mockingservicestests/src/com/android/server/display/WakelockControllerTest.java b/services/tests/mockingservicestests/src/com/android/server/display/WakelockControllerTest.java
index 288408c..07a81ff 100644
--- a/services/tests/mockingservicestests/src/com/android/server/display/WakelockControllerTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/display/WakelockControllerTest.java
@@ -21,6 +21,7 @@
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyZeroInteractions;
 
 import android.hardware.display.DisplayManagerInternal;
 
@@ -33,6 +34,8 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+import java.util.concurrent.Callable;
+
 @SmallTest
 @RunWith(AndroidJUnit4.class)
 public final class WakelockControllerTest {
@@ -64,25 +67,18 @@
     }
 
     @Test
-    public void acquireStateChangedSuspendBlockerAcquiresIfNotAcquired() {
-        // Acquire the suspend blocker
-        assertTrue(mWakelockController.acquireStateChangedSuspendBlocker());
-        assertTrue(mWakelockController.isOnStateChangedPending());
-
-        // Try to reacquire
-        assertFalse(mWakelockController.acquireStateChangedSuspendBlocker());
-        assertTrue(mWakelockController.isOnStateChangedPending());
+    public void acquireStateChangedSuspendBlockerAcquiresIfNotAcquired() throws Exception {
+        // Acquire
+        verifyWakelockAcquisitionAndReaquisition(WakelockController.WAKE_LOCK_STATE_CHANGED,
+                () -> mWakelockController.isOnStateChangedPending());
 
         // Verify acquire happened only once
         verify(mDisplayPowerCallbacks, times(1))
                 .acquireSuspendBlocker(mWakelockController.getSuspendBlockerOnStateChangedId());
 
         // Release
-        mWakelockController.releaseStateChangedSuspendBlocker();
-        assertFalse(mWakelockController.isOnStateChangedPending());
-
-        // Try to release again
-        mWakelockController.releaseStateChangedSuspendBlocker();
+        verifyWakelockReleaseAndRerelease(WakelockController.WAKE_LOCK_STATE_CHANGED,
+                () -> mWakelockController.isOnStateChangedPending());
 
         // Verify release happened only once
         verify(mDisplayPowerCallbacks, times(1))
@@ -90,25 +86,18 @@
     }
 
     @Test
-    public void acquireUnfinishedBusinessSuspendBlockerAcquiresIfNotAcquired() {
-        // Acquire the suspend blocker
-        mWakelockController.acquireUnfinishedBusinessSuspendBlocker();
-        assertTrue(mWakelockController.hasUnfinishedBusiness());
-
-        // Try to reacquire
-        mWakelockController.acquireUnfinishedBusinessSuspendBlocker();
-        assertTrue(mWakelockController.hasUnfinishedBusiness());
+    public void acquireUnfinishedBusinessSuspendBlockerAcquiresIfNotAcquired() throws Exception {
+        // Acquire
+        verifyWakelockAcquisitionAndReaquisition(WakelockController.WAKE_LOCK_UNFINISHED_BUSINESS,
+                () -> mWakelockController.hasUnfinishedBusiness());
 
         // Verify acquire happened only once
         verify(mDisplayPowerCallbacks, times(1))
                 .acquireSuspendBlocker(mWakelockController.getSuspendBlockerUnfinishedBusinessId());
 
-        // Release the suspend blocker
-        mWakelockController.releaseUnfinishedBusinessSuspendBlocker();
-        assertFalse(mWakelockController.hasUnfinishedBusiness());
-
-        // Try to release again
-        mWakelockController.releaseUnfinishedBusinessSuspendBlocker();
+        // Release
+        verifyWakelockReleaseAndRerelease(WakelockController.WAKE_LOCK_UNFINISHED_BUSINESS,
+                () -> mWakelockController.hasUnfinishedBusiness());
 
         // Verify release happened only once
         verify(mDisplayPowerCallbacks, times(1))
@@ -116,70 +105,56 @@
     }
 
     @Test
-    public void acquireProxPositiveSuspendBlockerAcquiresIfNotAcquired() {
-        // Acquire the suspend blocker
-        mWakelockController.acquireProxPositiveSuspendBlocker();
-        assertEquals(mWakelockController.getOnProximityPositiveMessages(), 1);
-
-        // Try to reacquire
-        mWakelockController.acquireProxPositiveSuspendBlocker();
-        assertEquals(mWakelockController.getOnProximityPositiveMessages(), 2);
+    public void acquireProxPositiveSuspendBlockerAcquiresIfNotAcquired() throws Exception {
+        // Acquire
+        verifyWakelockAcquisitionAndReaquisition(WakelockController.WAKE_LOCK_PROXIMITY_POSITIVE,
+                () -> mWakelockController.isProximityPositiveAcquired());
 
         // Verify acquire happened only once
-        verify(mDisplayPowerCallbacks, times(2))
+        verify(mDisplayPowerCallbacks, times(1))
                 .acquireSuspendBlocker(mWakelockController.getSuspendBlockerProxPositiveId());
 
-        // Release the suspend blocker
-        mWakelockController.releaseProxPositiveSuspendBlocker();
-        assertEquals(mWakelockController.getOnProximityPositiveMessages(), 0);
+        // Release
+        verifyWakelockReleaseAndRerelease(WakelockController.WAKE_LOCK_PROXIMITY_POSITIVE,
+                () -> mWakelockController.isProximityPositiveAcquired());
 
-        // Verify all suspend blockers were released
-        verify(mDisplayPowerCallbacks, times(2))
+        // Verify release happened only once
+        verify(mDisplayPowerCallbacks, times(1))
                 .releaseSuspendBlocker(mWakelockController.getSuspendBlockerProxPositiveId());
     }
 
     @Test
-    public void acquireProxNegativeSuspendBlockerAcquiresIfNotAcquired() {
-        // Acquire the suspend blocker
-        mWakelockController.acquireProxNegativeSuspendBlocker();
-        assertEquals(mWakelockController.getOnProximityNegativeMessages(), 1);
-
-        // Try to reacquire
-        mWakelockController.acquireProxNegativeSuspendBlocker();
-        assertEquals(mWakelockController.getOnProximityNegativeMessages(), 2);
+    public void acquireProxNegativeSuspendBlockerAcquiresIfNotAcquired() throws Exception {
+        // Acquire
+        verifyWakelockAcquisitionAndReaquisition(WakelockController.WAKE_LOCK_PROXIMITY_NEGATIVE,
+                () -> mWakelockController.isProximityNegativeAcquired());
 
         // Verify acquire happened only once
-        verify(mDisplayPowerCallbacks, times(2))
+        verify(mDisplayPowerCallbacks, times(1))
                 .acquireSuspendBlocker(mWakelockController.getSuspendBlockerProxNegativeId());
 
-        // Release the suspend blocker
-        mWakelockController.releaseProxNegativeSuspendBlocker();
-        assertEquals(mWakelockController.getOnProximityNegativeMessages(), 0);
+        // Release
+        verifyWakelockReleaseAndRerelease(WakelockController.WAKE_LOCK_PROXIMITY_NEGATIVE,
+                () -> mWakelockController.isProximityNegativeAcquired());
 
-        // Verify all suspend blockers were released
-        verify(mDisplayPowerCallbacks, times(2))
+        // Verify release happened only once
+        verify(mDisplayPowerCallbacks, times(1))
                 .releaseSuspendBlocker(mWakelockController.getSuspendBlockerProxNegativeId());
     }
 
     @Test
-    public void acquireProxDebounceSuspendBlockerAcquiresIfNotAcquired() {
+    public void acquireProxDebounceSuspendBlockerAcquiresIfNotAcquired() throws Exception {
         // Acquire the suspend blocker
-        mWakelockController.acquireProxDebounceSuspendBlocker();
-
-        // Try to reacquire
-        mWakelockController.acquireProxDebounceSuspendBlocker();
-        assertTrue(mWakelockController.hasProximitySensorDebounced());
+        verifyWakelockAcquisitionAndReaquisition(WakelockController.WAKE_LOCK_PROXIMITY_DEBOUNCE,
+                () -> mWakelockController.hasProximitySensorDebounced());
 
         // Verify acquire happened only once
         verify(mDisplayPowerCallbacks, times(1))
                 .acquireSuspendBlocker(mWakelockController.getSuspendBlockerProxDebounceId());
 
         // Release the suspend blocker
-        assertTrue(mWakelockController.releaseProxDebounceSuspendBlocker());
-
-        // Release again
-        assertFalse(mWakelockController.releaseProxDebounceSuspendBlocker());
-        assertFalse(mWakelockController.hasProximitySensorDebounced());
+        verifyWakelockReleaseAndRerelease(WakelockController.WAKE_LOCK_PROXIMITY_DEBOUNCE,
+                () -> mWakelockController.hasProximitySensorDebounced());
 
         // Verify suspend blocker was released only once
         verify(mDisplayPowerCallbacks, times(1))
@@ -189,52 +164,126 @@
     @Test
     public void proximityPositiveRunnableWorksAsExpected() {
         // Acquire the suspend blocker twice
-        mWakelockController.acquireProxPositiveSuspendBlocker();
-        mWakelockController.acquireProxPositiveSuspendBlocker();
+        assertTrue(mWakelockController.acquireWakelock(
+                WakelockController.WAKE_LOCK_PROXIMITY_POSITIVE));
 
         // Execute the runnable
         Runnable proximityPositiveRunnable = mWakelockController.getOnProximityPositiveRunnable();
         proximityPositiveRunnable.run();
 
         // Validate one suspend blocker was released
-        assertEquals(mWakelockController.getOnProximityPositiveMessages(), 1);
+        assertFalse(mWakelockController.isProximityPositiveAcquired());
         verify(mDisplayPowerCallbacks).onProximityPositive();
         verify(mDisplayPowerCallbacks).releaseSuspendBlocker(
                 mWakelockController.getSuspendBlockerProxPositiveId());
     }
 
     @Test
+    public void proximityPositiveRunnableDoesNothingIfNotAcquired() {
+        // Execute the runnable
+        Runnable proximityPositiveRunnable = mWakelockController.getOnProximityPositiveRunnable();
+        proximityPositiveRunnable.run();
+
+        // Validate one suspend blocker was released
+        assertFalse(mWakelockController.isProximityPositiveAcquired());
+        verifyZeroInteractions(mDisplayPowerCallbacks);
+    }
+
+    @Test
     public void proximityNegativeRunnableWorksAsExpected() {
         // Acquire the suspend blocker twice
-        mWakelockController.acquireProxNegativeSuspendBlocker();
-        mWakelockController.acquireProxNegativeSuspendBlocker();
+        assertTrue(mWakelockController.acquireWakelock(
+                WakelockController.WAKE_LOCK_PROXIMITY_NEGATIVE));
 
         // Execute the runnable
         Runnable proximityNegativeRunnable = mWakelockController.getOnProximityNegativeRunnable();
         proximityNegativeRunnable.run();
 
         // Validate one suspend blocker was released
-        assertEquals(mWakelockController.getOnProximityNegativeMessages(), 1);
+        assertFalse(mWakelockController.isProximityNegativeAcquired());
         verify(mDisplayPowerCallbacks).onProximityNegative();
         verify(mDisplayPowerCallbacks).releaseSuspendBlocker(
                 mWakelockController.getSuspendBlockerProxNegativeId());
     }
 
     @Test
+    public void proximityNegativeRunnableDoesNothingIfNotAcquired() {
+        // Execute the runnable
+        Runnable proximityNegativeRunnable = mWakelockController.getOnProximityNegativeRunnable();
+        proximityNegativeRunnable.run();
+
+        // Validate one suspend blocker was released
+        assertFalse(mWakelockController.isProximityNegativeAcquired());
+        verifyZeroInteractions(mDisplayPowerCallbacks);
+    }
+
+    @Test
     public void onStateChangeRunnableWorksAsExpected() {
         // Acquire the suspend blocker twice
-        mWakelockController.acquireStateChangedSuspendBlocker();
+        assertTrue(mWakelockController.acquireWakelock(WakelockController.WAKE_LOCK_STATE_CHANGED));
 
         // Execute the runnable
         Runnable stateChangeRunnable = mWakelockController.getOnStateChangedRunnable();
         stateChangeRunnable.run();
 
         // Validate one suspend blocker was released
-        assertEquals(mWakelockController.isOnStateChangedPending(), false);
+        assertFalse(mWakelockController.isOnStateChangedPending());
         verify(mDisplayPowerCallbacks).onStateChanged();
         verify(mDisplayPowerCallbacks).releaseSuspendBlocker(
                 mWakelockController.getSuspendBlockerOnStateChangedId());
     }
 
+    @Test
+    public void onStateChangeRunnableDoesNothingIfNotAcquired() {
+        // Execute the runnable
+        Runnable stateChangeRunnable = mWakelockController.getOnStateChangedRunnable();
+        stateChangeRunnable.run();
+
+        // Validate one suspend blocker was released
+        assertFalse(mWakelockController.isOnStateChangedPending());
+        verifyZeroInteractions(mDisplayPowerCallbacks);
+    }
+
+    private void verifyWakelockAcquisitionAndReaquisition(int wakelockId,
+            Callable<Boolean> isWakelockAcquiredCallable)
+            throws Exception {
+        verifyWakelockAcquisition(wakelockId, isWakelockAcquiredCallable);
+        verifyWakelockReacquisition(wakelockId, isWakelockAcquiredCallable);
+    }
+
+    private void verifyWakelockReleaseAndRerelease(int wakelockId,
+            Callable<Boolean> isWakelockAcquiredCallable)
+            throws Exception {
+        verifyWakelockRelease(wakelockId, isWakelockAcquiredCallable);
+        verifyWakelockRerelease(wakelockId, isWakelockAcquiredCallable);
+    }
+
+    private void verifyWakelockAcquisition(int wakelockId,
+            Callable<Boolean> isWakelockAcquiredCallable)
+            throws Exception {
+        assertTrue(mWakelockController.acquireWakelock(wakelockId));
+        assertTrue(isWakelockAcquiredCallable.call());
+    }
+
+    private void verifyWakelockReacquisition(int wakelockId,
+            Callable<Boolean> isWakelockAcquiredCallable)
+            throws Exception {
+        assertFalse(mWakelockController.acquireWakelock(wakelockId));
+        assertTrue(isWakelockAcquiredCallable.call());
+    }
+
+    private void verifyWakelockRelease(int wakelockId, Callable<Boolean> isWakelockAcquiredCallable)
+            throws Exception {
+        assertTrue(mWakelockController.releaseWakelock(wakelockId));
+        assertFalse(isWakelockAcquiredCallable.call());
+    }
+
+    private void verifyWakelockRerelease(int wakelockId,
+            Callable<Boolean> isWakelockAcquiredCallable)
+            throws Exception {
+        assertFalse(mWakelockController.releaseWakelock(wakelockId));
+        assertFalse(isWakelockAcquiredCallable.call());
+    }
+
 
 }