Reduce main thread usage in sysui tests

Make assert more testable and fix up the tests in simple ways so
they still work but not on the main thread because that may or
may not, but is definitely not not flaky.

Test: atest frameworks/base/packages/SystemUI/tests
Change-Id: Ie19d772570b92c4b7f467f5ef9519f81500e63bc

36 files changed, 144 insertions, 113 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/qs/QSSecurityFooter.java b/packages/SystemUI/src/com/android/systemui/qs/QSSecurityFooter.java
index e2e943a369c2..d7d3981d05d1 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QSSecurityFooter.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/QSSecurityFooter.java
@@ -79,7 +79,7 @@ public class QSSecurityFooter implements OnClickListener, DialogInterface.OnClic
         mFooterIcon = (ImageView) mRootView.findViewById(R.id.footer_icon);
         mFooterIconId = R.drawable.ic_info_outline;
         mContext = context;
-        mMainHandler = new Handler(Looper.getMainLooper());
+        mMainHandler = new Handler(Looper.myLooper());
         mActivityStarter = Dependency.get(ActivityStarter.class);
         mSecurityController = Dependency.get(SecurityController.class);
         mHandler = new H(Dependency.get(Dependency.BG_LOOPER));
diff --git a/packages/SystemUI/src/com/android/systemui/util/Assert.java b/packages/SystemUI/src/com/android/systemui/util/Assert.java
index 0f7c9a462c0a..096ac3fcee1d 100644
--- a/packages/SystemUI/src/com/android/systemui/util/Assert.java
+++ b/packages/SystemUI/src/com/android/systemui/util/Assert.java
@@ -18,19 +18,24 @@ package com.android.systemui.util;
 
 import android.os.Looper;
 
