Merge "Simplify the AttentionService API"
diff --git a/api/system-current.txt b/api/system-current.txt
index 7f6a005..6317ae0 100644
--- a/api/system-current.txt
+++ b/api/system-current.txt
@@ -6322,8 +6322,8 @@
     ctor public AttentionService();
     method public final void disableSelf();
     method @Nullable public final android.os.IBinder onBind(@NonNull android.content.Intent);
-    method public abstract void onCancelAttentionCheck(int);
-    method public abstract void onCheckAttention(int, @NonNull android.service.attention.AttentionService.AttentionCallback);
+    method public abstract void onCancelAttentionCheck(@NonNull android.service.attention.AttentionService.AttentionCallback);
+    method public abstract void onCheckAttention(@NonNull android.service.attention.AttentionService.AttentionCallback);
     field public static final int ATTENTION_FAILURE_CAMERA_PERMISSION_ABSENT = 6; // 0x6
     field public static final int ATTENTION_FAILURE_CANCELLED = 3; // 0x3
     field public static final int ATTENTION_FAILURE_PREEMPTED = 4; // 0x4
@@ -6335,8 +6335,8 @@
   }
 
   public static final class AttentionService.AttentionCallback {
-    method public void onFailure(int, int);
-    method public void onSuccess(int, int, long);
+    method public void onFailure(int);
+    method public void onSuccess(int, long);
   }
 
 }
diff --git a/core/java/android/attention/AttentionManagerInternal.java b/core/java/android/attention/AttentionManagerInternal.java
index ac19504..fa3d3b8 100644
--- a/core/java/android/attention/AttentionManagerInternal.java
+++ b/core/java/android/attention/AttentionManagerInternal.java
@@ -30,25 +30,21 @@
     /**
      * Checks whether user attention is at the screen and calls in the provided callback.
      *
-     * @param requestCode   a code associated with the attention check request; this code would be
-     *                      used to call back in {@link AttentionCallbackInternal#onSuccess} and
-     *                      {@link AttentionCallbackInternal#onFailure}
      * @param timeoutMillis a budget for the attention check; if it takes longer - {@link
      *                      AttentionCallbackInternal#onFailure} would be called with the {@link
      *                      android.service.attention.AttentionService#ATTENTION_FAILURE_TIMED_OUT}
      *                      code
      * @param callback      a callback for when the attention check has completed
-     * @return {@code true} if the attention check should succeed; {@false} otherwise.
+     * @return {@code true} if the attention check should succeed.
      */
-    public abstract boolean checkAttention(int requestCode,
-            long timeoutMillis, AttentionCallbackInternal callback);
+    public abstract boolean checkAttention(long timeoutMillis, AttentionCallbackInternal callback);
 
     /**
      * Cancels the specified attention check in case it's no longer needed.
      *
-     * @param requestCode a code provided during {@link #checkAttention}
+     * @param callback a callback that was used in {@link #checkAttention}
      */
-    public abstract void cancelAttentionCheck(int requestCode);
+    public abstract void cancelAttentionCheck(AttentionCallbackInternal callback);
 
     /**
      * Disables the dependants.
@@ -62,19 +58,17 @@
         /**
          * Provides the result of the attention check, if the check was successful.
          *
-         * @param requestCode a code provided in {@link #checkAttention}
          * @param result      an int with the result of the check
          * @param timestamp   a {@code SystemClock.uptimeMillis()} timestamp associated with the
          *                    attention check
          */
-        public abstract void onSuccess(int requestCode, int result, long timestamp);
+        public abstract void onSuccess(int result, long timestamp);
 
         /**
          * Provides the explanation for why the attention check had failed.
          *
-         * @param requestCode a code provided in {@link #checkAttention}
          * @param error       an int with the reason for failure
          */
-        public abstract void onFailure(int requestCode, int error);
+        public abstract void onFailure(int error);
     }
 }
