Revert "Revert "Integrate new looper apis into testables""

This reverts commit fd8f615953631863c94d7080441f8dd4a0a74f56.
+ some minor adjustments.

Test: runtest systemui
Change-Id: Ie4349fcfb2aadbbd8826506faa8a6fb1ecd5b589

5 files changed, 207 insertions, 172 deletions

diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/BluetoothControllerImplTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/BluetoothControllerImplTest.java
index 8808988406a9..f516d74f4062 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/BluetoothControllerImplTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/BluetoothControllerImplTest.java
@@ -20,7 +20,9 @@ import static org.mockito.Mockito.when;
 
 import android.bluetooth.BluetoothAdapter;
 import android.bluetooth.BluetoothProfile;
+import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
+import android.testing.TestableLooper.RunWithLooper;
 
 import com.android.settingslib.bluetooth.BluetoothEventManager;
 import com.android.settingslib.bluetooth.CachedBluetoothDevice;
@@ -31,10 +33,13 @@ import com.android.systemui.SysuiTestCase;
 
 import org.junit.Before;
 import org.junit.Test;
+import org.junit.runner.RunWith;
 
 import java.util.ArrayList;
 import java.util.List;
 
+@RunWith(AndroidTestingRunner.class)
+@RunWithLooper
 public class BluetoothControllerImplTest extends SysuiTestCase {
 
     private LocalBluetoothManager mMockBluetoothManager;
@@ -47,7 +52,7 @@ public class BluetoothControllerImplTest extends SysuiTestCase {
 
     @Before
     public void setup() throws Exception {
-        mTestableLooper = new TestableLooper();
+        mTestableLooper = TestableLooper.get(this);
         mMockBluetoothManager = mDependency.injectMockDependency(LocalBluetoothManager.class);
         mDevices = new ArrayList<>();
         mMockDeviceManager = mock(CachedBluetoothDeviceManager.class);
diff --git a/services/tests/notification/src/com/android/server/notification/NotificationManagerServiceTest.java b/services/tests/notification/src/com/android/server/notification/NotificationManagerServiceTest.java
index 92534a1480ac..d057eb5f1646 100644
--- a/services/tests/notification/src/com/android/server/notification/NotificationManagerServiceTest.java
+++ b/services/tests/notification/src/com/android/server/notification/NotificationManagerServiceTest.java
@@ -53,9 +53,10 @@ import android.os.Process;
 import android.os.UserHandle;
 import android.service.notification.NotificationListenerService;
 import android.service.notification.StatusBarNotification;
-import android.support.test.annotation.UiThreadTest;
 import android.support.test.InstrumentationRegistry;
+import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
+import android.testing.TestableLooper.RunWithLooper;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -63,6 +64,7 @@ import java.util.List;
 
 import org.junit.Before;
 import org.junit.Test;
+import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
@@ -70,6 +72,8 @@ import org.mockito.MockitoAnnotations;
 import com.android.server.lights.Light;
 import com.android.server.lights.LightsManager;
 
+@RunWith(AndroidTestingRunner.class)
+@RunWithLooper
 public class NotificationManagerServiceTest {
     private static final long WAIT_FOR_IDLE_TIMEOUT = 2;
     private static final String TEST_CHANNEL_ID = "NotificationManagerServiceTestChannelId";
@@ -109,7 +113,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Before
-    @UiThreadTest
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
         mNotificationManagerService = new TestableNotificationManagerService(mContext);
@@ -124,7 +127,7 @@ public class NotificationManagerServiceTest {
         final LightsManager mockLightsManager = mock(LightsManager.class);
         when(mockLightsManager.getLight(anyInt())).thenReturn(mock(Light.class));
         // Use this testable looper.
-        mTestableLooper = new TestableLooper(false);
+        mTestableLooper = TestableLooper.get(this);
 
         mListener = mNotificationListeners.new ManagedServiceInfo(
                 null, new ComponentName(PKG, "test_class"), uid, true, null, 0);
@@ -165,7 +168,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testCreateNotificationChannels_SingleChannel() throws Exception {
         final NotificationChannel channel =
                 new NotificationChannel("id", "name", NotificationManager.IMPORTANCE_DEFAULT);
@@ -177,7 +179,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testCreateNotificationChannels_NullChannelThrowsException() throws Exception {
         try {
             mBinderService.createNotificationChannels("test_pkg",
@@ -189,7 +190,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testCreateNotificationChannels_TwoChannels() throws Exception {
         final NotificationChannel channel1 =
                 new NotificationChannel("id1", "name", NotificationManager.IMPORTANCE_DEFAULT);
@@ -202,7 +202,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testCreateNotificationChannels_SecondCreateDoesNotChangeImportance()
             throws Exception {
         final NotificationChannel channel =
@@ -221,7 +220,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testCreateNotificationChannels_IdenticalChannelsInListIgnoresSecond()
             throws Exception {
         final NotificationChannel channel1 =
@@ -236,7 +234,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testBlockedNotifications_suspended() throws Exception {
         NotificationUsageStats usageStats = mock(NotificationUsageStats.class);
         when(mPackageManager.isPackageSuspendedForUser(anyString(), anyInt())).thenReturn(true);
@@ -249,7 +246,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testBlockedNotifications_blockedChannel() throws Exception {
         NotificationUsageStats usageStats = mock(NotificationUsageStats.class);
         when(mPackageManager.isPackageSuspendedForUser(anyString(), anyInt())).thenReturn(false);
@@ -263,7 +259,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testBlockedNotifications_blockedApp() throws Exception {
         NotificationUsageStats usageStats = mock(NotificationUsageStats.class);
         when(mPackageManager.isPackageSuspendedForUser(anyString(), anyInt())).thenReturn(false);
@@ -277,7 +272,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testEnqueueNotificationWithTag_PopulatesGetActiveNotifications() throws Exception {
         mBinderService.enqueueNotificationWithTag(PKG, "opPkg", "tag", 0,
                 generateNotificationRecord(null).getNotification(), 0);
@@ -288,7 +282,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testCancelNotificationImmediatelyAfterEnqueue() throws Exception {
         mBinderService.enqueueNotificationWithTag(PKG, "opPkg", "tag", 0,
                 generateNotificationRecord(null).getNotification(), 0);
@@ -300,7 +293,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testCancelNotificationWhilePostedAndEnqueued() throws Exception {
         mBinderService.enqueueNotificationWithTag(PKG, "opPkg", "tag", 0,
                 generateNotificationRecord(null).getNotification(), 0);
@@ -315,7 +307,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testCancelNotificationsFromListenerImmediatelyAfterEnqueue() throws Exception {
         final StatusBarNotification sbn = generateNotificationRecord(null).sbn;
         mBinderService.enqueueNotificationWithTag(PKG, "opPkg", "tag",
@@ -328,7 +319,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testCancelAllNotificationsImmediatelyAfterEnqueue() throws Exception {
         final StatusBarNotification sbn = generateNotificationRecord(null).sbn;
         mBinderService.enqueueNotificationWithTag(PKG, "opPkg", "tag",
@@ -341,7 +331,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testCancelAllNotifications_IgnoreForegroundService() throws Exception {
         final StatusBarNotification sbn = generateNotificationRecord(null).sbn;
         sbn.getNotification().flags |= Notification.FLAG_FOREGROUND_SERVICE;
@@ -355,7 +344,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testCancelAllNotifications_IgnoreOtherPackages() throws Exception {
         final StatusBarNotification sbn = generateNotificationRecord(null).sbn;
         sbn.getNotification().flags |= Notification.FLAG_FOREGROUND_SERVICE;
@@ -369,7 +357,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testCancelAllNotifications_NullPkgRemovesAll() throws Exception {
         final StatusBarNotification sbn = generateNotificationRecord(null).sbn;
         mBinderService.enqueueNotificationWithTag(PKG, "opPkg", "tag",
@@ -382,7 +369,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testCancelAllNotifications_NullPkgIgnoresUserAllNotifications() throws Exception {
         final StatusBarNotification sbn = generateNotificationRecord(null).sbn;
         mBinderService.enqueueNotificationWithTag(PKG, "opPkg", "tag",
@@ -396,7 +382,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testRemoveForegroundServiceFlag_ImmediatelyAfterEnqueue() throws Exception {
         final StatusBarNotification sbn = generateNotificationRecord(null).sbn;
         sbn.getNotification().flags |= Notification.FLAG_FOREGROUND_SERVICE;
@@ -411,7 +396,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testTvExtenderChannelOverride_onTv() throws Exception {
         mNotificationManagerService.setIsTelevision(true);
         mNotificationManagerService.setRankingHelper(mRankingHelper);
@@ -427,7 +411,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testTvExtenderChannelOverride_notOnTv() throws Exception {
         mNotificationManagerService.setIsTelevision(false);
         mNotificationManagerService.setRankingHelper(mRankingHelper);
@@ -443,7 +426,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testCreateChannelNotifyListener() throws Exception {
         List<String> associations = new ArrayList<>();
         associations.add("a");
@@ -469,7 +451,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testCreateChannelGroupNotifyListener() throws Exception {
         List<String> associations = new ArrayList<>();
         associations.add("a");
@@ -490,7 +471,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testUpdateChannelNotifyListener() throws Exception {
         List<String> associations = new ArrayList<>();
         associations.add("a");
@@ -509,7 +489,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testDeleteChannelNotifyListener() throws Exception {
         List<String> associations = new ArrayList<>();
         associations.add("a");
@@ -526,7 +505,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testDeleteChannelGroupNotifyListener() throws Exception {
         List<String> associations = new ArrayList<>();
         associations.add("a");
@@ -543,7 +521,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testUpdateNotificationChannelFromPrivilegedListener_success() throws Exception {
         mNotificationManagerService.setRankingHelper(mRankingHelper);
         List<String> associations = new ArrayList<>();
@@ -561,7 +538,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testUpdateNotificationChannelFromPrivilegedListener_noAccess() throws Exception {
         mNotificationManagerService.setRankingHelper(mRankingHelper);
         List<String> associations = new ArrayList<>();
@@ -583,7 +559,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testUpdateNotificationChannelFromPrivilegedListener_badUser() throws Exception {
         mNotificationManagerService.setRankingHelper(mRankingHelper);
         List<String> associations = new ArrayList<>();
@@ -609,7 +584,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testGetNotificationChannelFromPrivilegedListener_success() throws Exception {
         mNotificationManagerService.setRankingHelper(mRankingHelper);
         List<String> associations = new ArrayList<>();
@@ -624,7 +598,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testGetNotificationChannelFromPrivilegedListener_noAccess() throws Exception {
         mNotificationManagerService.setRankingHelper(mRankingHelper);
         List<String> associations = new ArrayList<>();
@@ -643,7 +616,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testGetNotificationChannelFromPrivilegedListener_badUser() throws Exception {
         mNotificationManagerService.setRankingHelper(mRankingHelper);
         List<String> associations = new ArrayList<>();
@@ -666,7 +638,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testGetNotificationChannelGroupsFromPrivilegedListener_success() throws Exception {
         mNotificationManagerService.setRankingHelper(mRankingHelper);
         List<String> associations = new ArrayList<>();
@@ -680,7 +651,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testGetNotificationChannelGroupsFromPrivilegedListener_noAccess() throws Exception {
         mNotificationManagerService.setRankingHelper(mRankingHelper);
         List<String> associations = new ArrayList<>();
@@ -698,7 +668,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testGetNotificationChannelGroupsFromPrivilegedListener_badUser() throws Exception {
         mNotificationManagerService.setRankingHelper(mRankingHelper);
         List<String> associations = new ArrayList<>();
@@ -719,7 +688,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testHasCompanionDevice_failure() throws Exception {
         when(mCompanionMgr.getAssociations(anyString(), anyInt())).thenThrow(
                 new IllegalArgumentException());
@@ -727,7 +695,6 @@ public class NotificationManagerServiceTest {
     }
 
     @Test
-    @UiThreadTest
     public void testHasCompanionDevice_noService() throws Exception {
         mNotificationManagerService = new TestableNotificationManagerService(mContext);
 
diff --git a/tests/testables/src/android/testing/AndroidTestingRunner.java b/tests/testables/src/android/testing/AndroidTestingRunner.java
index 816ed033a3e2..a425f70e836c 100644
--- a/tests/testables/src/android/testing/AndroidTestingRunner.java
+++ b/tests/testables/src/android/testing/AndroidTestingRunner.java
@@ -18,7 +18,7 @@ import android.support.test.internal.runner.junit4.statement.RunAfters;
 import android.support.test.internal.runner.junit4.statement.RunBefores;
 import android.support.test.internal.runner.junit4.statement.UiThreadStatement;
 
-import android.testing.TestableLooper.LooperStatement;
+import android.testing.TestableLooper.LooperFrameworkMethod;
 import android.testing.TestableLooper.RunWithLooper;
 
 import org.junit.After;
@@ -30,6 +30,7 @@ import org.junit.runners.model.FrameworkMethod;
 import org.junit.runners.model.InitializationError;
 import org.junit.runners.model.Statement;
 
+import java.util.ArrayList;
 import java.util.List;
 
 /**
@@ -49,28 +50,21 @@ public class AndroidTestingRunner extends BlockJUnit4ClassRunner {
 
     @Override
     protected Statement methodInvoker(FrameworkMethod method, Object test) {
-        return shouldRunOnUiThread(method) ? new UiThreadStatement(
-                methodInvokerInt(method, test), true) : methodInvokerInt(method, test);
-    }
-
-    protected Statement methodInvokerInt(FrameworkMethod method, Object test) {
-        RunWithLooper annotation = method.getAnnotation(RunWithLooper.class);
-        if (annotation == null) annotation = mKlass.getAnnotation(RunWithLooper.class);
-        if (annotation != null) {
-            return new LooperStatement(super.methodInvoker(method, test),
-                    annotation.setAsMainLooper(), test);
-        }
-        return super.methodInvoker(method, test);
+        method = looperWrap(method, test, method);
+        final Statement statement = super.methodInvoker(method, test);
+        return shouldRunOnUiThread(method) ? new UiThreadStatement(statement, true) : statement;
     }
 
     protected Statement withBefores(FrameworkMethod method, Object target, Statement statement) {
-        List befores = this.getTestClass().getAnnotatedMethods(Before.class);
+        List befores = looperWrap(method, target,
+                this.getTestClass().getAnnotatedMethods(Before.class));
         return befores.isEmpty() ? statement : new RunBefores(method, statement,
                 befores, target);
     }
 
     protected Statement withAfters(FrameworkMethod method, Object target, Statement statement) {
-        List afters = this.getTestClass().getAnnotatedMethods(After.class);
+        List afters = looperWrap(method, target,
+                this.getTestClass().getAnnotatedMethods(After.class));
         return afters.isEmpty() ? statement : new RunAfters(method, statement, afters,
                 target);
     }
@@ -88,6 +82,30 @@ public class AndroidTestingRunner extends BlockJUnit4ClassRunner {
         return annotation == null ? 0L : annotation.timeout();
     }
 
+    protected List<FrameworkMethod> looperWrap(FrameworkMethod method, Object test,
+            List<FrameworkMethod> methods) {
+        RunWithLooper annotation = method.getAnnotation(RunWithLooper.class);
+        if (annotation == null) annotation = mKlass.getAnnotation(RunWithLooper.class);
+        if (annotation != null) {
+            methods = new ArrayList<>(methods);
+            for (int i = 0; i < methods.size(); i++) {
+                methods.set(i, LooperFrameworkMethod.get(methods.get(i),
+                        annotation.setAsMainLooper(), test));
+            }
+        }
+        return methods;
+    }
+
+    protected FrameworkMethod looperWrap(FrameworkMethod method, Object test,
+            FrameworkMethod base) {
+        RunWithLooper annotation = method.getAnnotation(RunWithLooper.class);
+        if (annotation == null) annotation = mKlass.getAnnotation(RunWithLooper.class);
+        if (annotation != null) {
+            return LooperFrameworkMethod.get(base, annotation.setAsMainLooper(), test);
+        }
+        return base;
+    }
+
     public boolean shouldRunOnUiThread(FrameworkMethod method) {
         if (mKlass.getAnnotation(UiThreadTest.class) != null) {
             return true;
diff --git a/tests/testables/src/android/testing/TestableLooper.java b/tests/testables/src/android/testing/TestableLooper.java
index 8a33cf918646..9eddc5112d0c 100644
--- a/tests/testables/src/android/testing/TestableLooper.java
+++ b/tests/testables/src/android/testing/TestableLooper.java
@@ -15,20 +15,21 @@
 package android.testing;
 
 import android.os.Handler;
+import android.os.HandlerThread;
 import android.os.Looper;
 import android.os.Message;
 import android.os.MessageQueue;
+import android.os.TestLooperManager;
+import android.support.test.InstrumentationRegistry;
 import android.util.ArrayMap;
 
-import org.junit.runners.model.Statement;
+import org.junit.runners.model.FrameworkMethod;
 
 import java.lang.annotation.ElementType;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
-import java.lang.reflect.Constructor;
 import java.lang.reflect.Field;
-import java.lang.reflect.Method;
 import java.util.Map;
 
 /**
@@ -38,65 +39,35 @@ import java.util.Map;
  */
 public class TestableLooper {
 
-    private final Method mNext;
-    private final Method mRecycleUnchecked;
-
     private Looper mLooper;
     private MessageQueue mQueue;
     private boolean mMain;
     private Object mOriginalMain;
     private MessageHandler mMessageHandler;
 
-    private int mParsedCount;
     private Handler mHandler;
-    private Message mEmptyMessage;
-
-    public TestableLooper() throws Exception {
-        this(true);
-    }
+    private Runnable mEmptyMessage;
+    private TestLooperManager mQueueWrapper;
 
-    public TestableLooper(boolean setMyLooper) throws Exception {
-        setupQueue(setMyLooper);
-        mNext = mQueue.getClass().getDeclaredMethod("next");
-        mNext.setAccessible(true);
-        mRecycleUnchecked = Message.class.getDeclaredMethod("recycleUnchecked");
-        mRecycleUnchecked.setAccessible(true);
+    public TestableLooper(Looper l) throws Exception {
+        this(InstrumentationRegistry.getInstrumentation().acquireLooperManager(l), l);
     }
 
-    public Looper getLooper() {
-        return mLooper;
+    private TestableLooper(TestLooperManager wrapper, Looper l) throws Exception {
+        mQueueWrapper = wrapper;
+        setupQueue(l);
     }
 
-    private void clearLooper() throws NoSuchFieldException, IllegalAccessException {
-        Field field = Looper.class.getDeclaredField("sThreadLocal");
-        field.setAccessible(true);
-        ThreadLocal<Looper> sThreadLocal = (ThreadLocal<Looper>) field.get(null);
-        sThreadLocal.set(null);
+    private TestableLooper(Looper looper, boolean b) throws Exception {
+        setupQueue(looper);
     }
 
-    private boolean setForCurrentThread() throws NoSuchFieldException, IllegalAccessException {
-        if (Looper.myLooper() != mLooper) {
-            Field field = Looper.class.getDeclaredField("sThreadLocal");
-            field.setAccessible(true);
-            ThreadLocal<Looper> sThreadLocal = (ThreadLocal<Looper>) field.get(null);
-            sThreadLocal.set(mLooper);
-            return true;
-        }
-        return false;
+    public Looper getLooper() {
+        return mLooper;
     }
 
-    private void setupQueue(boolean setMyLooper) throws Exception {
-        if (setMyLooper) {
-            clearLooper();
-            Looper.prepare();
-            mLooper = Looper.myLooper();
-        } else {
-            Constructor<Looper> constructor = Looper.class.getDeclaredConstructor(
-                    boolean.class);
-            constructor.setAccessible(true);
-            mLooper = constructor.newInstance(true);
-        }
-
+    private void setupQueue(Looper l) throws Exception {
+        mLooper = l;
         mQueue = mLooper.getQueue();
         mHandler = new Handler(mLooper);
     }
@@ -121,9 +92,7 @@ public class TestableLooper {
      * tests.
      */
     public void destroy() throws NoSuchFieldException, IllegalAccessException {
-        if (Looper.myLooper() == mLooper) {
-            clearLooper();
-        }
+        mQueueWrapper.release();
         if (mMain && mOriginalMain != null) {
             Field field = mLooper.getClass().getDeclaredField("sMainLooper");
             field.setAccessible(true);
@@ -156,34 +125,35 @@ public class TestableLooper {
 
     private int processQueuedMessages() {
         int count = 0;
-        mEmptyMessage = mHandler.obtainMessage(1);
-        mHandler.sendMessageDelayed(mEmptyMessage, 1);
+        mEmptyMessage = () -> { };
+        mHandler.post(mEmptyMessage);
+        waitForMessage(mQueueWrapper, mHandler, mEmptyMessage);
         while (parseMessageInt()) count++;
         return count;
     }
 
     private boolean parseMessageInt() {
         try {
-            Message result = (Message) mNext.invoke(mQueue);
+            Message result = mQueueWrapper.next();
             if (result != null) {
                 // This is a break message.
-                if (result == mEmptyMessage) {
-                    mRecycleUnchecked.invoke(result);
+                if (result.getCallback() == mEmptyMessage) {
+                    mQueueWrapper.recycle(result);
                     return false;
                 }
 
                 if (mMessageHandler != null) {
                     if (mMessageHandler.onMessageHandled(result)) {
                         result.getTarget().dispatchMessage(result);
-                        mRecycleUnchecked.invoke(result);
+                        mQueueWrapper.recycle(result);
                     } else {
-                        mRecycleUnchecked.invoke(result);
+                        mQueueWrapper.recycle(result);
                         // Message handler indicated it doesn't want us to continue.
                         return false;
                     }
                 } else {
                     result.getTarget().dispatchMessage(result);
-                    mRecycleUnchecked.invoke(result);
+                    mQueueWrapper.recycle(result);
                 }
             } else {
                 // No messages, don't continue parsing
@@ -199,10 +169,14 @@ public class TestableLooper {
      * Runs an executable with myLooper set and processes all messages added.
      */
     public void runWithLooper(RunnableWithException runnable) throws Exception {
-        boolean set = setForCurrentThread();
-        runnable.run();
+        new Handler(getLooper()).post(() -> {
+            try {
+                runnable.run();
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+        });
         processAllMessages();
-        if (set) clearLooper();
     }
 
     public interface RunnableWithException {
@@ -215,39 +189,132 @@ public class TestableLooper {
         boolean setAsMainLooper() default false;
     }
 
+    private static void waitForMessage(TestLooperManager queueWrapper, Handler handler,
+            Runnable execute) {
+        for (int i = 0; i < 10; i++) {
+            if (!queueWrapper.hasMessages(handler, null, execute)) {
+                try {
+                    Thread.sleep(1);
+                } catch (InterruptedException e) {
+                }
+            }
+        }
+        if (!queueWrapper.hasMessages(handler, null, execute)) {
+            throw new RuntimeException("Message didn't queue...");
+        }
+    }
+
     private static final Map<Object, TestableLooper> sLoopers = new ArrayMap<>();
 
     public static TestableLooper get(Object test) {
         return sLoopers.get(test);
     }
 
-    public static class LooperStatement extends Statement {
-        private final boolean mSetAsMain;
-        private final Statement mBase;
-        private final TestableLooper mLooper;
+    public static class LooperFrameworkMethod extends FrameworkMethod {
+        private HandlerThread mHandlerThread;
+
+        private final TestableLooper mTestableLooper;
+        private final Looper mLooper;
+        private final Handler mHandler;
 
-        public LooperStatement(Statement base, boolean setAsMain, Object test) {
-            mBase = base;
+        public LooperFrameworkMethod(FrameworkMethod base, boolean setAsMain, Object test) {
+            super(base.getMethod());
             try {
-                mLooper = new TestableLooper(false);
-                sLoopers.put(test, mLooper);
-                mSetAsMain = setAsMain;
+                mLooper = setAsMain ? Looper.getMainLooper() : createLooper();
+                mTestableLooper = new TestableLooper(mLooper, false);
             } catch (Exception e) {
                 throw new RuntimeException(e);
             }
+            sLoopers.put(test, mTestableLooper);
+            mHandler = new Handler(mLooper);
         }
 
-        @Override
-        public void evaluate() throws Throwable {
-            mLooper.setForCurrentThread();
-            if (mSetAsMain) {
-                mLooper.setAsMainLooper();
+        public LooperFrameworkMethod(TestableLooper other, FrameworkMethod base) {
+            super(base.getMethod());
+            mLooper = other.mLooper;
+            mTestableLooper = other;
+            mHandler = new Handler(mLooper);
+        }
+
+        public static FrameworkMethod get(FrameworkMethod base, boolean setAsMain, Object test) {
+            if (sLoopers.containsKey(test)) {
+                return new LooperFrameworkMethod(sLoopers.get(test), base);
             }
+            return new LooperFrameworkMethod(base, setAsMain, test);
+        }
 
+        @Override
+        public Object invokeExplosively(Object target, Object... params) throws Throwable {
+            if (Looper.myLooper() == mLooper) {
+                // Already on the right thread from another statement, just execute then.
+                return super.invokeExplosively(target, params);
+            }
+            boolean set = mTestableLooper.mQueueWrapper == null;
+            if (set) {
+                mTestableLooper.mQueueWrapper = InstrumentationRegistry.getInstrumentation()
+                        .acquireLooperManager(mLooper);
+            }
             try {
-                mBase.evaluate();
+                Object[] ret = new Object[1];
+                // Run the execution on the looper thread.
+                Runnable execute = () -> {
+                    try {
+                        ret[0] = super.invokeExplosively(target, params);
+                    } catch (Throwable throwable) {
+                        throw new LooperException(throwable);
+                    }
+                };
+                Message m = Message.obtain(mHandler, execute);
+
+                // Dispatch our message.
+                try {
+                    mTestableLooper.mQueueWrapper.execute(m);
+                } catch (LooperException e) {
+                    throw e.getSource();
+                } catch (RuntimeException re) {
+                    // If the TestLooperManager has to post, it will wrap what it throws in a
+                    // RuntimeException, make sure we grab the actual source.
+                    if (re.getCause() instanceof LooperException) {
+                        throw ((LooperException) re.getCause()).getSource();
+                    } else {
+                        throw re.getCause();
+                    }
+                } finally {
+                    m.recycle();
+                }
+                return ret[0];
             } finally {
-                mLooper.destroy();
+                if (set) {
+                    mTestableLooper.mQueueWrapper.release();
+                    mTestableLooper.mQueueWrapper = null;
+                }
+            }
+        }
+
+        private Looper createLooper() {
+            // TODO: Find way to share these.
+            mHandlerThread = new HandlerThread(TestableLooper.class.getSimpleName());
+            mHandlerThread.start();
+            return mHandlerThread.getLooper();
+        }
+
+        @Override
+        protected void finalize() throws Throwable {
+            super.finalize();
+            if (mHandlerThread != null) {
+                mHandlerThread.quit();
+            }
+        }
+
+        private static class LooperException extends RuntimeException {
+            private final Throwable mSource;
+
+            public LooperException(Throwable t) {
+                mSource = t;
+            }
+
+            public Throwable getSource() {
+                return mSource;
             }
         }
     }
diff --git a/tests/testables/tests/src/android/testing/TestableLooperTest.java b/tests/testables/tests/src/android/testing/TestableLooperTest.java
index 18e5fffef992..12f1d0a5f414 100644
--- a/tests/testables/tests/src/android/testing/TestableLooperTest.java
+++ b/tests/testables/tests/src/android/testing/TestableLooperTest.java
@@ -24,17 +24,16 @@ import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
 import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
 import android.testing.TestableLooper.MessageHandler;
 import android.testing.TestableLooper.RunWithLooper;
 
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-
 @RunWith(AndroidTestingRunner.class)
 @RunWithLooper
 public class TestableLooperTest {
@@ -46,11 +45,6 @@ public class TestableLooperTest {
         mTestableLooper = TestableLooper.get(this);
     }
 
-    @After
-    public void tearDown() throws Exception {
-        mTestableLooper.destroy();
-    }
-
     @Test
     public void testMessageExecuted() throws Exception {
         Handler h = new Handler();
@@ -133,39 +127,23 @@ public class TestableLooperTest {
     @Test
     public void testMainLooper() throws Exception {
         assertNotEquals(Looper.myLooper(), Looper.getMainLooper());
-
-        Looper originalMain = Looper.getMainLooper();
-        mTestableLooper.setAsMainLooper();
-        assertEquals(Looper.myLooper(), Looper.getMainLooper());
-        Runnable r = mock(Runnable.class);
-
-        new Handler(Looper.getMainLooper()).post(r);
-        mTestableLooper.processAllMessages();
-
-        verify(r).run();
-        mTestableLooper.destroy();
-
-        assertEquals(originalMain, Looper.getMainLooper());
-    }
-
-    @Test
-    public void testNotMyLooper() throws Exception {
-        TestableLooper looper = new TestableLooper(false);
-
-        assertEquals(Looper.myLooper(), mTestableLooper.getLooper());
-        assertNotEquals(Looper.myLooper(), looper.getLooper());
-
         Runnable r = mock(Runnable.class);
         Runnable r2 = mock(Runnable.class);
-        new Handler().post(r);
-        new Handler(looper.getLooper()).post(r2);
-
-        looper.processAllMessages();
-        verify(r2).run();
-        verify(r, never()).run();
-
-        mTestableLooper.processAllMessages();
-        verify(r).run();
+        TestableLooper testableLooper = new TestableLooper(Looper.getMainLooper());
+
+        try {
+            testableLooper.setMessageHandler(m -> {
+                if (m.getCallback() == r) return true;
+                return false;
+            });
+            new Handler(Looper.getMainLooper()).post(r);
+            testableLooper.processAllMessages();
+
+            verify(r).run();
+            verify(r2, never()).run();
+        } finally {
+            testableLooper.destroy();
+        }
     }
 
     @Test