+import com.android.internal.annotations.VisibleForTesting;
+
 /**
  * Helper providing common assertions.
  */
 public class Assert {
 
+    @VisibleForTesting
+    public static Looper sMainLooper = Looper.getMainLooper();
+
     public static void isMainThread() {
-        if (!Looper.getMainLooper().isCurrentThread()) {
+        if (!sMainLooper.isCurrentThread()) {
             throw new IllegalStateException("should be called from the main thread.");
         }
     }
 
     public static void isNotMainThread() {
-        if (Looper.getMainLooper().isCurrentThread()) {
+        if (sMainLooper.isCurrentThread()) {
             throw new IllegalStateException("should not be called from the main thread.");
         }
     }
diff --git a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardClockSwitchTest.java b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardClockSwitchTest.java
index 9cbe4152b5c4..7ca54231fe7b 100644
--- a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardClockSwitchTest.java
+++ b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardClockSwitchTest.java
@@ -51,7 +51,7 @@ import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
 @SmallTest
-@RunWithLooper(setAsMainLooper = true)
+@RunWithLooper
 @RunWith(AndroidTestingRunner.class)
 public class KeyguardClockSwitchTest extends SysuiTestCase {
     private PluginManager mPluginManager;
diff --git a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardPinBasedInputViewTest.java b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardPinBasedInputViewTest.java
index 359832f7a542..58870e4acbd0 100644
--- a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardPinBasedInputViewTest.java
+++ b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardPinBasedInputViewTest.java
@@ -38,7 +38,7 @@ import org.mockito.MockitoAnnotations;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-@RunWithLooper(setAsMainLooper = true)
+@RunWithLooper
 public class KeyguardPinBasedInputViewTest extends SysuiTestCase {
 
     @Mock
diff --git a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardSliceViewTest.java b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardSliceViewTest.java
index b98ce39f5ed3..77895c97051c 100644
--- a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardSliceViewTest.java
+++ b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardSliceViewTest.java
@@ -19,9 +19,14 @@ import android.graphics.Color;
 import android.net.Uri;
 import android.test.suitebuilder.annotation.SmallTest;
 import android.testing.AndroidTestingRunner;
+import android.testing.TestableLooper;
 import android.testing.TestableLooper.RunWithLooper;
 import android.view.LayoutInflater;
 
+import androidx.slice.SliceProvider;
+import androidx.slice.SliceSpecs;
+import androidx.slice.builders.ListBuilder;
+
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.keyguard.KeyguardSliceProvider;
 
@@ -34,12 +39,8 @@ import java.util.Collections;
 import java.util.HashSet;
 import java.util.concurrent.atomic.AtomicBoolean;
 
-import androidx.slice.SliceProvider;
-import androidx.slice.SliceSpecs;
-import androidx.slice.builders.ListBuilder;
-
 @SmallTest
-@RunWithLooper(setAsMainLooper = true)
+@RunWithLooper
 @RunWith(AndroidTestingRunner.class)
 public class KeyguardSliceViewTest extends SysuiTestCase {
     private KeyguardSliceView mKeyguardSliceView;
@@ -47,6 +48,7 @@ public class KeyguardSliceViewTest extends SysuiTestCase {
 
     @Before
     public void setUp() throws Exception {
+        com.android.systemui.util.Assert.sMainLooper = TestableLooper.get(this).getLooper();
         mKeyguardSliceView = (KeyguardSliceView) LayoutInflater.from(getContext())
                 .inflate(R.layout.keyguard_status_area, null);
         mSliceUri = Uri.parse(KeyguardSliceProvider.KEYGUARD_SLICE_URI);
diff --git a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardStatusViewTest.java b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardStatusViewTest.java
index 9e96df2c30cf..3582ab010413 100644
--- a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardStatusViewTest.java
+++ b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardStatusViewTest.java
@@ -20,10 +20,12 @@ import static org.mockito.Mockito.verify;
 
 import android.test.suitebuilder.annotation.SmallTest;
 import android.testing.AndroidTestingRunner;
+import android.testing.TestableLooper;
 import android.testing.TestableLooper.RunWithLooper;
 import android.view.LayoutInflater;
 
 import com.android.systemui.SysuiTestCase;
+import com.android.systemui.util.Assert;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -32,7 +34,7 @@ import org.mockito.InjectMocks;
 import org.mockito.Mock;
 
 @SmallTest
-@RunWithLooper(setAsMainLooper = true)
+@RunWithLooper
 @RunWith(AndroidTestingRunner.class)
 public class KeyguardStatusViewTest extends SysuiTestCase {
 
@@ -45,6 +47,7 @@ public class KeyguardStatusViewTest extends SysuiTestCase {
 
     @Before
     public void setUp() {
+        Assert.sMainLooper = TestableLooper.get(this).getLooper();
         LayoutInflater layoutInflater = LayoutInflater.from(getContext());
         mKeyguardStatusView =
                 (KeyguardStatusView) layoutInflater.inflate(R.layout.keyguard_status_view, null);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/ExpandHelperTest.java b/packages/SystemUI/tests/src/com/android/systemui/ExpandHelperTest.java
index c180ff8fd30c..948e0014a451 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/ExpandHelperTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/ExpandHelperTest.java
@@ -16,29 +16,29 @@
 
 package com.android.systemui;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
 import android.animation.ObjectAnimator;
 import android.content.Context;
-import android.support.test.InstrumentationRegistry;
 import android.support.test.annotation.UiThreadTest;
 import android.support.test.filters.SmallTest;
-import android.support.test.runner.AndroidJUnit4;
 import android.testing.AndroidTestingRunner;
+import android.testing.TestableLooper;
 import android.testing.TestableLooper.RunWithLooper;
 
-import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.NotificationTestHelper;
+import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
+import com.android.systemui.util.Assert;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
-
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-@RunWithLooper(setAsMainLooper = true)
+@RunWithLooper
 public class ExpandHelperTest extends SysuiTestCase {
 
     private ExpandableNotificationRow mRow;
@@ -47,6 +47,7 @@ public class ExpandHelperTest extends SysuiTestCase {
 
     @Before
     public void setUp() throws Exception {
+        Assert.sMainLooper = TestableLooper.get(this).getLooper();
         Context context = getContext();
         mRow = new NotificationTestHelper(context).createRow();
         mCallback = mock(ExpandHelper.Callback.class);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/SysuiTestCase.java b/packages/SystemUI/tests/src/com/android/systemui/SysuiTestCase.java
index a58bc8548bd4..9c5a59243d05 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/SysuiTestCase.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/SysuiTestCase.java
@@ -29,6 +29,7 @@ import android.testing.LeakCheck;
 import android.util.Log;
 
 import com.android.keyguard.KeyguardUpdateMonitor;
+import com.android.systemui.util.Assert;
 
 import org.junit.After;
 import org.junit.Before;
@@ -78,6 +79,8 @@ public abstract class SysuiTestCase {
     public void SysuiTeardown() {
         InstrumentationRegistry.registerInstance(mRealInstrumentation,
                 InstrumentationRegistry.getArguments());
+        // Reset the assert's main looper.
+        Assert.sMainLooper = Looper.getMainLooper();
     }
 
     protected LeakCheck getLeakCheck() {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeMachineTest.java b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeMachineTest.java
index 368c814f8e0a..e1c481e4ef28 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeMachineTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeMachineTest.java
@@ -28,8 +28,8 @@ import static com.android.systemui.doze.DozeMachine.State.UNINITIALIZED;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
-import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
@@ -40,14 +40,12 @@ import static org.mockito.Mockito.when;
 
 import android.support.test.filters.SmallTest;
 import android.testing.AndroidTestingRunner;
-import android.view.Display;
+import android.testing.UiThreadTest;
 
 import com.android.internal.hardware.AmbientDisplayConfiguration;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.util.wakelock.WakeLockFake;
 
-import android.testing.UiThreadTest;
-
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeTriggersTest.java b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeTriggersTest.java
index 5c8336c8dee1..31fc625d34dd 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeTriggersTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeTriggersTest.java
@@ -48,7 +48,7 @@ import org.junit.runner.RunWith;
 @SmallTest
 @Ignore("failing")
 @RunWith(AndroidTestingRunner.class)
-@RunWithLooper(setAsMainLooper = true)
+@RunWithLooper
 public class DozeTriggersTest extends SysuiTestCase {
     private DozeTriggers mTriggers;
     private DozeMachine mMachine;
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardSliceProviderTest.java b/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardSliceProviderTest.java
index 095395185f86..c1c80ce4a70a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardSliceProviderTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardSliceProviderTest.java
@@ -55,7 +55,7 @@ import java.util.concurrent.TimeUnit;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-@RunWithLooper(setAsMainLooper = true)
+@RunWithLooper
 public class KeyguardSliceProviderTest extends SysuiTestCase {
 
     @Mock
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/QSFragmentTest.java b/packages/SystemUI/tests/src/com/android/systemui/qs/QSFragmentTest.java
index 4e24354d9878..bc7d9836d6f8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/QSFragmentTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/QSFragmentTest.java
@@ -18,39 +18,32 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.mock;
 
-import android.app.FragmentController;
-import android.app.FragmentManagerNonConfig;
+import android.content.Context;
 import android.os.Looper;
-import android.support.test.filters.FlakyTest;
+import android.support.test.filters.SmallTest;
+import android.testing.AndroidTestingRunner;
+import android.testing.LayoutInflaterBuilder;
+import android.testing.TestableLooper;
+import android.testing.TestableLooper.RunWithLooper;
+import android.view.View;
+import android.widget.FrameLayout;
 
 import com.android.internal.logging.MetricsLogger;
 import com.android.keyguard.CarrierText;
 import com.android.systemui.Dependency;
 import com.android.systemui.R;
-
-import android.os.Parcelable;
-import android.support.test.filters.SmallTest;
-import android.testing.AndroidTestingRunner;
-
 import com.android.systemui.SysuiBaseFragmentTest;
 import com.android.systemui.statusbar.phone.StatusBarIconController;
 import com.android.systemui.statusbar.policy.Clock;
 import com.android.systemui.statusbar.policy.UserSwitcherController;
-import android.testing.LayoutInflaterBuilder;
-import android.testing.TestableLooper;
-import android.testing.TestableLooper.RunWithLooper;
 
 import org.junit.Before;
 import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
-import android.content.Context;
-import android.view.View;
-import android.widget.FrameLayout;
-
 @RunWith(AndroidTestingRunner.class)
-@RunWithLooper(setAsMainLooper = true)
+@RunWithLooper
 @SmallTest
 @Ignore
 public class QSFragmentTest extends SysuiBaseFragmentTest {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/QSSecurityFooterTest.java b/packages/SystemUI/tests/src/com/android/systemui/qs/QSSecurityFooterTest.java
index 33b347a66b33..fd31013db429 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/QSSecurityFooterTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/QSSecurityFooterTest.java
@@ -22,7 +22,6 @@ import static org.mockito.Mockito.when;
 
 import android.content.Context;
 import android.content.pm.UserInfo;
-import android.os.Looper;
 import android.os.UserManager;
 import android.provider.Settings;
 import android.test.suitebuilder.annotation.SmallTest;
@@ -58,7 +57,7 @@ import org.mockito.Mockito;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-@RunWithLooper(setAsMainLooper = true)
+@RunWithLooper
 public class QSSecurityFooterTest extends SysuiTestCase {
 
     private final String MANAGING_ORGANIZATION = "organization";
@@ -76,7 +75,8 @@ public class QSSecurityFooterTest extends SysuiTestCase {
     @Before
     public void setUp() {
         mDependency.injectTestDependency(SecurityController.class, mSecurityController);
-        mDependency.injectTestDependency(Dependency.BG_LOOPER, Looper.getMainLooper());
+        mDependency.injectTestDependency(Dependency.BG_LOOPER,
+                TestableLooper.get(this).getLooper());
         mContext.addMockSystemService(Context.LAYOUT_INFLATER_SERVICE,
                 new LayoutInflaterBuilder(mContext)
                         .replace("ImageView", TestableImageView.class)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/external/TileServicesTest.java b/packages/SystemUI/tests/src/com/android/systemui/qs/external/TileServicesTest.java
index c016a851010a..c6597b9bd534 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/external/TileServicesTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/external/TileServicesTest.java
@@ -24,16 +24,15 @@ import android.content.ComponentName;
 import android.os.Looper;
 import android.service.quicksettings.Tile;
 import android.test.suitebuilder.annotation.SmallTest;
-
 import android.testing.AndroidTestingRunner;
+import android.testing.TestableLooper;
+import android.testing.TestableLooper.RunWithLooper;
+
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.qs.QSTileHost;
 import com.android.systemui.statusbar.phone.StatusBarIconController;
 import com.android.systemui.statusbar.policy.BluetoothController;
 
-import android.testing.TestableLooper;
-import android.testing.TestableLooper.RunWithLooper;
-
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -45,7 +44,7 @@ import java.util.ArrayList;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-@RunWithLooper(setAsMainLooper = true)
+@RunWithLooper
 public class TileServicesTest extends SysuiTestCase {
     private static int NUM_FAKES = TileServices.DEFAULT_MAX_BOUND * 2;
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NonPhoneDependencyTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NonPhoneDependencyTest.java
index 2e280d336aab..9449e297fcc5 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NonPhoneDependencyTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NonPhoneDependencyTest.java
@@ -50,7 +50,7 @@ import org.mockito.MockitoAnnotations;
  */
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-@TestableLooper.RunWithLooper(setAsMainLooper = true)
+@TestableLooper.RunWithLooper
 @Ignore("b/118400112")
 public class NonPhoneDependencyTest extends SysuiTestCase {
     @Mock private NotificationPresenter mPresenter;
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationListenerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationListenerTest.java
index 63ececbe2994..65c04fe4bcd3 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationListenerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationListenerTest.java
@@ -16,15 +16,12 @@
 
 package com.android.systemui.statusbar;
 
-import static junit.framework.Assert.assertTrue;
-
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.app.Notification;
 import android.os.Handler;
-import android.os.Looper;
 import android.os.UserHandle;
 import android.service.notification.NotificationListenerService;
 import android.service.notification.StatusBarNotification;
@@ -45,7 +42,7 @@ import org.mockito.MockitoAnnotations;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-@TestableLooper.RunWithLooper(setAsMainLooper = true)
+@TestableLooper.RunWithLooper
 public class NotificationListenerTest extends SysuiTestCase {
     private static final String TEST_PACKAGE_NAME = "test";
     private static final int TEST_UID = 0;
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationViewHierarchyManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationViewHierarchyManagerTest.java
index 72d6cd8eaeea..520a927d1ab0 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationViewHierarchyManagerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationViewHierarchyManagerTest.java
@@ -46,6 +46,7 @@ import com.android.systemui.statusbar.notification.row.ExpandableView;
 import com.android.systemui.statusbar.notification.stack.NotificationListContainer;
 import com.android.systemui.statusbar.phone.NotificationGroupManager;
 import com.android.systemui.statusbar.phone.ShadeController;
+import com.android.systemui.util.Assert;
 
 import com.google.android.collect.Lists;
 
@@ -60,7 +61,7 @@ import java.util.List;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-@TestableLooper.RunWithLooper(setAsMainLooper = true)
+@TestableLooper.RunWithLooper
 public class NotificationViewHierarchyManagerTest extends SysuiTestCase {
     @Mock private NotificationPresenter mPresenter;
     @Mock private NotificationData mNotificationData;
@@ -79,6 +80,7 @@ public class NotificationViewHierarchyManagerTest extends SysuiTestCase {
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
+        Assert.sMainLooper = TestableLooper.get(this).getLooper();
         mDependency.injectTestDependency(NotificationEntryManager.class, mEntryManager);
         mDependency.injectTestDependency(NotificationLockscreenUserManager.class,
                 mLockscreenUserManager);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/AboveShelfObserverTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/AboveShelfObserverTest.java
index b7aa21b86d33..db2c8780e783 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/AboveShelfObserverTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/AboveShelfObserverTest.java
@@ -19,15 +19,15 @@ package com.android.systemui.statusbar.notification;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 
-import android.support.test.runner.AndroidJUnit4;
 import android.test.suitebuilder.annotation.SmallTest;
 import android.testing.AndroidTestingRunner;
+import android.testing.TestableLooper;
 import android.testing.TestableLooper.RunWithLooper;
 import android.widget.FrameLayout;
 
 import com.android.systemui.SysuiTestCase;
-import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.NotificationTestHelper;
+import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 
 import org.junit.Assert;
 import org.junit.Before;
@@ -36,7 +36,7 @@ import org.junit.runner.RunWith;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-@RunWithLooper(setAsMainLooper = true)
+@RunWithLooper
 public class AboveShelfObserverTest extends SysuiTestCase {
 
     private AboveShelfObserver mObserver;
@@ -46,6 +46,7 @@ public class AboveShelfObserverTest extends SysuiTestCase {
 
     @Before
     public void setUp() throws Exception {
+        com.android.systemui.util.Assert.sMainLooper = TestableLooper.get(this).getLooper();
         mNotificationTestHelper = new NotificationTestHelper(getContext());
         mHostLayout = new FrameLayout(getContext());
         mObserver = new AboveShelfObserver(mHostLayout);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/NotificationDataTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/NotificationDataTest.java
index f94ba95999bf..8e88ed0556bf 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/NotificationDataTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/NotificationDataTest.java
@@ -53,6 +53,7 @@ import android.service.notification.StatusBarNotification;
 import android.support.test.annotation.UiThreadTest;
 import android.support.test.filters.SmallTest;
 import android.testing.AndroidTestingRunner;
+import android.testing.TestableLooper;
 import android.testing.TestableLooper.RunWithLooper;
 import android.util.ArraySet;
 
@@ -78,7 +79,7 @@ import java.util.List;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-@RunWithLooper(setAsMainLooper = true)
+@RunWithLooper
 public class NotificationDataTest extends SysuiTestCase {
 
     private static final int UID_NORMAL = 123;
@@ -101,6 +102,7 @@ public class NotificationDataTest extends SysuiTestCase {
 
     @Before
     public void setUp() throws Exception {
+        com.android.systemui.util.Assert.sMainLooper = TestableLooper.get(this).getLooper();
         MockitoAnnotations.initMocks(this);
         when(mMockStatusBarNotification.getUid()).thenReturn(UID_NORMAL);
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/NotificationViewWrapperTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/NotificationViewWrapperTest.java
index 63d1e8dbc954..24aa772e2fc1 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/NotificationViewWrapperTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/NotificationViewWrapperTest.java
@@ -17,28 +17,29 @@
 package com.android.systemui.statusbar.notification;
 
 import android.content.Context;
-import android.support.test.InstrumentationRegistry;
 import android.support.test.filters.SmallTest;
-import android.support.test.runner.AndroidJUnit4;
 import android.testing.AndroidTestingRunner;
+import android.testing.TestableLooper;
 import android.testing.TestableLooper.RunWithLooper;
 import android.view.View;
 
 import com.android.systemui.SysuiTestCase;
-import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.NotificationTestHelper;
+import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.notification.row.wrapper.NotificationViewWrapper;
+import com.android.systemui.util.Assert;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
 @RunWith(AndroidTestingRunner.class)
 @SmallTest
-@RunWithLooper(setAsMainLooper = true)
+@RunWithLooper
 public class NotificationViewWrapperTest extends SysuiTestCase {
 
     @Test
     public void constructor_doesntUseViewContext() throws Exception {
+        Assert.sMainLooper = TestableLooper.get(this).getLooper();
         new TestableNotificationViewWrapper(mContext,
                 new View(mContext),
                 new NotificationTestHelper(getContext()).createRow());
@@ -50,4 +51,4 @@ public class NotificationViewWrapperTest extends SysuiTestCase {
             super(ctx, view, row);
         }
     }
-}
\ No newline at end of file
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationLoggerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationLoggerTest.java
index 3710fa833d50..512acd073a84 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationLoggerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationLoggerTest.java
@@ -37,13 +37,13 @@ import android.testing.TestableLooper;
 import com.android.internal.statusbar.IStatusBarService;
 import com.android.internal.statusbar.NotificationVisibility;
 import com.android.systemui.SysuiTestCase;
-
-import com.android.systemui.statusbar.notification.NotificationEntryManager;
 import com.android.systemui.statusbar.NotificationListener;
 import com.android.systemui.statusbar.NotificationPresenter;
 import com.android.systemui.statusbar.notification.NotificationData;
+import com.android.systemui.statusbar.notification.NotificationEntryManager;
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.notification.stack.NotificationListContainer;
+
 import com.google.android.collect.Lists;
 
 import org.junit.Before;
@@ -58,7 +58,7 @@ import java.util.concurrent.ConcurrentLinkedQueue;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-@TestableLooper.RunWithLooper(setAsMainLooper = true)
+@TestableLooper.RunWithLooper
 public class NotificationLoggerTest extends SysuiTestCase {
     private static final String TEST_PACKAGE_NAME = "test";
     private static final int TEST_UID = 0;
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowTest.java
index 2da72e7858c8..6d3553912701 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowTest.java
@@ -40,6 +40,7 @@ import android.app.AppOpsManager;
 import android.app.NotificationChannel;
 import android.support.test.filters.SmallTest;
 import android.testing.AndroidTestingRunner;
+import android.testing.TestableLooper;
 import android.testing.TestableLooper.RunWithLooper;
 import android.util.ArraySet;
 import android.view.NotificationHeaderView;
@@ -64,7 +65,7 @@ import java.util.List;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-@RunWithLooper(setAsMainLooper = true)
+@RunWithLooper
 public class ExpandableNotificationRowTest extends SysuiTestCase {
 
     private ExpandableNotificationRow mGroupRow;
@@ -77,6 +78,7 @@ public class ExpandableNotificationRowTest extends SysuiTestCase {
 
     @Before
     public void setUp() throws Exception {
+        com.android.systemui.util.Assert.sMainLooper = TestableLooper.get(this).getLooper();
         mNotificationTestHelper = new NotificationTestHelper(mContext);
         mGroupRow = mNotificationTestHelper.createGroup();
         mGroupRow.setHeadsUpAnimatingAwayListener(
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationBlockingHelperManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationBlockingHelperManagerTest.java
index 4efab5385c0a..669b98e1b279 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationBlockingHelperManagerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationBlockingHelperManagerTest.java
@@ -16,10 +16,18 @@
 
 package com.android.systemui.statusbar.notification.row;
 
-import com.android.systemui.SysuiTestCase;
-import com.android.systemui.plugins.statusbar.NotificationMenuRowPlugin;
-import com.android.systemui.statusbar.notification.NotificationEntryManager;
-import com.android.systemui.statusbar.NotificationTestHelper;
+import static android.service.notification.NotificationListenerService.Ranking.USER_SENTIMENT_NEGATIVE;
+import static android.service.notification.NotificationListenerService.Ranking.USER_SENTIMENT_NEUTRAL;
+import static android.service.notification.NotificationListenerService.Ranking.USER_SENTIMENT_POSITIVE;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import android.content.Context;
 import android.support.test.filters.FlakyTest;
@@ -28,30 +36,24 @@ import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 import android.view.View;
 
+import com.android.systemui.SysuiTestCase;
+import com.android.systemui.plugins.statusbar.NotificationMenuRowPlugin;
+import com.android.systemui.statusbar.NotificationTestHelper;
+import com.android.systemui.statusbar.notification.NotificationEntryManager;
+import com.android.systemui.util.Assert;
+
 import org.junit.Before;
 import org.junit.Test;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import static android.service.notification.NotificationListenerService.Ranking.USER_SENTIMENT_NEGATIVE;
-import static android.service.notification.NotificationListenerService.Ranking.USER_SENTIMENT_NEUTRAL;
-import static android.service.notification.NotificationListenerService.Ranking.USER_SENTIMENT_POSITIVE;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.Mockito.spy;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
 /**
  * Tests for {@link NotificationBlockingHelperManager}.
  */
 @SmallTest
 @FlakyTest
 @org.junit.runner.RunWith(AndroidTestingRunner.class)
-@TestableLooper.RunWithLooper(setAsMainLooper = true)
+@TestableLooper.RunWithLooper
 public class NotificationBlockingHelperManagerTest extends SysuiTestCase {
 
     private NotificationBlockingHelperManager mBlockingHelperManager;
@@ -65,6 +67,7 @@ public class NotificationBlockingHelperManagerTest extends SysuiTestCase {
 
     @Before
     public void setUp() {
+        Assert.sMainLooper = TestableLooper.get(this).getLooper();
         MockitoAnnotations.initMocks(this);
         when(mGutsManager.openGuts(
                 any(View.class),
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationGutsManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationGutsManagerTest.java
index 766c5d2377c6..8966aca3069c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationGutsManagerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationGutsManagerTest.java
@@ -64,6 +64,7 @@ import com.android.systemui.statusbar.notification.NotificationEntryManager;
 import com.android.systemui.statusbar.notification.row.NotificationGutsManager.OnSettingsClickListener;
 import com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayout;
 import com.android.systemui.statusbar.policy.DeviceProvisionedController;
+import com.android.systemui.util.Assert;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -79,7 +80,7 @@ import org.mockito.junit.MockitoRule;
  */
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-@TestableLooper.RunWithLooper(setAsMainLooper = true)
+@TestableLooper.RunWithLooper
 public class NotificationGutsManagerTest extends SysuiTestCase {
     private static final String TEST_CHANNEL_ID = "NotificationManagerServiceTestChannelId";
 
@@ -101,6 +102,7 @@ public class NotificationGutsManagerTest extends SysuiTestCase {
     @Before
     public void setUp() {
         mTestableLooper = TestableLooper.get(this);
+        Assert.sMainLooper = TestableLooper.get(this).getLooper();
         mDependency.injectTestDependency(DeviceProvisionedController.class,
                 mDeviceProvisionedController);
         mHandler = Handler.createAsync(mTestableLooper.getLooper());
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationMenuRowTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationMenuRowTest.java
index bb9911b24f0a..e4d019656072 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationMenuRowTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationMenuRowTest.java
@@ -44,7 +44,7 @@ import org.junit.runner.RunWith;
 import org.mockito.Mockito;
 
 @RunWith(AndroidTestingRunner.class)
-@RunWithLooper()
+@RunWithLooper
 @SmallTest
 public class NotificationMenuRowTest extends LeakCheckedTest {
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationCustomViewWrapperTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationCustomViewWrapperTest.java
index 4b94a2523cfe..fed66af07bcc 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationCustomViewWrapperTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationCustomViewWrapperTest.java
@@ -18,6 +18,7 @@ package com.android.systemui.statusbar.notification.row.wrapper;
 
 import android.support.test.filters.SmallTest;
 import android.testing.AndroidTestingRunner;
+import android.testing.TestableLooper;
 import android.testing.TestableLooper.RunWithLooper;
 import android.view.View;
 import android.widget.RemoteViews;
@@ -34,13 +35,14 @@ import org.junit.runner.RunWith;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-@RunWithLooper(setAsMainLooper = true)
+@RunWithLooper
 public class NotificationCustomViewWrapperTest extends SysuiTestCase {
 
     private ExpandableNotificationRow mRow;
 
     @Before
     public void setUp() throws Exception {
+        com.android.systemui.util.Assert.sMainLooper = TestableLooper.get(this).getLooper();
         mRow = new NotificationTestHelper(mContext).createRow();
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationChildrenContainerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationChildrenContainerTest.java
index 272845396e27..bbafb4e4a211 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationChildrenContainerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationChildrenContainerTest.java
@@ -18,13 +18,14 @@ package com.android.systemui.statusbar.notification.stack;
 
 import android.support.test.filters.SmallTest;
 import android.testing.AndroidTestingRunner;
+import android.testing.TestableLooper;
 import android.testing.TestableLooper.RunWithLooper;
 import android.view.NotificationHeaderView;
 import android.view.View;
 
 import com.android.systemui.SysuiTestCase;
-import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.NotificationTestHelper;
+import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 
 import org.junit.Assert;
 import org.junit.Before;
@@ -33,16 +34,16 @@ import org.junit.runner.RunWith;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-@RunWithLooper(setAsMainLooper = true)
+@RunWithLooper
 public class NotificationChildrenContainerTest extends SysuiTestCase {
 
     private ExpandableNotificationRow mGroup;
-    private int mId;
     private NotificationTestHelper mNotificationTestHelper;
     private NotificationChildrenContainer mChildrenContainer;
 
     @Before
     public void setUp() throws Exception {
+        com.android.systemui.util.Assert.sMainLooper = TestableLooper.get(this).getLooper();
         mNotificationTestHelper = new NotificationTestHelper(mContext);
         mGroup = mNotificationTestHelper.createGroup();
         mChildrenContainer = mGroup.getChildrenContainer();
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationRoundnessManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationRoundnessManagerTest.java
index 662e016b77ff..8ae7d52a5da5 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationRoundnessManagerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationRoundnessManagerTest.java
@@ -25,6 +25,7 @@ import static org.mockito.Mockito.when;
 
 import android.support.test.filters.SmallTest;
 import android.testing.AndroidTestingRunner;
+import android.testing.TestableLooper;
 import android.testing.TestableLooper.RunWithLooper;
 
 import com.android.systemui.SysuiTestCase;
@@ -41,7 +42,7 @@ import java.util.HashSet;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-@RunWithLooper(setAsMainLooper = true)
+@RunWithLooper
 public class NotificationRoundnessManagerTest extends SysuiTestCase {
 
     private NotificationRoundnessManager mRoundnessManager = new NotificationRoundnessManager();
@@ -52,6 +53,7 @@ public class NotificationRoundnessManagerTest extends SysuiTestCase {
 
     @Before
     public void setUp() throws Exception {
+        com.android.systemui.util.Assert.sMainLooper = TestableLooper.get(this).getLooper();
         NotificationTestHelper testHelper = new NotificationTestHelper(getContext());
         mFirst = testHelper.createRow();
         mFirst.setHeadsUpAnimatingAwayListener(animatingAway
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/HeadsUpAppearanceControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/HeadsUpAppearanceControllerTest.java
index c99e766ac697..4f6329cb0c57 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/HeadsUpAppearanceControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/HeadsUpAppearanceControllerTest.java
@@ -23,18 +23,18 @@ import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.support.test.filters.SmallTest;
-import android.support.test.runner.AndroidJUnit4;
 import android.testing.AndroidTestingRunner;
+import android.testing.TestableLooper;
 import android.testing.TestableLooper.RunWithLooper;
 import android.view.View;
 import android.widget.TextView;
 
 import com.android.systemui.SysuiTestCase;
-import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.HeadsUpStatusBarView;
 import com.android.systemui.statusbar.NotificationTestHelper;
-import com.android.systemui.statusbar.policy.DarkIconDispatcher;
+import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayout;
+import com.android.systemui.statusbar.policy.DarkIconDispatcher;
 
 import org.junit.Assert;
 import org.junit.Before;
@@ -43,7 +43,7 @@ import org.junit.runner.RunWith;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-@RunWithLooper(setAsMainLooper = true)
+@RunWithLooper
 public class HeadsUpAppearanceControllerTest extends SysuiTestCase {
 
     private final NotificationStackScrollLayout mStackScroller =
@@ -58,6 +58,7 @@ public class HeadsUpAppearanceControllerTest extends SysuiTestCase {
 
     @Before
     public void setUp() throws Exception {
+        com.android.systemui.util.Assert.sMainLooper = TestableLooper.get(this).getLooper();
         NotificationTestHelper testHelper = new NotificationTestHelper(getContext());
         mFirst = testHelper.createRow();
         mDependency.injectTestDependency(DarkIconDispatcher.class, mDarkIconDispatcher);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardBouncerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardBouncerTest.java
index 7f8668fa064c..f7a95c50fee8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardBouncerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardBouncerTest.java
@@ -28,6 +28,7 @@ import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.verifyZeroInteractions;
 import static org.mockito.Mockito.when;
 
+import android.content.res.ColorStateList;
 import android.graphics.Color;
 import android.support.test.filters.SmallTest;
 import android.testing.AndroidTestingRunner;
@@ -35,7 +36,6 @@ import android.testing.TestableLooper;
 import android.view.ViewGroup;
 import android.view.ViewTreeObserver;
 import android.widget.FrameLayout;
-import android.content.res.ColorStateList;
 
 import com.android.internal.widget.LockPatternUtils;
 import com.android.keyguard.KeyguardHostView;
@@ -56,7 +56,7 @@ import org.mockito.MockitoAnnotations;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-@TestableLooper.RunWithLooper(setAsMainLooper = true)
+@TestableLooper.RunWithLooper
 public class KeyguardBouncerTest extends SysuiTestCase {
 
     @Mock
@@ -78,6 +78,7 @@ public class KeyguardBouncerTest extends SysuiTestCase {
 
     @Before
     public void setup() {
+        com.android.systemui.util.Assert.sMainLooper = TestableLooper.get(this).getLooper();
         MockitoAnnotations.initMocks(this);
         DejankUtils.setImmediate(true);
         final ViewGroup container = new FrameLayout(getContext());
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardClockPositionAlgorithmTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardClockPositionAlgorithmTest.java
index 31014b8a62e9..27ed9c5a14c8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardClockPositionAlgorithmTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardClockPositionAlgorithmTest.java
@@ -30,7 +30,7 @@ import org.junit.runner.RunWith;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-@TestableLooper.RunWithLooper(setAsMainLooper = true)
+@TestableLooper.RunWithLooper
 public class KeyguardClockPositionAlgorithmTest extends SysuiTestCase {
 
     private static final int SCREEN_HEIGHT = 2000;
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardPresentationTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardPresentationTest.java
index 54291536093c..c5875554b073 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardPresentationTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardPresentationTest.java
@@ -20,21 +20,20 @@ import android.support.test.filters.SmallTest;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 import android.view.LayoutInflater;
-import android.view.View;
 
 import com.android.systemui.R;
 import com.android.systemui.SysuiTestCase;
 
-import org.junit.Assert;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-@TestableLooper.RunWithLooper(setAsMainLooper = true)
+@TestableLooper.RunWithLooper
 public class KeyguardPresentationTest extends SysuiTestCase {
     @Test
     public void testInflation_doesntCrash() {
+        com.android.systemui.util.Assert.sMainLooper = TestableLooper.get(this).getLooper();
         LayoutInflater inflater = LayoutInflater.from(getContext());
         inflater.inflate(R.layout.keyguard_presentation, null);
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/QuickStepControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/QuickStepControllerTest.java
index 1783d0cf2519..cdaa2420186f 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/QuickStepControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/QuickStepControllerTest.java
@@ -29,7 +29,6 @@ import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.anyFloat;
-import static org.mockito.Mockito.atLeast;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -38,11 +37,6 @@ import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
-import com.android.systemui.R;
-import com.android.systemui.recents.OverviewProxyService;
-import com.android.systemui.shared.recents.IOverviewProxy;
-import com.android.systemui.SysuiTestCase;
-
 import android.content.Context;
 import android.content.res.Resources;
 import android.support.test.filters.SmallTest;
@@ -51,6 +45,11 @@ import android.testing.TestableLooper.RunWithLooper;
 import android.view.MotionEvent;
 import android.view.View;
 
+import com.android.systemui.R;
+import com.android.systemui.SysuiTestCase;
+import com.android.systemui.recents.OverviewProxyService;
+import com.android.systemui.shared.recents.IOverviewProxy;
+
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -59,7 +58,7 @@ import org.mockito.MockitoAnnotations;
 
 /** atest QuickStepControllerTest */
 @RunWith(AndroidTestingRunner.class)
-@RunWithLooper(setAsMainLooper = true)
+@RunWithLooper
 @SmallTest
 public class QuickStepControllerTest extends SysuiTestCase {
     private QuickStepController mController;
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/SystemUIDialogTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/SystemUIDialogTest.java
index dcd531dc9eb7..090963be7ac7 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/SystemUIDialogTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/SystemUIDialogTest.java
@@ -15,6 +15,7 @@
 package com.android.systemui.statusbar.phone;
 
 import static junit.framework.Assert.assertTrue;
+
 import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
@@ -35,7 +36,7 @@ import org.mockito.ArgumentCaptor;
 
 
 @RunWith(AndroidTestingRunner.class)
-@RunWithLooper(setAsMainLooper = true)
+@RunWithLooper
 @SmallTest
 public class SystemUIDialogTest extends SysuiTestCase {
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/RemoteInputViewTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/RemoteInputViewTest.java
index 3164c0469c40..b3ac6be65d36 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/RemoteInputViewTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/RemoteInputViewTest.java
@@ -32,9 +32,10 @@ import android.widget.ImageButton;
 import com.android.systemui.Dependency;
 import com.android.systemui.R;
 import com.android.systemui.SysuiTestCase;
-import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.NotificationTestHelper;
 import com.android.systemui.statusbar.RemoteInputController;
+import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
+import com.android.systemui.util.Assert;
 
 import org.junit.After;
 import org.junit.Before;
@@ -44,7 +45,7 @@ import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
 @RunWith(AndroidTestingRunner.class)
-@TestableLooper.RunWithLooper(setAsMainLooper = true)
+@TestableLooper.RunWithLooper
 @SmallTest
 public class RemoteInputViewTest extends SysuiTestCase {
 
@@ -60,6 +61,7 @@ public class RemoteInputViewTest extends SysuiTestCase {
 
     @Before
     public void setUp() throws Exception {
+        Assert.sMainLooper = TestableLooper.get(this).getLooper();
         MockitoAnnotations.initMocks(this);
 
         mDependency.injectTestDependency(RemoteInputQuickSettingsDisabler.class,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/util/wakelock/KeepAwakeAnimationListenerTest.java b/packages/SystemUI/tests/src/com/android/systemui/util/wakelock/KeepAwakeAnimationListenerTest.java
index 43942f72993e..ab9b0c979fd3 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/util/wakelock/KeepAwakeAnimationListenerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/util/wakelock/KeepAwakeAnimationListenerTest.java
@@ -20,11 +20,13 @@ import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 
 import android.animation.Animator;
+import android.os.Looper;
 import android.support.test.filters.SmallTest;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 
 import com.android.systemui.SysuiTestCase;
+import com.android.systemui.util.Assert;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -33,7 +35,7 @@ import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
 @SmallTest
-@TestableLooper.RunWithLooper(setAsMainLooper = true)
+@TestableLooper.RunWithLooper
 @RunWith(AndroidTestingRunner.class)
 public class KeepAwakeAnimationListenerTest extends SysuiTestCase {
     @Mock WakeLock mWakeLock;
@@ -41,6 +43,7 @@ public class KeepAwakeAnimationListenerTest extends SysuiTestCase {
 
     @Before
     public void setup() {
+        Assert.sMainLooper = TestableLooper.get(this).getLooper();
         MockitoAnnotations.initMocks(this);
         KeepAwakeAnimationListener.sWakeLock = mWakeLock;
         mKeepAwakeAnimationListener = new KeepAwakeAnimationListener(getContext());
@@ -55,4 +58,10 @@ public class KeepAwakeAnimationListenerTest extends SysuiTestCase {
         mKeepAwakeAnimationListener.onAnimationEnd((Animator) null);
         verify(mWakeLock).release();
     }
+
+    @Test(expected = IllegalStateException.class)
+    public void initThrows_onNonMainThread() {
+        Assert.sMainLooper = Looper.getMainLooper();
+        new KeepAwakeAnimationListener(getContext());
+    }
 }