2 files changed, 42 insertions, 5 deletions

diff --git a/core/java/com/android/internal/util/WakeupMessage.java b/core/java/com/android/internal/util/WakeupMessage.java
index 46098c58108f..70b6f96e3717 100644
--- a/core/java/com/android/internal/util/WakeupMessage.java
+++ b/core/java/com/android/internal/util/WakeupMessage.java
@@ -47,17 +47,19 @@ public class WakeupMessage implements AlarmManager.OnAlarmListener {
     protected final int mCmd, mArg1, mArg2;
     @VisibleForTesting
     protected final Object mObj;
+    private final Runnable mRunnable;
     private boolean mScheduled;
 
     public WakeupMessage(Context context, Handler handler,
             String cmdName, int cmd, int arg1, int arg2, Object obj) {
-        mAlarmManager = (AlarmManager) context.getSystemService(Context.ALARM_SERVICE);
+        mAlarmManager = getAlarmManager(context);
         mHandler = handler;
         mCmdName = cmdName;
         mCmd = cmd;
         mArg1 = arg1;
         mArg2 = arg2;
         mObj = obj;
+        mRunnable = null;
     }
 
     public WakeupMessage(Context context, Handler handler, String cmdName, int cmd, int arg1) {
@@ -73,6 +75,21 @@ public class WakeupMessage implements AlarmManager.OnAlarmListener {
         this(context, handler, cmdName, cmd, 0, 0, null);
     }
 
+    public WakeupMessage(Context context, Handler handler, String cmdName, Runnable runnable) {
+        mAlarmManager = getAlarmManager(context);
+        mHandler = handler;
+        mCmdName = cmdName;
+        mCmd = 0;
+        mArg1 = 0;
+        mArg2 = 0;
+        mObj = null;
+        mRunnable = runnable;
+    }
+
+    private static AlarmManager getAlarmManager(Context context) {
+        return (AlarmManager) context.getSystemService(Context.ALARM_SERVICE);
+    }
+
     /**
      * Schedule the message to be delivered at the time in milliseconds of the
      * {@link android.os.SystemClock#elapsedRealtime SystemClock.elapsedRealtime()} clock and wakeup
@@ -107,7 +124,12 @@ public class WakeupMessage implements AlarmManager.OnAlarmListener {
             mScheduled = false;
         }
         if (stillScheduled) {
-            Message msg = mHandler.obtainMessage(mCmd, mArg1, mArg2, mObj);
+            Message msg;
+            if (mRunnable == null) {
+                msg = mHandler.obtainMessage(mCmd, mArg1, mArg2, mObj);
+            } else {
+                msg = Message.obtain(mHandler, mRunnable);
+            }
             mHandler.dispatchMessage(msg);
             msg.recycle();
         }
diff --git a/core/tests/utiltests/src/com/android/internal/util/WakeupMessageTest.java b/core/tests/utiltests/src/com/android/internal/util/WakeupMessageTest.java
index 7935880d12a9..734ebef0007e 100644
--- a/core/tests/utiltests/src/com/android/internal/util/WakeupMessageTest.java
+++ b/core/tests/utiltests/src/com/android/internal/util/WakeupMessageTest.java
@@ -47,6 +47,7 @@ public class WakeupMessageTest {
     private static final int TEST_ARG2 = 182;
     private static final Object TEST_OBJ = "hello";
 
+    @Mock Context mContext;
     @Mock AlarmManager mAlarmManager;
     WakeupMessage mMessage;
     // Make a spy so that we can verify calls to it
@@ -86,13 +87,12 @@ public class WakeupMessageTest {
     public void setUp() {
         MockitoAnnotations.initMocks(this);
 
-        Context context = mock(Context.class);
-        when(context.getSystemService(Context.ALARM_SERVICE)).thenReturn(mAlarmManager);
+        when(mContext.getSystemService(Context.ALARM_SERVICE)).thenReturn(mAlarmManager);
         // capture the listener for each AlarmManager.setExact call
         doNothing().when(mAlarmManager).setExact(anyInt(), anyLong(), any(String.class),
                 mListenerCaptor.capture(), any(Handler.class));
 
-        mMessage = new WakeupMessage(context, mHandler, TEST_CMD_NAME, TEST_CMD, TEST_ARG1,
+        mMessage = new WakeupMessage(mContext, mHandler, TEST_CMD_NAME, TEST_CMD, TEST_ARG1,
                 TEST_ARG2, TEST_OBJ);
     }
 
@@ -168,4 +168,19 @@ public class WakeupMessageTest {
         verifyMessageDispatchedOnce();
     }
 
+    /**
+     * Verify that a Runnable is scheduled and dispatched.
+     */
+    @Test
+    public void scheduleRunnable() {
+        final long when = 1011;
+        final Runnable runnable = mock(Runnable.class);
+        WakeupMessage dut = new WakeupMessage(mContext, mHandler, TEST_CMD_NAME, runnable);
+        dut.schedule(when);
+        verify(mAlarmManager).setExact(eq(AlarmManager.ELAPSED_REALTIME_WAKEUP), eq(when),
+                eq(TEST_CMD_NAME), any(AlarmManager.OnAlarmListener.class), eq(mHandler));
+        mListenerCaptor.getValue().onAlarm();
+        verify(runnable, times(1)).run();
+    }
+
 }