diff --git a/core/java/android/service/wallpaper/WallpaperService.java b/core/java/android/service/wallpaper/WallpaperService.java
index f548d3b77bf3d..1c6275fb8dc13 100644
--- a/core/java/android/service/wallpaper/WallpaperService.java
+++ b/core/java/android/service/wallpaper/WallpaperService.java
@@ -195,7 +195,8 @@ public abstract class WallpaperService extends Service {
         // Needed for throttling onComputeColors.
         private long mLastColorInvalidation;
         private final Runnable mNotifyColorsChanged = this::notifyColorsChanged;
-        private Supplier<Long> mClockFunction = SystemClock::elapsedRealtime;
+        private final Supplier<Long> mClockFunction;
+        private final Handler mHandler;
 
         DisplayManager mDisplayManager;
         Display mDisplay;
@@ -362,6 +363,26 @@ public abstract class WallpaperService extends Service {
                 }
             }
         };
+
+        /**
+         * Default constructor
+         */
+        public Engine() {
+            this(SystemClock::elapsedRealtime, Handler.getMain());
+        }
+
+        /**
+         * Constructor used for test purposes.
+         *
+         * @param clockFunction Supplies current times in millis.
+         * @param handler Used for posting/deferring asynchronous calls.
+         * @hide
+         */
+        @VisibleForTesting
+        public Engine(Supplier<Long> clockFunction, Handler handler) {
+           mClockFunction = clockFunction;
+           mHandler = handler;
+        }
         
         /**
          * Provides access to the surface in which this wallpaper is drawn.
@@ -563,18 +584,17 @@ public abstract class WallpaperService extends Service {
          */
         public void notifyColorsChanged() {
             final long now = mClockFunction.get();
-            final Handler mainHandler = Handler.getMain();
             if (now - mLastColorInvalidation < NOTIFY_COLORS_RATE_LIMIT_MS) {
                 Log.w(TAG, "This call has been deferred. You should only call "
                         + "notifyColorsChanged() once every "
                         + (NOTIFY_COLORS_RATE_LIMIT_MS / 1000f) + " seconds.");
-                if (!mainHandler.hasCallbacks(mNotifyColorsChanged)) {
-                    mainHandler.postDelayed(mNotifyColorsChanged, NOTIFY_COLORS_RATE_LIMIT_MS);
+                if (!mHandler.hasCallbacks(mNotifyColorsChanged)) {
+                    mHandler.postDelayed(mNotifyColorsChanged, NOTIFY_COLORS_RATE_LIMIT_MS);
                 }
                 return;
             }
             mLastColorInvalidation = now;
-            mainHandler.removeCallbacks(mNotifyColorsChanged);
+            mHandler.removeCallbacks(mNotifyColorsChanged);
 
             try {
                 final WallpaperColors newColors = onComputeColors();
@@ -662,14 +682,6 @@ public abstract class WallpaperService extends Service {
             }
         }
 
-        /**
-         * @hide
-         */
-        @VisibleForTesting
-        public void setClockFunction(Supplier<Long> clockFunction) {
-            mClockFunction = clockFunction;
-        }
-
         void updateSurface(boolean forceRelayout, boolean forceReport, boolean redrawNeeded) {
             if (mDestroyed) {
                 Log.w(TAG, "Ignoring updateSurface: destroyed");
diff --git a/services/tests/servicestests/src/com/android/server/wallpaper/WallpaperServiceTests.java b/services/tests/servicestests/src/com/android/server/wallpaper/WallpaperServiceTests.java
index 0d29e89637286..9c010a07135d4 100644
--- a/services/tests/servicestests/src/com/android/server/wallpaper/WallpaperServiceTests.java
+++ b/services/tests/servicestests/src/com/android/server/wallpaper/WallpaperServiceTests.java
@@ -17,10 +17,14 @@
 package com.android.server.wallpaper;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
 import android.app.WallpaperColors;
+import android.os.Handler;
+import android.os.Message;
 import android.os.SystemClock;
 import android.service.wallpaper.WallpaperService;
+import android.support.test.annotation.UiThreadTest;
 import android.support.test.filters.SmallTest;
 import android.support.test.runner.AndroidJUnit4;
 
@@ -28,18 +32,31 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 
 import java.util.concurrent.CountDownLatch;
+import java.util.function.Supplier;
 
 @SmallTest
 @RunWith(AndroidJUnit4.class)
 public class WallpaperServiceTests {
 
+    @UiThreadTest
     @Test
     public void testNotifyColorsChanged_rateLimit() throws Exception {
+        long[] clockOffset = {0};
+        boolean[] postDelayed = {false};
+        Supplier<Long> clockFunction = () -> SystemClock.elapsedRealtime() + clockOffset[0];
+        Handler handler = new Handler() {
+            @Override
+            public boolean sendMessageAtTime(Message msg, long uptimeMillis) {
+                postDelayed[0] = true;
+                return super.sendMessageAtTime(msg, uptimeMillis);
+            }
+        };
+
         CountDownLatch eventCountdown = new CountDownLatch(2);
         WallpaperService service = new WallpaperService() {
             @Override
             public Engine onCreateEngine() {
-                return new WallpaperService.Engine() {
+                return new WallpaperService.Engine(clockFunction, handler) {
                     @Override
                     public WallpaperColors onComputeColors() {
                         eventCountdown.countDown();
@@ -59,8 +76,11 @@ public class WallpaperServiceTests {
         engine.notifyColorsChanged();
         assertEquals("OnComputeColors should have been throttled.",
                 1, eventCountdown.getCount());
-        // Called after being deferred.
-        engine.setClockFunction(() ->  SystemClock.elapsedRealtime() + 1500);
+        // Should have been posted to the handler.
+        assertTrue("Event should have been delayed", postDelayed[0]);
+
+        // Called again after being deferred.
+        clockOffset[0] = 1500;
         engine.notifyColorsChanged();
         assertEquals("OnComputeColors should have been deferred.",
                 0, eventCountdown.getCount());