diff --git a/packages/DocumentsUI/tests/src/com/android/documentsui/ThumbnailCacheTest.java b/packages/DocumentsUI/tests/src/com/android/documentsui/ThumbnailCacheTest.java
index dda4918475914..a6951cfc216be 100644
--- a/packages/DocumentsUI/tests/src/com/android/documentsui/ThumbnailCacheTest.java
+++ b/packages/DocumentsUI/tests/src/com/android/documentsui/ThumbnailCacheTest.java
@@ -25,13 +25,16 @@ import android.content.ComponentCallbacks2;
 import android.graphics.Bitmap;
 import android.graphics.Point;
 import android.net.Uri;
+import android.support.test.annotation.UiThreadTest;
 import android.support.test.filters.SmallTest;
+import android.support.test.rule.UiThreadTestRule;
 import android.support.test.runner.AndroidJUnit4;
 
 import com.android.documentsui.ThumbnailCache.Result;
 import com.android.documentsui.testing.Bitmaps;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -52,6 +55,11 @@ public class ThumbnailCacheTest {
 
     private static final long LAST_MODIFIED = 100;
 
+    // We need this rule to help us run tests in main thread. It'll be applied to each test
+    // annotated with @UiThreadTest.
+    @Rule
+    public UiThreadTestRule uiThreadRule = new UiThreadTestRule();
+
     private static final int CACHE_SIZE_LIMIT =
             MIDSIZE_BITMAP.getByteCount() + LARGE_BITMAP.getByteCount();
 
@@ -63,6 +71,7 @@ public class ThumbnailCacheTest {
     }
 
     @Test
+    @UiThreadTest
     public void testMiss() {
         mCache.putThumbnail(URI_0, MID_SIZE, MIDSIZE_BITMAP, LAST_MODIFIED);
 
@@ -72,6 +81,7 @@ public class ThumbnailCacheTest {
     }
 
     @Test
+    @UiThreadTest
     public void testHit_Exact() {
         mCache.putThumbnail(URI_0, MID_SIZE, MIDSIZE_BITMAP, LAST_MODIFIED);
 
@@ -82,6 +92,7 @@ public class ThumbnailCacheTest {
     }
 
     @Test
+    @UiThreadTest
     public void testHit_Smaller() {
         mCache.putThumbnail(URI_0, MID_SIZE, MIDSIZE_BITMAP, LAST_MODIFIED);
 
@@ -92,6 +103,7 @@ public class ThumbnailCacheTest {
     }
 
     @Test
+    @UiThreadTest
     public void testHit_Larger() {
         mCache.putThumbnail(URI_0, MID_SIZE, MIDSIZE_BITMAP, LAST_MODIFIED);
 
@@ -102,6 +114,7 @@ public class ThumbnailCacheTest {
     }
 
     @Test
+    @UiThreadTest
     public void testHit_Larger_HasBothSize() {
         mCache.putThumbnail(URI_0, LARGE_SIZE, LARGE_BITMAP, LAST_MODIFIED);
         mCache.putThumbnail(URI_0, SMALL_SIZE, SMALL_BITMAP, LAST_MODIFIED);
@@ -113,6 +126,7 @@ public class ThumbnailCacheTest {
     }
 
     @Test
+    @UiThreadTest
     public void testHit_Exact_MultiplePut() {
         mCache.putThumbnail(URI_0, MID_SIZE, MIDSIZE_BITMAP, LAST_MODIFIED);
 
@@ -127,6 +141,7 @@ public class ThumbnailCacheTest {
     }
 
     @Test
+    @UiThreadTest
     public void testHit_EqualLastModified() {
         mCache.putThumbnail(URI_0, MID_SIZE, MIDSIZE_BITMAP, LAST_MODIFIED);
 
@@ -136,6 +151,7 @@ public class ThumbnailCacheTest {
     }
 
     @Test
+    @UiThreadTest
     public void testEvictOldest_SizeExceeded() {
         mCache.putThumbnail(URI_0, MID_SIZE, MIDSIZE_BITMAP, LAST_MODIFIED);
         mCache.putThumbnail(URI_1, SMALL_SIZE, SMALL_BITMAP, LAST_MODIFIED);
@@ -147,6 +163,7 @@ public class ThumbnailCacheTest {
     }
 
     @Test
+    @UiThreadTest
     public void testCacheShrink_OnTrimMemory_Moderate() {
         mCache.putThumbnail(URI_0, MID_SIZE, MIDSIZE_BITMAP, LAST_MODIFIED);
         mCache.putThumbnail(URI_0, SMALL_SIZE, SMALL_BITMAP, LAST_MODIFIED);
@@ -159,6 +176,7 @@ public class ThumbnailCacheTest {
     }
 
     @Test
+    @UiThreadTest
     public void testCacheShrink_OnTrimMemory_Background() {
         mCache.putThumbnail(URI_0, LARGE_SIZE, LARGE_BITMAP, LAST_MODIFIED);
         mCache.putThumbnail(URI_0, SMALL_SIZE, SMALL_BITMAP, LAST_MODIFIED);