diff --git a/core/java/android/service/attention/AttentionService.java b/core/java/android/service/attention/AttentionService.java
index 84f440f..6172ce5 100644
--- a/core/java/android/service/attention/AttentionService.java
+++ b/core/java/android/service/attention/AttentionService.java
@@ -40,7 +40,7 @@
  * The system's default AttentionService implementation is configured in
  * {@code config_AttentionComponent}. If this config has no value, a stub is returned.
  *
- * See: {@link AttentionManagerService}.
+ * See: {@link com.android.server.attention.AttentionManagerService}.
  *
  * <pre>
  * {@literal
@@ -109,15 +109,16 @@
 
         /** {@inheritDoc} */
         @Override
-        public void checkAttention(int requestCode, IAttentionCallback callback) {
+        public void checkAttention(IAttentionCallback callback) {
             Preconditions.checkNotNull(callback);
-            AttentionService.this.onCheckAttention(requestCode, new AttentionCallback(callback));
+            AttentionService.this.onCheckAttention(new AttentionCallback(callback));
         }
 
         /** {@inheritDoc} */
         @Override
-        public void cancelAttentionCheck(int requestCode) {
-            AttentionService.this.onCancelAttentionCheck(requestCode);
+        public void cancelAttentionCheck(IAttentionCallback callback) {
+            Preconditions.checkNotNull(callback);
+            AttentionService.this.onCancelAttentionCheck(new AttentionCallback(callback));
         }
     };
 
@@ -146,35 +147,43 @@
     /**
      * Checks the user attention and calls into the provided callback.
      *
-     * @param requestCode an identifier that could be used to cancel the request
-     * @param callback    the callback to return the result to
+     * @param callback the callback to return the result to
      */
-    public abstract void onCheckAttention(int requestCode, @NonNull AttentionCallback callback);
+    public abstract void onCheckAttention(@NonNull AttentionCallback callback);
 
-    /** Cancels the attention check for a given request code. */
-    public abstract void onCancelAttentionCheck(int requestCode);
+    /**
+     * Cancels pending work for a given callback.
+     *
+     * Implementation must call back with a failure code of {@link #ATTENTION_FAILURE_CANCELLED}.
+     */
+    public abstract void onCancelAttentionCheck(@NonNull AttentionCallback callback);
 
     /** Callbacks for AttentionService results. */
     public static final class AttentionCallback {
-        private final IAttentionCallback mCallback;
+        @NonNull private final IAttentionCallback mCallback;
 
-        private AttentionCallback(IAttentionCallback callback) {
+        private AttentionCallback(@NonNull IAttentionCallback callback) {
             mCallback = callback;
         }
 
-        /** Returns the result. */
-        public void onSuccess(int requestCode, @AttentionSuccessCodes int result, long timestamp) {
+        /**
+         * Signals a success and provides the result code.
+         *
+         * @param timestamp of when the attention signal was computed; system throttles the requests
+         *                  so this is useful to know how fresh the result is.
+         */
+        public void onSuccess(@AttentionSuccessCodes int result, long timestamp) {
             try {
-                mCallback.onSuccess(requestCode, result, timestamp);
+                mCallback.onSuccess(result, timestamp);
             } catch (RemoteException e) {
                 e.rethrowFromSystemServer();
             }
         }
 
-        /** Signals a failure. */
-        public void onFailure(int requestCode, @AttentionFailureCodes int error) {
+        /** Signals a failure and provides the error code. */
+        public void onFailure(@AttentionFailureCodes int error) {
             try {
-                mCallback.onFailure(requestCode, error);
+                mCallback.onFailure(error);
             } catch (RemoteException e) {
                 e.rethrowFromSystemServer();
             }
diff --git a/core/java/android/service/attention/IAttentionCallback.aidl b/core/java/android/service/attention/IAttentionCallback.aidl
index 0e8a1e7..f65b9c0 100644
--- a/core/java/android/service/attention/IAttentionCallback.aidl
+++ b/core/java/android/service/attention/IAttentionCallback.aidl
@@ -22,6 +22,6 @@
  * @hide
  */
 oneway interface IAttentionCallback {
-    void onSuccess(int requestCode, int result, long timestamp);
-    void onFailure(int requestCode, int error);
+    void onSuccess(int result, long timestamp);
+    void onFailure(int error);
 }
diff --git a/core/java/android/service/attention/IAttentionService.aidl b/core/java/android/service/attention/IAttentionService.aidl
index c3b6f48..99e7997 100644
--- a/core/java/android/service/attention/IAttentionService.aidl
+++ b/core/java/android/service/attention/IAttentionService.aidl
@@ -24,6 +24,6 @@
  * @hide
  */
 oneway interface IAttentionService {
-    void checkAttention(int requestCode, IAttentionCallback callback);
-    void cancelAttentionCheck(int requestCode);
+    void checkAttention(IAttentionCallback callback);
+    void cancelAttentionCheck(IAttentionCallback callback);
 }
\ No newline at end of file
diff --git a/services/core/java/com/android/server/attention/AttentionManagerService.java b/services/core/java/com/android/server/attention/AttentionManagerService.java
index b447c53..1681c5b 100644
--- a/services/core/java/com/android/server/attention/AttentionManagerService.java
+++ b/services/core/java/com/android/server/attention/AttentionManagerService.java
@@ -18,6 +18,7 @@
 
 import static android.provider.DeviceConfig.NAMESPACE_ATTENTION_MANAGER_SERVICE;
 import static android.provider.Settings.System.ADAPTIVE_SLEEP;
+import static android.service.attention.AttentionService.ATTENTION_FAILURE_CANCELLED;
 
 import android.Manifest;
 import android.annotation.NonNull;
@@ -156,9 +157,8 @@
      *
      * @return {@code true} if the framework was able to send the provided callback to the service
      */
-    private boolean checkAttention(int requestCode, long timeout,
-            AttentionCallbackInternal callback) {
-        Preconditions.checkNotNull(callback);
+    private boolean checkAttention(long timeout, AttentionCallbackInternal callbackInternal) {
+        Preconditions.checkNotNull(callbackInternal);
 
         if (!isAttentionServiceSupported()) {
             Slog.w(LOG_TAG, "Trying to call checkAttention() on an unsupported device.");
@@ -172,6 +172,7 @@
 
         synchronized (mLock) {
             final long now = SystemClock.uptimeMillis();
+            // schedule shutting down the connection if no one resets this timer
             freeIfInactiveLocked();
 
             final UserState userState = getOrCreateCurrentUserStateLocked();
@@ -184,46 +185,50 @@
                 // make sure every callback is called back
                 if (userState.mPendingAttentionCheck != null) {
                     userState.mPendingAttentionCheck.cancel(
-                            AttentionService.ATTENTION_FAILURE_UNKNOWN);
+                            ATTENTION_FAILURE_CANCELLED);
                 }
-                userState.mPendingAttentionCheck = new PendingAttentionCheck(requestCode,
-                        callback, () -> checkAttention(requestCode, timeout, callback));
+                // fire the check when the service is started
+                userState.mPendingAttentionCheck = new PendingAttentionCheck(
+                        callbackInternal, () -> checkAttention(timeout, callbackInternal));
             } else {
                 try {
                     // throttle frequent requests
-                    final AttentionCheckCache attentionCheckCache = userState.mAttentionCheckCache;
-                    if (attentionCheckCache != null && now
-                            < attentionCheckCache.mLastComputed + STALE_AFTER_MILLIS) {
-                        callback.onSuccess(requestCode, attentionCheckCache.mResult,
-                                attentionCheckCache.mTimestamp);
+                    final AttentionCheckCache cache = userState.mAttentionCheckCache;
+                    if (cache != null && now < cache.mLastComputed + STALE_AFTER_MILLIS) {
+                        callbackInternal.onSuccess(cache.mResult, cache.mTimestamp);
                         return true;
                     }
 
+                    // schedule request cancellation if not returned by that point yet
                     cancelAfterTimeoutLocked(timeout);
 
-                    userState.mCurrentAttentionCheckRequestCode = requestCode;
-                    userState.mService.checkAttention(requestCode, new IAttentionCallback.Stub() {
-                        @Override
-                        public void onSuccess(int requestCode, int result, long timestamp) {
-                            callback.onSuccess(requestCode, result, timestamp);
-                            synchronized (mLock) {
-                                userState.mAttentionCheckCache = new AttentionCheckCache(
-                                        SystemClock.uptimeMillis(), result,
-                                        timestamp);
-                                userState.mCurrentAttentionCheckIsFulfilled = true;
-                            }
-                            StatsLog.write(StatsLog.ATTENTION_MANAGER_SERVICE_RESULT_REPORTED,
-                                    result);
-                        }
+                    userState.mCurrentAttentionCheck = new AttentionCheck(callbackInternal,
+                            new IAttentionCallback.Stub() {
+                                @Override
+                                public void onSuccess(int result, long timestamp) {
+                                    callbackInternal.onSuccess(result, timestamp);
+                                    synchronized (mLock) {
+                                        userState.mAttentionCheckCache = new AttentionCheckCache(
+                                                SystemClock.uptimeMillis(), result,
+                                                timestamp);
+                                        userState.mCurrentAttentionCheckIsFulfilled = true;
+                                    }
+                                    StatsLog.write(
+                                            StatsLog.ATTENTION_MANAGER_SERVICE_RESULT_REPORTED,
+                                            result);
+                                }
 
-                        @Override
-                        public void onFailure(int requestCode, int error) {
-                            callback.onFailure(requestCode, error);
-                            userState.mCurrentAttentionCheckIsFulfilled = true;
-                            StatsLog.write(StatsLog.ATTENTION_MANAGER_SERVICE_RESULT_REPORTED,
-                                    error);
-                        }
-                    });
+                                @Override
+                                public void onFailure(int error) {
+                                    callbackInternal.onFailure(error);
+                                    userState.mCurrentAttentionCheckIsFulfilled = true;
+                                    StatsLog.write(
+                                            StatsLog.ATTENTION_MANAGER_SERVICE_RESULT_REPORTED,
+                                            error);
+                                }
+                            });
+                    userState.mService.checkAttention(
+                            userState.mCurrentAttentionCheck.mIAttentionCallback);
                 } catch (RemoteException e) {
                     Slog.e(LOG_TAG, "Cannot call into the AttentionService");
                     return false;
@@ -234,7 +239,7 @@
     }
 
     /** Cancels the specified attention check. */
-    private void cancelAttentionCheck(int requestCode) {
+    private void cancelAttentionCheck(AttentionCallbackInternal callbackInternal) {
         synchronized (mLock) {
             final UserState userState = peekCurrentUserStateLocked();
             if (userState == null) {
@@ -242,15 +247,21 @@
             }
             if (userState.mService == null) {
                 if (userState.mPendingAttentionCheck != null
-                        && userState.mPendingAttentionCheck.mRequestCode == requestCode) {
+                        && userState.mPendingAttentionCheck.mCallbackInternal.equals(
+                        callbackInternal)) {
                     userState.mPendingAttentionCheck = null;
                 }
                 return;
             }
-            try {
-                userState.mService.cancelAttentionCheck(requestCode);
-            } catch (RemoteException e) {
-                Slog.e(LOG_TAG, "Cannot call into the AttentionService");
+            if (userState.mCurrentAttentionCheck.mCallbackInternal.equals(callbackInternal)) {
+                try {
+                    userState.mService.cancelAttentionCheck(
+                            userState.mCurrentAttentionCheck.mIAttentionCallback);
+                } catch (RemoteException e) {
+                    Slog.e(LOG_TAG, "Cannot call into the AttentionService");
+                }
+            } else {
+                Slog.e(LOG_TAG, "Cannot cancel a non-current request");
             }
         }
     }
@@ -387,14 +398,13 @@
         }
 
         @Override
-        public boolean checkAttention(int requestCode, long timeout,
-                AttentionCallbackInternal callback) {
-            return AttentionManagerService.this.checkAttention(requestCode, timeout, callback);
+        public boolean checkAttention(long timeout, AttentionCallbackInternal callbackInternal) {
+            return AttentionManagerService.this.checkAttention(timeout, callbackInternal);
         }
 
         @Override
-        public void cancelAttentionCheck(int requestCode) {
-            AttentionManagerService.this.cancelAttentionCheck(requestCode);
+        public void cancelAttentionCheck(AttentionCallbackInternal callbackInternal) {
+            AttentionManagerService.this.cancelAttentionCheck(callbackInternal);
         }
 
         @Override
@@ -417,19 +427,17 @@
     }
 
     private static final class PendingAttentionCheck {
-        private final int mRequestCode;
-        private final AttentionCallbackInternal mCallback;
+        private final AttentionCallbackInternal mCallbackInternal;
         private final Runnable mRunnable;
 
-        PendingAttentionCheck(int requestCode, AttentionCallbackInternal callback,
+        PendingAttentionCheck(AttentionCallbackInternal callbackInternal,
                 Runnable runnable) {
-            mRequestCode = requestCode;
-            mCallback = callback;
+            mCallbackInternal = callbackInternal;
             mRunnable = runnable;
         }
 
         void cancel(@AttentionFailureCodes int failureCode) {
-            mCallback.onFailure(mRequestCode, failureCode);
+            mCallbackInternal.onFailure(failureCode);
         }
 
         void run() {
@@ -437,6 +445,17 @@
         }
     }
 
+    private static final class AttentionCheck {
+        private final AttentionCallbackInternal mCallbackInternal;
+        private final IAttentionCallback mIAttentionCallback;
+
+        AttentionCheck(AttentionCallbackInternal callbackInternal,
+                IAttentionCallback iAttentionCallback) {
+            mCallbackInternal = callbackInternal;
+            mIAttentionCallback = iAttentionCallback;
+        }
+    }
+
     private static final class UserState {
         final ComponentName mComponentName;
         final AttentionServiceConnection mConnection = new AttentionServiceConnection();
@@ -446,12 +465,12 @@
         @GuardedBy("mLock")
         boolean mBinding;
         @GuardedBy("mLock")
-        int mCurrentAttentionCheckRequestCode;
+        AttentionCheck mCurrentAttentionCheck;
         @GuardedBy("mLock")
         boolean mCurrentAttentionCheckIsFulfilled;
+
         @GuardedBy("mLock")
         PendingAttentionCheck mPendingAttentionCheck;
-
         @GuardedBy("mLock")
         AttentionCheckCache mAttentionCheckCache;
 
@@ -569,8 +588,7 @@
                         if (userState != null) {
                             // If not called back already.
                             if (!userState.mCurrentAttentionCheckIsFulfilled) {
-                                cancel(userState,
-                                        AttentionService.ATTENTION_FAILURE_TIMED_OUT);
+                                cancel(userState, AttentionService.ATTENTION_FAILURE_TIMED_OUT);
                             }
 
                         }
@@ -588,13 +606,14 @@
         if (userState.mService != null) {
             try {
                 userState.mService.cancelAttentionCheck(
-                        userState.mCurrentAttentionCheckRequestCode);
+                        userState.mCurrentAttentionCheck.mIAttentionCallback);
             } catch (RemoteException e) {
                 Slog.e(LOG_TAG, "Unable to cancel attention check");
             }
 
             if (userState.mPendingAttentionCheck != null) {
                 userState.mPendingAttentionCheck.cancel(failureCode);
+                userState.mPendingAttentionCheck = null;
             }
         }
     }
diff --git a/services/core/java/com/android/server/power/AttentionDetector.java b/services/core/java/com/android/server/power/AttentionDetector.java
index 701e5af..d9d21ba 100644
--- a/services/core/java/com/android/server/power/AttentionDetector.java
+++ b/services/core/java/com/android/server/power/AttentionDetector.java
@@ -34,6 +34,7 @@
 import com.android.server.LocalServices;
 
 import java.io.PrintWriter;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 
 /**
@@ -65,6 +66,11 @@
     private final Object mLock;
 
     /**
+     * If we're currently waiting for an attention callback
+     */
+    private final AtomicBoolean mRequested;
+
+    /**
      * {@link android.service.attention.AttentionService} API timeout.
      */
     private long mMaxAttentionApiTimeoutMillis;
@@ -78,11 +84,6 @@
     protected AttentionManagerInternal mAttentionManager;
 
     /**
-     * If we're currently waiting for an attention callback
-     */
-    private boolean mRequested;
-
-    /**
      * Current wakefulness of the device. {@see PowerManagerInternal}
      */
     private int mWakefulness;
@@ -94,14 +95,11 @@
 
     @VisibleForTesting
     final AttentionCallbackInternal mCallback = new AttentionCallbackInternal() {
-
         @Override
-        public void onSuccess(int requestCode, int result, long timestamp) {
-            Slog.v(TAG, "onSuccess: " + requestCode + ", " + result
-                    + " - current requestCode: " + getRequestCode());
-            synchronized (mLock) {
-                if (requestCode == getRequestCode() && mRequested) {
-                    mRequested = false;
+        public void onSuccess(int result, long timestamp) {
+            Slog.v(TAG, "onSuccess: " + result);
+            if (mRequested.getAndSet(false)) {
+                synchronized (mLock) {
                     if (mWakefulness != PowerManagerInternal.WAKEFULNESS_AWAKE) {
                         if (DEBUG) Slog.d(TAG, "Device slept before receiving callback.");
                         return;
@@ -116,19 +114,16 @@
         }
 
         @Override
-        public void onFailure(int requestCode, int error) {
+        public void onFailure(int error) {
             Slog.i(TAG, "Failed to check attention: " + error);
-            synchronized (mLock) {
-                if (requestCode == getRequestCode()) {
-                    mRequested = false;
-                }
-            }
+            mRequested.set(false);
         }
     };
 
     public AttentionDetector(Runnable onUserAttention, Object lock) {
         mOnUserAttention = onUserAttention;
         mLock = lock;
+        mRequested = new AtomicBoolean(false);
 
         // Device starts with an awake state upon boot.
         mWakefulness = PowerManagerInternal.WAKEFULNESS_AWAKE;
@@ -181,9 +176,11 @@
                         + (whenToCheck - whenToStopExtending));
             }
             return nextScreenDimming;
-        } else if (mRequested) {
+        } else if (mRequested.get()) {
             if (DEBUG) {
-                Slog.d(TAG, "Pending attention callback, wait. " + getRequestCode());
+                // TODO(b/128134941): consider adding a member ID increasing counter in
+                //  AttentionCallbackInternal to track this better.
+                Slog.d(TAG, "Pending attention callback, wait.");
             }
             return whenToCheck;
         }
@@ -192,14 +189,13 @@
         // callback might arrive before #checkAttention returns (if there are cached results.)
         // This means that we must assume that the request was successful, and then cancel it
         // afterwards if AttentionManager couldn't deliver it.
-        mRequested = true;
-        final boolean sent = mAttentionManager.checkAttention(getRequestCode(),
-                getAttentionTimeout(), mCallback);
+        mRequested.set(true);
+        final boolean sent = mAttentionManager.checkAttention(getAttentionTimeout(), mCallback);
         if (!sent) {
-            mRequested = false;
+            mRequested.set(false);
         }
 
-        Slog.v(TAG, "Checking user attention with request code: " + getRequestCode());
+        Slog.v(TAG, "Checking user attention");
         return whenToCheck;
     }
 
@@ -241,9 +237,9 @@
     }
 
     private void cancelCurrentRequestIfAny() {
-        if (mRequested) {
-            mAttentionManager.cancelAttentionCheck(getRequestCode());
-            mRequested = false;
+        if (mRequested.get()) {
+            mAttentionManager.cancelAttentionCheck(mCallback);
+            mRequested.set(false);
         }
     }
 
@@ -255,11 +251,6 @@
     }
 
     @VisibleForTesting
-    int getRequestCode() {
-        return (int) (mLastUserActivityTime % Integer.MAX_VALUE);
-    }
-
-    @VisibleForTesting
     long getAttentionTimeout() {
         return mMaxAttentionApiTimeoutMillis;
     }
diff --git a/services/tests/servicestests/src/com/android/server/power/AttentionDetectorTest.java b/services/tests/servicestests/src/com/android/server/power/AttentionDetectorTest.java
index a1a58b4..5de41ea 100644
--- a/services/tests/servicestests/src/com/android/server/power/AttentionDetectorTest.java
+++ b/services/tests/servicestests/src/com/android/server/power/AttentionDetectorTest.java
@@ -21,7 +21,6 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.reset;
@@ -29,6 +28,7 @@
 import static org.mockito.Mockito.when;
 
 import android.attention.AttentionManagerInternal;
+import android.attention.AttentionManagerInternal.AttentionCallbackInternal;
 import android.os.PowerManager;
 import android.os.PowerManagerInternal;
 import android.os.SystemClock;
@@ -41,14 +41,17 @@
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
 @SmallTest
 public class AttentionDetectorTest extends AndroidTestCase {
 
-    private @Mock AttentionManagerInternal mAttentionManagerInternal;
-    private @Mock Runnable mOnUserAttention;
+    @Mock
+    private AttentionManagerInternal mAttentionManagerInternal;
+    @Mock
+    private Runnable mOnUserAttention;
     private TestableAttentionDetector mAttentionDetector;
     private long mAttentionTimeout;
     private long mNextDimming;
@@ -57,7 +60,7 @@
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        when(mAttentionManagerInternal.checkAttention(anyInt(), anyLong(), any()))
+        when(mAttentionManagerInternal.checkAttention(anyLong(), any()))
                 .thenReturn(true);
         mAttentionDetector = new TestableAttentionDetector();
         mAttentionDetector.onWakefulnessChangeStarted(PowerManagerInternal.WAKEFULNESS_AWAKE);
@@ -82,7 +85,7 @@
     @Test
     public void testOnUserActivity_checksAttention() {
         long when = registerAttention();
-        verify(mAttentionManagerInternal).checkAttention(anyInt(), anyLong(), any());
+        verify(mAttentionManagerInternal).checkAttention(anyLong(), any());
         assertThat(when).isLessThan(mNextDimming);
     }
 
@@ -92,7 +95,7 @@
                 Settings.System.ADAPTIVE_SLEEP, 0, UserHandle.USER_CURRENT);
         mAttentionDetector.updateEnabledFromSettings(getContext());
         long when = registerAttention();
-        verify(mAttentionManagerInternal, never()).checkAttention(anyInt(), anyLong(), any());
+        verify(mAttentionManagerInternal, never()).checkAttention(anyLong(), any());
         assertThat(mNextDimming).isEqualTo(when);
     }
 
@@ -100,7 +103,7 @@
     public void testOnUserActivity_doesntCheckIfNotSupported() {
         mAttentionDetector.setAttentionServiceSupported(false);
         long when = registerAttention();
-        verify(mAttentionManagerInternal, never()).checkAttention(anyInt(), anyLong(), any());
+        verify(mAttentionManagerInternal, never()).checkAttention(anyLong(), any());
         assertThat(mNextDimming).isEqualTo(when);
     }
 
@@ -129,7 +132,7 @@
         mNextDimming = now;
         mAttentionDetector.onUserActivity(now, PowerManager.USER_ACTIVITY_EVENT_TOUCH);
         mAttentionDetector.updateUserActivity(mNextDimming + 5000L);
-        verify(mAttentionManagerInternal, never()).checkAttention(anyInt(), anyLong(), any());
+        verify(mAttentionManagerInternal, never()).checkAttention(anyLong(), any());
     }
 
     @Test
@@ -146,7 +149,7 @@
         long now = SystemClock.uptimeMillis();
         mAttentionDetector.onUserActivity(now - 15000L, PowerManager.USER_ACTIVITY_EVENT_TOUCH);
         mAttentionDetector.updateUserActivity(now + 2000L);
-        verify(mAttentionManagerInternal, never()).checkAttention(anyInt(), anyLong(), any());
+        verify(mAttentionManagerInternal, never()).checkAttention(anyLong(), any());
     }
 
     @Test
@@ -154,7 +157,7 @@
         registerAttention();
         reset(mAttentionManagerInternal);
         long when = mAttentionDetector.updateUserActivity(mNextDimming);
-        verify(mAttentionManagerInternal, never()).checkAttention(anyInt(), anyLong(), any());
+        verify(mAttentionManagerInternal, never()).checkAttention(anyLong(), any());
         assertThat(when).isLessThan(mNextDimming);
     }
 
@@ -162,32 +165,35 @@
     public void testOnWakefulnessChangeStarted_cancelsRequestWhenNotAwake() {
         registerAttention();
         mAttentionDetector.onWakefulnessChangeStarted(PowerManagerInternal.WAKEFULNESS_ASLEEP);
-        verify(mAttentionManagerInternal).cancelAttentionCheck(anyInt());
+
+        ArgumentCaptor<AttentionCallbackInternal> callbackCaptor = ArgumentCaptor.forClass(
+                AttentionCallbackInternal.class);
+        verify(mAttentionManagerInternal).cancelAttentionCheck(callbackCaptor.capture());
+        assertEquals(callbackCaptor.getValue(), mAttentionDetector.mCallback);
     }
 
     @Test
     public void testCallbackOnSuccess_ignoresIfNoAttention() {
         registerAttention();
-        mAttentionDetector.mCallback.onSuccess(mAttentionDetector.getRequestCode(),
-                AttentionService.ATTENTION_SUCCESS_ABSENT, SystemClock.uptimeMillis());
+        mAttentionDetector.mCallback.onSuccess(AttentionService.ATTENTION_SUCCESS_ABSENT,
+                SystemClock.uptimeMillis());
         verify(mOnUserAttention, never()).run();
     }
 
     @Test
     public void testCallbackOnSuccess_callsCallback() {
         registerAttention();
-        mAttentionDetector.mCallback.onSuccess(mAttentionDetector.getRequestCode(),
-                AttentionService.ATTENTION_SUCCESS_PRESENT, SystemClock.uptimeMillis());
+        mAttentionDetector.mCallback.onSuccess(AttentionService.ATTENTION_SUCCESS_PRESENT,
+                SystemClock.uptimeMillis());
         verify(mOnUserAttention).run();
     }
 
     @Test
     public void testCallbackOnFailure_unregistersCurrentRequestCode() {
         registerAttention();
-        mAttentionDetector.mCallback.onFailure(mAttentionDetector.getRequestCode(),
-                AttentionService.ATTENTION_FAILURE_UNKNOWN);
-        mAttentionDetector.mCallback.onSuccess(mAttentionDetector.getRequestCode(),
-                AttentionService.ATTENTION_SUCCESS_PRESENT, SystemClock.uptimeMillis());
+        mAttentionDetector.mCallback.onFailure(AttentionService.ATTENTION_FAILURE_UNKNOWN);
+        mAttentionDetector.mCallback.onSuccess(AttentionService.ATTENTION_SUCCESS_PRESENT,
+                SystemClock.uptimeMillis());
         verify(mOnUserAttention, never()).run();
     }