diff --git a/tests/robotests/src/com/android/settings/security/SecurityFeatureProviderImplTest.java b/tests/robotests/src/com/android/settings/security/SecurityFeatureProviderImplTest.java
index 7df9af7..f64054c 100644
--- a/tests/robotests/src/com/android/settings/security/SecurityFeatureProviderImplTest.java
+++ b/tests/robotests/src/com/android/settings/security/SecurityFeatureProviderImplTest.java
@@ -16,14 +16,12 @@
 
 package com.android.settings.security;
 
+import android.content.Context;
 import android.content.IContentProvider;
 import android.content.Intent;
-import android.content.Context;
 import android.content.pm.PackageManager;
 import android.content.res.Resources;
-import android.graphics.Color;
 import android.graphics.drawable.Drawable;
-import android.graphics.drawable.PaintDrawable;
 import android.os.Bundle;
 import android.support.v7.preference.Preference;
 import android.support.v7.preference.PreferenceScreen;
@@ -44,17 +42,14 @@
 import org.robolectric.annotation.Implementation;
 import org.robolectric.annotation.Implements;
 
-import java.util.Collections;
-import java.util.List;
 import java.util.Map;
 
-import static com.google.common.truth.Truth.assertThat;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyInt;
 import static org.mockito.Matchers.anyString;
-import static org.mockito.Mockito.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
-import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
@@ -68,8 +63,8 @@
     private static final String URI_GET_SUMMARY = "content://package/text/summary";
     private static final String URI_GET_ICON = "content://package/icon/my_icon";
 
-    private static final Drawable MOCK_DRAWABLE = new PaintDrawable(Color.BLUE);
-
+    @Mock
+    private Drawable mMockDrawable;
     @Mock
     private Context mContext;
     @Mock
@@ -101,7 +96,7 @@
         mImpl = new SecurityFeatureProviderImpl();
         when(mContext.getPackageManager()).thenReturn(mPackageManager);
         when(mPackageManager.getResourcesForApplication(anyString())).thenReturn(mResources);
-        when(mResources.getDrawable(anyInt(), any())).thenReturn(MOCK_DRAWABLE);
+        when(mResources.getDrawable(anyInt(), any())).thenReturn(mMockDrawable);
     }
 
     @Test
@@ -139,9 +134,9 @@
         dashboardCategory.getTile(0).intent = new Intent().setPackage("package");
         dashboardCategory.getTile(0).metaData = bundle;
 
-        mImpl.updatePreferences(mContext, getPreferenceScreen(), dashboardCategory);
-        assertThat(screen.findPreference(MOCK_KEY).getIcon()).isEqualTo(MOCK_DRAWABLE);
-        assertThat(screen.findPreference(MOCK_KEY).getSummary()).isEqualTo(MOCK_SUMMARY);
+        mImpl.updatePreferences(mContext, screen, dashboardCategory);
+        verify(screen.findPreference(MOCK_KEY)).setIcon(mMockDrawable);
+        verify(screen.findPreference(MOCK_KEY)).setSummary(MOCK_SUMMARY);
     }
 
     @Test
@@ -154,25 +149,22 @@
         bundle.putString(TileUtils.META_DATA_PREFERENCE_SUMMARY_URI, URI_GET_SUMMARY);
 
         PreferenceScreen screen = getPreferenceScreen();
-        screen.findPreference(MOCK_KEY).setIcon(MOCK_DRAWABLE);
-        screen.findPreference(MOCK_KEY).setSummary(MOCK_SUMMARY);
+        when(screen.findPreference(MOCK_KEY).getSummary()).thenReturn(MOCK_SUMMARY);
+        when(screen.findPreference(MOCK_KEY).getIcon()).thenReturn(mMockDrawable);
 
         DashboardCategory dashboardCategory = getDashboardCategory();
         dashboardCategory.getTile(0).intent = new Intent().setPackage("package");
         dashboardCategory.getTile(0).metaData = bundle;
 
         mImpl.updatePreferences(mContext, screen, dashboardCategory);
-        verify(screen.findPreference(MOCK_KEY), times(0)).setIcon(any());
-        verify(screen.findPreference(MOCK_KEY), times(0)).setSummary(anyString());
-        assertThat(screen.findPreference(MOCK_KEY).getIcon()).isEqualTo(MOCK_DRAWABLE);
-        assertThat(screen.findPreference(MOCK_KEY).getSummary()).isEqualTo(MOCK_SUMMARY);
+        verify(screen.findPreference(MOCK_KEY), never()).setSummary(anyString());
     }
 
     private PreferenceScreen getPreferenceScreen() {
-        PreferenceScreen screen = new PreferenceScreen(mContext, null);
-        Preference preference = spy(new Preference(mContext));
-        preference.setKey(MOCK_KEY);
-        screen.addPreference(preference);
+        final PreferenceScreen screen = mock(PreferenceScreen.class);
+        final Preference pref = mock(Preference.class);
+        when(screen.findPreference(MOCK_KEY)).thenReturn(pref);
+        when(pref.getKey()).thenReturn(MOCK_KEY);
         return screen;
     }