Merge "Speed up sysui tests"

10 files changed, 193 insertions, 142 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/qs/customize/QSCustomizer.java b/packages/SystemUI/src/com/android/systemui/qs/customize/QSCustomizer.java
index 9bbead46c33b..0be53b4b7c26 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/customize/QSCustomizer.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/customize/QSCustomizer.java
@@ -60,8 +60,6 @@ public class QSCustomizer extends LinearLayout implements OnMenuItemClickListene
 
     private final QSDetailClipper mClipper;
 
-    private PhoneStatusBar mPhoneStatusBar;
-
     private boolean isShown;
     private QSTileHost mHost;
     private RecyclerView mRecyclerView;
@@ -119,7 +117,6 @@ public class QSCustomizer extends LinearLayout implements OnMenuItemClickListene
 
     public void setHost(QSTileHost host) {
         mHost = host;
-        mPhoneStatusBar = host.getPhoneStatusBar();
         mTileAdapter.setHost(host);
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/QSTileHost.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/QSTileHost.java
index 567ab3b5966d..227ebdfacbda 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/QSTileHost.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/QSTileHost.java
@@ -165,10 +165,6 @@ public class QSTileHost implements QSTile.Host, Tunable {
         mHeader = view;
     }
 
-    public PhoneStatusBar getPhoneStatusBar() {
-        return mStatusBar;
-    }
-
     public void destroy() {
         mHandlerThread.quitSafely();
         mTiles.values().forEach(tile -> tile.destroy());
diff --git a/packages/SystemUI/src/com/android/systemui/util/LayoutInflaterBuilder.java b/packages/SystemUI/src/com/android/systemui/util/LayoutInflaterBuilder.java
index f42092130af1..5cfe677efd57 100644
--- a/packages/SystemUI/src/com/android/systemui/util/LayoutInflaterBuilder.java
+++ b/packages/SystemUI/src/com/android/systemui/util/LayoutInflaterBuilder.java
@@ -81,11 +81,23 @@ public class LayoutInflaterBuilder {
      * @return Builder object post-modification.
      */
     public LayoutInflaterBuilder replace(@NonNull Class from, @NonNull Class to) {
+        return replace(from.getName(), to);
+    }
+
+    /**
+     * Instructs the Builder to configure the LayoutInflater such that all instances
+     * of one {@link View} will be replaced with instances of another during inflation.
+     *
+     * @param tag Instances of this tag will be replaced during inflation.
+     * @param to Instances of this class will be inflated as replacements.
+     * @return Builder object post-modification.
+     */
+    public LayoutInflaterBuilder replace(@NonNull String tag, @NonNull Class to) {
         assertIfAlreadyBuilt();
         if (mReplaceMap == null) {
             mReplaceMap = new ArrayMap<String, String>();
         }
-        mReplaceMap.put(from.getName(), to.getName());
+        mReplaceMap.put(tag, to.getName());
         return this;
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/QSFooterTest.java b/packages/SystemUI/tests/src/com/android/systemui/qs/QSFooterTest.java
index 391c45fb8487..8acd6ba5e1d7 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/QSFooterTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/QSFooterTest.java
@@ -14,37 +14,31 @@
 
 package com.android.systemui.qs;
 
+import static junit.framework.Assert.assertEquals;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
 import android.content.Context;
-import android.content.res.Resources;
+import android.os.Handler;
 import android.os.Looper;
 import android.support.test.runner.AndroidJUnit4;
 import android.test.suitebuilder.annotation.SmallTest;
 import android.text.SpannableStringBuilder;
-import android.view.LayoutInflater;
 import android.view.View;
 import android.view.ViewGroup;
-import android.widget.ImageView;
 import android.widget.TextView;
 
 import com.android.systemui.R;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.statusbar.policy.SecurityController;
+import com.android.systemui.util.LayoutInflaterBuilder;
+import com.android.systemui.utils.TestableImageView;
+
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
-import static junit.framework.Assert.assertEquals;
-import static org.mockito.Matchers.anyBoolean;
-import static org.mockito.Matchers.anyInt;
-import static org.mockito.Matchers.anyObject;
-import static org.mockito.Matchers.eq;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.reset;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.verifyNoMoreInteractions;
-import static org.mockito.Mockito.when;
-
 @SmallTest
 @RunWith(AndroidJUnit4.class)
 public class QSFooterTest extends SysuiTestCase {
@@ -53,28 +47,26 @@ public class QSFooterTest extends SysuiTestCase {
     private final String DEVICE_OWNER_PACKAGE = "TestDPC";
     private final String VPN_PACKAGE = "TestVPN";
 
-    private ViewGroup mRootView = mock(ViewGroup.class);
-    private TextView mFooterText = mock(TextView.class);
-    private ImageView mFooterIcon = mock(ImageView.class);
-    private ImageView mFooterIcon2 = mock(ImageView.class);
+    private ViewGroup mRootView;
+    private TextView mFooterText;
+    private TestableImageView mFooterIcon;
+    private TestableImageView mFooterIcon2;
     private QSFooter mFooter;
-    private Resources mResources;
     private SecurityController mSecurityController = mock(SecurityController.class);
 
     @Before
     public void setUp() {
-        when(mRootView.findViewById(R.id.footer_text)).thenReturn(mFooterText);
-        when(mRootView.findViewById(R.id.footer_icon)).thenReturn(mFooterIcon);
-        when(mRootView.findViewById(R.id.footer_icon2)).thenReturn(mFooterIcon2);
-        final LayoutInflater layoutInflater = mock(LayoutInflater.class);
-        when(layoutInflater.inflate(eq(R.layout.quick_settings_footer), anyObject(), anyBoolean()))
-                .thenReturn(mRootView);
-        final Context context = mock(Context.class);
-        when(context.getSystemService(Context.LAYOUT_INFLATER_SERVICE)).thenReturn(layoutInflater);
-        mResources = mContext.getResources();
-        when(context.getResources()).thenReturn(mResources);
-        mFooter = new QSFooter(null, context);
-        reset(mRootView);
+        mContext.addMockSystemService(Context.LAYOUT_INFLATER_SERVICE,
+                new LayoutInflaterBuilder(mContext)
+                        .replace("ImageView", TestableImageView.class)
+                        .build());
+        Handler h = new Handler(Looper.getMainLooper());
+        h.post(() -> mFooter = new QSFooter(null, mContext));
+        waitForIdleSync(h);
+        mRootView = (ViewGroup) mFooter.getView();
+        mFooterText = (TextView) mRootView.findViewById(R.id.footer_text);
+        mFooterIcon = (TestableImageView) mRootView.findViewById(R.id.footer_icon);
+        mFooterIcon2 = (TestableImageView) mRootView.findViewById(R.id.footer_icon2);
         mFooter.setHostEnvironment(null, mSecurityController, Looper.getMainLooper());
     }
 
@@ -86,8 +78,7 @@ public class QSFooterTest extends SysuiTestCase {
         mFooter.refreshState();
 
         waitForIdleSync(mFooter.mHandler);
-        verify(mRootView).setVisibility(View.GONE);
-        verifyNoMoreInteractions(mRootView);
+        assertEquals(View.GONE, mRootView.getVisibility());
     }
 
     @Test
@@ -97,10 +88,8 @@ public class QSFooterTest extends SysuiTestCase {
         mFooter.refreshState();
 
         waitForIdleSync(mFooter.mHandler);
-        verify(mFooterText).setText(mResources.getString(R.string.do_disclosure_generic));
-        verifyNoMoreInteractions(mFooterText);
-        verify(mRootView).setVisibility(View.VISIBLE);
-        verifyNoMoreInteractions(mRootView);
+        assertEquals(mContext.getString(R.string.do_disclosure_generic), mFooterText.getText());
+        assertEquals(View.VISIBLE, mRootView.getVisibility());
     }
 
     @Test
@@ -111,11 +100,9 @@ public class QSFooterTest extends SysuiTestCase {
         mFooter.refreshState();
 
         waitForIdleSync(mFooter.mHandler);
-        verify(mFooterText).setText(mResources.getString(R.string.do_disclosure_with_name,
-                MANAGING_ORGANIZATION));
-        verifyNoMoreInteractions(mFooterText);
-        verify(mRootView).setVisibility(View.VISIBLE);
-        verifyNoMoreInteractions(mRootView);
+        assertEquals(mContext.getString(R.string.do_disclosure_with_name, MANAGING_ORGANIZATION),
+                mFooterText.getText());
+        assertEquals(View.VISIBLE, mRootView.getVisibility());
     }
 
     @Test
@@ -126,9 +113,9 @@ public class QSFooterTest extends SysuiTestCase {
         mFooter.refreshState();
 
         waitForIdleSync(mFooter.mHandler);
-        verify(mFooterIcon).setVisibility(View.VISIBLE);
-        verify(mFooterIcon).setImageResource(R.drawable.ic_qs_network_logging);
-        verify(mFooterIcon2).setVisibility(View.INVISIBLE);
+        assertEquals(View.VISIBLE, mFooterIcon.getVisibility());
+        assertEquals(R.drawable.ic_qs_network_logging, mFooterIcon.getLastImageResource());
+        assertEquals(View.INVISIBLE, mFooterIcon2.getVisibility());
     }
 
     @Test
@@ -140,9 +127,10 @@ public class QSFooterTest extends SysuiTestCase {
         mFooter.refreshState();
 
         waitForIdleSync(mFooter.mHandler);
-        verify(mFooterIcon).setVisibility(View.VISIBLE);
-        verify(mFooterIcon, never()).setImageResource(anyInt());
-        verify(mFooterIcon2).setVisibility(View.INVISIBLE);
+        assertEquals(View.VISIBLE, mFooterIcon.getVisibility());
+        // -1 == never set.
+        assertEquals(-1, mFooterIcon.getLastImageResource());
+        assertEquals(View.INVISIBLE, mFooterIcon2.getVisibility());
     }
 
     @Test
@@ -154,10 +142,11 @@ public class QSFooterTest extends SysuiTestCase {
         mFooter.refreshState();
 
         waitForIdleSync(mFooter.mHandler);
-        verify(mFooterIcon).setVisibility(View.VISIBLE);
-        verify(mFooterIcon, never()).setImageResource(anyInt());
-        verify(mFooterIcon2).setVisibility(View.VISIBLE);
-        verify(mFooterIcon2, never()).setImageResource(anyInt());
+        assertEquals(View.VISIBLE, mFooterIcon.getVisibility());
+        assertEquals(View.VISIBLE, mFooterIcon2.getVisibility());
+        // -1 == never set.
+        assertEquals(-1, mFooterIcon.getLastImageResource());
+        assertEquals(-1, mFooterIcon2.getLastImageResource());
     }
 
     @Test
@@ -211,15 +200,15 @@ public class QSFooterTest extends SysuiTestCase {
     private CharSequence getExpectedMessage(boolean hasDeviceOwnerOrganization, boolean hasVPN) {
         final SpannableStringBuilder message = new SpannableStringBuilder();
         message.append(hasDeviceOwnerOrganization ?
-                mResources.getString(R.string.monitoring_description_do_header_with_name,
+                mContext.getString(R.string.monitoring_description_do_header_with_name,
                         MANAGING_ORGANIZATION, DEVICE_OWNER_PACKAGE) :
-                mResources.getString(R.string.monitoring_description_do_header_generic,
+                mContext.getString(R.string.monitoring_description_do_header_generic,
                         DEVICE_OWNER_PACKAGE));
         message.append("\n\n");
-        message.append(mResources.getString(R.string.monitoring_description_do_body));
-        message.append(mResources.getString(
+        message.append(mContext.getString(R.string.monitoring_description_do_body));
+        message.append(mContext.getString(
                 R.string.monitoring_description_do_learn_more_separator));
-        message.append(mResources.getString(R.string.monitoring_description_do_learn_more),
+        message.append(mContext.getString(R.string.monitoring_description_do_learn_more),
                 mFooter.new EnterprisePrivacySpan(), 0);
         return message;
     }
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 350a95ff66f4..c0d5bbd35690 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/QSFragmentTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/QSFragmentTest.java
@@ -63,7 +63,7 @@ public class QSFragmentTest extends FragmentTestCase {
         when(userSwitcher.getKeyguardMonitor()).thenReturn(keyguardMonitor);
         when(userSwitcher.getUsers()).thenReturn(new ArrayList<>());
         QSTileHost host = new QSTileHost(mContext,
-                mock(PhoneStatusBar.class),
+                null,
                 getLeakChecker(BluetoothController.class),
                 getLeakChecker(LocationController.class),
                 getLeakChecker(RotationLockController.class),
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/external/TileLifecycleManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/qs/external/TileLifecycleManagerTest.java
index 782a4890ba55..66ec7dd3f270 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/external/TileLifecycleManagerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/external/TileLifecycleManagerTest.java
@@ -15,8 +15,10 @@
  */
 package com.android.systemui.qs.external;
 
+import static junit.framework.Assert.assertFalse;
 import static junit.framework.Assert.assertTrue;
 
+import static org.junit.Assert.assertEquals;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.anyInt;
 import static org.mockito.Mockito.anyString;
@@ -26,9 +28,7 @@ import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.content.ComponentName;
-import android.content.Context;
 import android.content.Intent;
-import android.content.ServiceConnection;
 import android.content.pm.PackageInfo;
 import android.content.pm.ServiceInfo;
 import android.net.Uri;
@@ -50,15 +50,12 @@ import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mockito;
-import org.mockito.invocation.InvocationOnMock;
-import org.mockito.stubbing.Answer;
 
 @SmallTest
 @RunWith(AndroidJUnit4.class)
 public class TileLifecycleManagerTest extends SysuiTestCase {
     private static final int TEST_FAIL_TIMEOUT = 5000;
 
-    private final Context mMockContext = Mockito.mock(Context.class);
     private final PackageManagerAdapter mMockPackageManagerAdapter =
             Mockito.mock(PackageManagerAdapter.class);
     private final IQSTileService.Stub mMockTileService = Mockito.mock(IQSTileService.Stub.class);
@@ -77,19 +74,7 @@ public class TileLifecycleManagerTest extends SysuiTestCase {
         // Stub.asInterface will just return itself.
         when(mMockTileService.queryLocalInterface(anyString())).thenReturn(mMockTileService);
 
-        // Default behavior for bind is success and connects as mMockTileService.
-        when(mMockContext.bindServiceAsUser(any(), any(), anyInt(), any()))
-                .thenAnswer(
-                        new Answer<Boolean>() {
-                            @Override
-                            public Boolean answer(InvocationOnMock invocation) {
-                                ServiceConnection connection =
-                                        (ServiceConnection) invocation.getArguments()[1];
-                                connection.onServiceConnected(
-                                        mTileServiceComponentName, mMockTileService);
-                                return true;
-                            }
-                        });
+        mContext.addMockService(mTileServiceComponentName, mMockTileService);
 
 
         mTileServiceIntent = new Intent().setComponent(mTileServiceComponentName);
@@ -97,7 +82,7 @@ public class TileLifecycleManagerTest extends SysuiTestCase {
         mThread = new HandlerThread("TestThread");
         mThread.start();
         mHandler = new Handler(mThread.getLooper());
-        mStateManager = new TileLifecycleManager(mHandler, mMockContext,
+        mStateManager = new TileLifecycleManager(mHandler, mContext,
                 Mockito.mock(IQSService.class), new Tile(),
                 mTileServiceIntent,
                 mUser,
@@ -126,11 +111,7 @@ public class TileLifecycleManagerTest extends SysuiTestCase {
     }
 
     private void verifyBind(int times) {
-        verify(mMockContext, times(times)).bindServiceAsUser(
-                mTileServiceIntent,
-                mStateManager,
-                Context.BIND_AUTO_CREATE | Context.BIND_FOREGROUND_SERVICE_WHILE_AWAKE,
-                mUser);
+        assertEquals(times > 0, mContext.isBound(mTileServiceComponentName));
     }
 
     @Test
@@ -143,7 +124,7 @@ public class TileLifecycleManagerTest extends SysuiTestCase {
     public void testUnbind() {
         mStateManager.setBindService(true);
         mStateManager.setBindService(false);
-        verify(mMockContext).unbindService(mStateManager);
+        assertFalse(mContext.isBound(mTileServiceComponentName));
     }
 
     @Test
@@ -203,7 +184,7 @@ public class TileLifecycleManagerTest extends SysuiTestCase {
 
         verifyBind(1);
         mStateManager.setBindService(false);
-        verify(mMockContext).unbindService(mStateManager);
+        assertFalse(mContext.isBound(mTileServiceComponentName));
         verify(mMockTileService, never()).onStartListening();
     }
 
@@ -217,7 +198,7 @@ public class TileLifecycleManagerTest extends SysuiTestCase {
 
         verifyBind(1);
         mStateManager.setBindService(false);
-        verify(mMockContext).unbindService(mStateManager);
+        assertFalse(mContext.isBound(mTileServiceComponentName));
         verify(mMockTileService, never()).onClick(null);
     }
 
@@ -233,7 +214,7 @@ public class TileLifecycleManagerTest extends SysuiTestCase {
         // Package is re-enabled.
         setPackageEnabled(true);
         mStateManager.onReceive(
-                mMockContext,
+                mContext,
                 new Intent(
                         Intent.ACTION_PACKAGE_CHANGED,
                         Uri.fromParts(
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/KeyguardIndicationControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/KeyguardIndicationControllerTest.java
index 639c8daf586b..7335af3fff9d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/KeyguardIndicationControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/KeyguardIndicationControllerTest.java
@@ -16,15 +16,18 @@
 
 package com.android.systemui.statusbar;
 
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.when;
+
 import android.app.admin.DevicePolicyManager;
 import android.app.trust.TrustManager;
-import android.content.ContentResolver;
 import android.content.Context;
-import android.content.pm.PackageManager;
-import android.content.res.Resources;
+import android.hardware.fingerprint.FingerprintManager;
 import android.os.Looper;
 import android.support.test.runner.AndroidJUnit4;
-import android.telephony.SubscriptionManager;
 import android.test.suitebuilder.annotation.SmallTest;
 import android.view.View;
 import android.view.ViewGroup;
@@ -37,22 +40,15 @@ import com.android.systemui.statusbar.phone.KeyguardIndicationTextView;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.ArgumentCaptor;
-
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.reset;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.verifyNoMoreInteractions;
-import static org.mockito.Mockito.when;
 
 @SmallTest
 @RunWith(AndroidJUnit4.class)
 public class KeyguardIndicationControllerTest extends SysuiTestCase {
 
     private final String ORGANIZATION_NAME = "organization";
-    private final String DISCLOSURE_WITH_ORGANIZATION_NAME = "managed by organization";
 
-    private Context mContext = mock(Context.class);
+    private String mDisclosureWithOrganization;
+
     private DevicePolicyManager mDevicePolicyManager = mock(DevicePolicyManager.class);
     private ViewGroup mIndicationArea = mock(ViewGroup.class);
     private KeyguardIndicationTextView mDisclosure = mock(KeyguardIndicationTextView.class);
@@ -61,19 +57,11 @@ public class KeyguardIndicationControllerTest extends SysuiTestCase {
 
     @Before
     public void setUp() throws Exception {
-        final Resources resources = mock(Resources.class);
-        when(mContext.getResources()).thenReturn(resources);
-        when(mContext.getContentResolver()).thenReturn(mock(ContentResolver.class));
-        when(mContext.getPackageManager()).thenReturn(mock(PackageManager.class));
-        when(mContext.getSystemService(Context.TELEPHONY_SUBSCRIPTION_SERVICE)).thenReturn(
-                mock(SubscriptionManager.class));
-        when(mContext.getSystemService(Context.TRUST_SERVICE)).thenReturn(
-                mock(TrustManager.class));
-        when(mContext.getSystemService(Context.DEVICE_POLICY_SERVICE)).thenReturn(
-                mDevicePolicyManager);
-
-        when(resources.getString(R.string.do_disclosure_with_name, ORGANIZATION_NAME))
-                .thenReturn(DISCLOSURE_WITH_ORGANIZATION_NAME);
+        mContext.addMockSystemService(Context.DEVICE_POLICY_SERVICE, mDevicePolicyManager);
+        mContext.addMockSystemService(Context.TRUST_SERVICE, mock(TrustManager.class));
+        mContext.addMockSystemService(Context.FINGERPRINT_SERVICE, mock(FingerprintManager.class));
+        mDisclosureWithOrganization = mContext.getString(R.string.do_disclosure_with_name,
+                ORGANIZATION_NAME);
 
         when(mIndicationArea.findViewById(R.id.keyguard_indication_enterprise_disclosure))
                 .thenReturn(mDisclosure);
@@ -113,7 +101,7 @@ public class KeyguardIndicationControllerTest extends SysuiTestCase {
         createController();
 
         verify(mDisclosure).setVisibility(View.VISIBLE);
-        verify(mDisclosure).switchIndication(DISCLOSURE_WITH_ORGANIZATION_NAME);
+        verify(mDisclosure).switchIndication(mDisclosureWithOrganization);
         verifyNoMoreInteractions(mDisclosure);
     }
 
@@ -140,7 +128,7 @@ public class KeyguardIndicationControllerTest extends SysuiTestCase {
         monitor.onKeyguardVisibilityChanged(true);
 
         verify(mDisclosure).setVisibility(View.VISIBLE);
-        verify(mDisclosure).switchIndication(DISCLOSURE_WITH_ORGANIZATION_NAME);
+        verify(mDisclosure).switchIndication(mDisclosureWithOrganization);
         verifyNoMoreInteractions(mDisclosure);
         reset(mDisclosure);
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SecurityControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SecurityControllerTest.java
index 9a697eedc2e1..87c4c664f352 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SecurityControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SecurityControllerTest.java
@@ -16,26 +16,24 @@
 
 package com.android.systemui.statusbar.policy;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
 import android.app.admin.DevicePolicyManager;
 import android.content.Context;
-import android.content.pm.UserInfo;
 import android.net.ConnectivityManager;
-import android.os.UserManager;
 import android.support.test.runner.AndroidJUnit4;
 import android.test.suitebuilder.annotation.SmallTest;
 
 import com.android.systemui.SysuiTestCase;
+
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
-import static org.mockito.Mockito.anyInt;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
-
 
 @SmallTest
 @RunWith(AndroidJUnit4.class)
@@ -45,16 +43,9 @@ public class SecurityControllerTest extends SysuiTestCase {
 
     @Before
     public void setUp() throws Exception {
-        final Context context = mock(Context.class);
-        when(context.getSystemService(Context.DEVICE_POLICY_SERVICE))
-                .thenReturn(mDevicePolicyManager);
-        when(context.getSystemService(Context.CONNECTIVITY_SERVICE))
-                .thenReturn(mock(ConnectivityManager.class));
-        final UserManager userManager = mock(UserManager.class);
-        when(userManager.getUserInfo(anyInt())).thenReturn(mock(UserInfo.class));
-        when(context.getSystemService(Context.USER_SERVICE))
-                .thenReturn(userManager);
-        mSecurityController = new SecurityControllerImpl(context);
+        mContext.addMockSystemService(Context.DEVICE_POLICY_SERVICE, mDevicePolicyManager);
+        mContext.addMockSystemService(Context.CONNECTIVITY_SERVICE, mock(ConnectivityManager.class));
+        mSecurityController = new SecurityControllerImpl(mContext);
     }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/utils/TestableContext.java b/packages/SystemUI/tests/src/com/android/systemui/utils/TestableContext.java
index bf73416ffcff..a95280641aa1 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/utils/TestableContext.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/utils/TestableContext.java
@@ -16,15 +16,19 @@ package com.android.systemui.utils;
 
 import android.content.BroadcastReceiver;
 import android.content.ComponentCallbacks;
+import android.content.ComponentName;
 import android.content.ContentProviderClient;
 import android.content.Context;
 import android.content.ContextWrapper;
 import android.content.Intent;
 import android.content.IntentFilter;
 import android.content.ServiceConnection;
+import android.content.res.Resources;
 import android.os.Handler;
+import android.os.IBinder;
 import android.os.UserHandle;
 import android.provider.Settings;
+import android.util.ArrayMap;
 
 import com.android.systemui.utils.leaks.Tracker;
 import com.android.systemui.SysuiTestCase;
@@ -34,6 +38,10 @@ public class TestableContext extends ContextWrapper {
     private final FakeContentResolver mFakeContentResolver;
     private final FakeSettingsProvider mSettingsProvider;
 
+    private ArrayMap<String, Object> mMockSystemServices;
+    private ArrayMap<ComponentName, IBinder> mMockServices;
+    private ArrayMap<ServiceConnection, ComponentName> mActiveServices;
+
     private Tracker mReceiver;
     private Tracker mService;
     private Tracker mComponent;
@@ -51,6 +59,33 @@ public class TestableContext extends ContextWrapper {
         mComponent = test.getTracker("component");
     }
 
+    @Override
+    public Resources getResources() {
+        return super.getResources();
+    }
+
+    public void addMockSystemService(String name, Object service) {
+        mMockSystemServices = lazyInit(mMockSystemServices);
+        mMockSystemServices.put(name, service);
+    }
+
+    public void addMockService(ComponentName component, IBinder service) {
+        mMockServices = lazyInit(mMockServices);
+        mMockServices.put(component, service);
+    }
+
+    private <T, V> ArrayMap<T, V> lazyInit(ArrayMap<T, V> services) {
+        return services != null ? services : new ArrayMap<T, V>();
+    }
+
+    @Override
+    public Object getSystemService(String name) {
+        if (mMockSystemServices != null && mMockSystemServices.containsKey(name)) {
+            return mMockSystemServices.get(name);
+        }
+        return super.getSystemService(name);
+    }
+
     public FakeSettingsProvider getSettingsProvider() {
         return mSettingsProvider;
     }
@@ -96,6 +131,7 @@ public class TestableContext extends ContextWrapper {
     @Override
     public boolean bindService(Intent service, ServiceConnection conn, int flags) {
         if (mService != null) mService.getLeakInfo(conn).addAllocation(new Throwable());
+        if (checkMocks(service.getComponent(), conn)) return true;
         return super.bindService(service, conn, flags);
     }
 
@@ -103,6 +139,7 @@ public class TestableContext extends ContextWrapper {
     public boolean bindServiceAsUser(Intent service, ServiceConnection conn, int flags,
             Handler handler, UserHandle user) {
         if (mService != null) mService.getLeakInfo(conn).addAllocation(new Throwable());
+        if (checkMocks(service.getComponent(), conn)) return true;
         return super.bindServiceAsUser(service, conn, flags, handler, user);
     }
 
@@ -110,15 +147,35 @@ public class TestableContext extends ContextWrapper {
     public boolean bindServiceAsUser(Intent service, ServiceConnection conn, int flags,
             UserHandle user) {
         if (mService != null) mService.getLeakInfo(conn).addAllocation(new Throwable());
+        if (checkMocks(service.getComponent(), conn)) return true;
         return super.bindServiceAsUser(service, conn, flags, user);
     }
 
+    private boolean checkMocks(ComponentName component, ServiceConnection conn) {
+        if (mMockServices != null && component != null && mMockServices.containsKey(component)) {
+            mActiveServices = lazyInit(mActiveServices);
+            mActiveServices.put(conn, component);
+            conn.onServiceConnected(component, mMockServices.get(component));
+            return true;
+        }
+        return false;
+    }
+
     @Override
     public void unbindService(ServiceConnection conn) {
         if (mService != null) mService.getLeakInfo(conn).clearAllocations();
+        if (mActiveServices != null && mActiveServices.containsKey(conn)) {
+            conn.onServiceDisconnected(mActiveServices.get(conn));
+            mActiveServices.remove(conn);
+            return;
+        }
         super.unbindService(conn);
     }
 
+    public boolean isBound(ComponentName component) {
+        return mActiveServices != null && mActiveServices.containsValue(component);
+    }
+
     @Override
     public void registerComponentCallbacks(ComponentCallbacks callback) {
         if (mComponent != null) mComponent.getLeakInfo(callback).addAllocation(new Throwable());
diff --git a/packages/SystemUI/tests/src/com/android/systemui/utils/TestableImageView.java b/packages/SystemUI/tests/src/com/android/systemui/utils/TestableImageView.java
new file mode 100644
index 000000000000..b131460e1eff
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/utils/TestableImageView.java
@@ -0,0 +1,40 @@
+/*
+ * Copyright (C) 2016 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
+ * except in compliance with the License. You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package com.android.systemui.utils;
+
+import android.annotation.DrawableRes;
+import android.annotation.Nullable;
+import android.content.Context;
+import android.util.AttributeSet;
+import android.widget.ImageView;
+
+public class TestableImageView extends ImageView {
+
+    private int mLastResId = -1;
+
+    public TestableImageView(Context context, @Nullable AttributeSet attrs) {
+        super(context, attrs);
+    }
+
+    @Override
+    public void setImageResource(@DrawableRes int resId) {
+        mLastResId = resId;
+        super.setImageResource(resId);
+    }
+
+    public int getLastImageResource() {
+        return mLastResId;
+    }
+}