Fix WifiDialogActivityTest broken

- The ShadowWifiManager class is not working as expected

- Use Mockito class instead of Shadow class

Bug: 175369329
Test: manual test
make RunSettingsRoboTests ROBOTEST_FILTER=WifiDialogActivityTest

Change-Id: Idbc851c4d9702be23561c2c48c9d7d077e5f452e
diff --git a/src/com/android/settings/wifi/WifiDialogActivity.java b/src/com/android/settings/wifi/WifiDialogActivity.java
index a423bd8..67f291d 100644
--- a/src/com/android/settings/wifi/WifiDialogActivity.java
+++ b/src/com/android/settings/wifi/WifiDialogActivity.java
@@ -105,6 +105,7 @@
     private Intent mIntent;
     private NetworkDetailsTracker mNetworkDetailsTracker;
     private HandlerThread mWorkerThread;
+    private WifiManager mWifiManager;
 
     @Override
     protected void onCreate(Bundle savedInstanceState) {
@@ -150,22 +151,12 @@
     @Override
     protected void onStart() {
         super.onStart();
-        if (mDialog2 != null || mDialog != null) {
+        if (mDialog2 != null || mDialog != null || !hasWifiManager()) {
             return;
         }
 
         if (WizardManagerHelper.isAnySetupWizard(getIntent())) {
-            final int targetStyle = ThemeHelper.isSetupWizardDayNightEnabled(this)
-                    ? R.style.SuwAlertDialogThemeCompat_DayNight :
-                    R.style.SuwAlertDialogThemeCompat_Light;
-            if (mIsWifiTrackerLib) {
-                mDialog2 = WifiDialog2.createModal(this, this,
-                        mNetworkDetailsTracker.getWifiEntry(),
-                        WifiConfigUiBase2.MODE_CONNECT, targetStyle);
-            } else {
-                mDialog = WifiDialog.createModal(this, this, mAccessPoint,
-                        WifiConfigUiBase.MODE_CONNECT, targetStyle);
-            }
+            createDialogWithSuwTheme();
         } else {
             if (mIsWifiTrackerLib) {
                 mDialog2 = WifiDialog2.createModal(this, this,
@@ -177,11 +168,30 @@
         }
 
         if (mIsWifiTrackerLib) {
-            mDialog2.show();
-            mDialog2.setOnDismissListener(this);
+            if (mDialog2 != null) {
+                mDialog2.show();
+                mDialog2.setOnDismissListener(this);
+            }
         } else {
-            mDialog.show();
-            mDialog.setOnDismissListener(this);
+            if (mDialog != null) {
+                mDialog.show();
+                mDialog.setOnDismissListener(this);
+            }
+        }
+    }
+
+    @VisibleForTesting
+    protected void createDialogWithSuwTheme() {
+        final int targetStyle = ThemeHelper.isSetupWizardDayNightEnabled(this)
+                ? R.style.SuwAlertDialogThemeCompat_DayNight :
+                R.style.SuwAlertDialogThemeCompat_Light;
+        if (mIsWifiTrackerLib) {
+            mDialog2 = WifiDialog2.createModal(this, this,
+                    mNetworkDetailsTracker.getWifiEntry(),
+                    WifiConfigUiBase2.MODE_CONNECT, targetStyle);
+        } else {
+            mDialog = WifiDialog.createModal(this, this, mAccessPoint,
+                    WifiConfigUiBase.MODE_CONNECT, targetStyle);
         }
     }
 
@@ -221,21 +231,21 @@
 
     @Override
     public void onForget(WifiDialog dialog) {
-        final WifiManager wifiManager = getSystemService(WifiManager.class);
+        if (!hasWifiManager()) return;
         final AccessPoint accessPoint = dialog.getController().getAccessPoint();
         if (accessPoint != null) {
             if (!accessPoint.isSaved()) {
                 if (accessPoint.getNetworkInfo() != null &&
                         accessPoint.getNetworkInfo().getState() != NetworkInfo.State.DISCONNECTED) {
                     // Network is active but has no network ID - must be ephemeral.
-                    wifiManager.disableEphemeralNetwork(
+                    mWifiManager.disableEphemeralNetwork(
                             AccessPoint.convertToQuotedString(accessPoint.getSsidStr()));
                 } else {
                     // Should not happen, but a monkey seems to trigger it
                     Log.e(TAG, "Failed to forget invalid network " + accessPoint.getConfig());
                 }
             } else {
-                wifiManager.forget(accessPoint.getConfig().networkId, null /* listener */);
+                mWifiManager.forget(accessPoint.getConfig().networkId, null /* listener */);
             }
         }
 
@@ -251,6 +261,7 @@
 
     @Override
     public void onSubmit(WifiDialog2 dialog) {
+        if (!hasWifiManager()) return;
         final WifiEntry wifiEntry = dialog.getController().getWifiEntry();
         final WifiConfiguration config = dialog.getController().getConfig();
 
@@ -258,7 +269,7 @@
             if (config == null && wifiEntry != null && wifiEntry.canConnect()) {
                 wifiEntry.connect(null /* callback */);
             } else {
-                getSystemService(WifiManager.class).connect(config, null /* listener */);
+                mWifiManager.connect(config, null /* listener */);
             }
         }
 
@@ -269,22 +280,22 @@
 
     @Override
     public void onSubmit(WifiDialog dialog) {
+        if (!hasWifiManager()) return;
         final WifiConfiguration config = dialog.getController().getConfig();
         final AccessPoint accessPoint = dialog.getController().getAccessPoint();
-        final WifiManager wifiManager = getSystemService(WifiManager.class);
 
         if (getIntent().getBooleanExtra(KEY_CONNECT_FOR_CALLER, true)) {
             if (config == null) {
                 if (accessPoint != null && accessPoint.isSaved()) {
-                    wifiManager.connect(accessPoint.getConfig(), null /* listener */);
+                    mWifiManager.connect(accessPoint.getConfig(), null /* listener */);
                 }
             } else {
-                wifiManager.save(config, null /* listener */);
+                mWifiManager.save(config, null /* listener */);
                 if (accessPoint != null) {
                     // accessPoint is null for "Add network"
                     NetworkInfo networkInfo = accessPoint.getNetworkInfo();
                     if (networkInfo == null || !networkInfo.isConnected()) {
-                        wifiManager.connect(config, null /* listener */);
+                        mWifiManager.connect(config, null /* listener */);
                     }
                 }
             }
@@ -350,6 +361,12 @@
         }
     }
 
+    private boolean hasWifiManager() {
+        if (mWifiManager != null) return true;
+        mWifiManager = getSystemService(WifiManager.class);
+        return (mWifiManager != null);
+    }
+
     protected boolean hasPermissionForResult() {
         final String callingPackage = getCallingPackage();
         if (callingPackage == null) {
diff --git a/tests/robotests/src/com/android/settings/wifi/WifiDialogActivityTest.java b/tests/robotests/src/com/android/settings/wifi/WifiDialogActivityTest.java
index 9a61b6b..4ceb1e3 100644
--- a/tests/robotests/src/com/android/settings/wifi/WifiDialogActivityTest.java
+++ b/tests/robotests/src/com/android/settings/wifi/WifiDialogActivityTest.java
@@ -26,7 +26,8 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -36,41 +37,25 @@
 import android.net.wifi.WifiConfiguration;
 import android.net.wifi.WifiManager;
 
-import androidx.lifecycle.Lifecycle.State;
-import androidx.test.core.app.ActivityScenario;
-
-import com.android.settings.R;
-import com.android.settings.testutils.shadow.ShadowAlertDialogCompat;
-import com.android.settings.testutils.shadow.ShadowConnectivityManager;
-import com.android.settings.testutils.shadow.ShadowNetworkDetailsTracker;
-import com.android.settings.testutils.shadow.ShadowWifiManager;
+import com.android.settings.testutils.FakeFeatureFactory;
 import com.android.settingslib.wifi.AccessPoint;
+import com.android.wifitrackerlib.WifiEntry;
 
 import com.google.android.setupcompat.util.WizardManagerHelper;
 
-import org.junit.After;
 import org.junit.Before;
-import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.robolectric.Robolectric;
 import org.robolectric.RobolectricTestRunner;
-import org.robolectric.annotation.Config;
-import org.robolectric.util.ReflectionHelpers;
 
 @RunWith(RobolectricTestRunner.class)
-@Config(shadows = {
-        ShadowAlertDialogCompat.class,
-        ShadowConnectivityManager.class,
-        ShadowNetworkDetailsTracker.class,
-        ShadowWifiManager.class
-})
 public class WifiDialogActivityTest {
 
-    private static final String CALLING_PACKAGE = "calling_package";
-    private static final String AP1_SSID = "\"ap1\"";
+    static final String CALLING_PACKAGE = "calling_package";
+    static final int REQUEST_CODE = REQUEST_CODE_WIFI_DPP_ENROLLEE_QR_CODE_SCANNER;
 
     @Mock
     PackageManager mPackageManager;
@@ -87,13 +72,13 @@
     @Mock
     WifiConfigController2 mWifiConfiguration2;
     @Mock
+    WifiEntry mWifiEntry;
+    @Mock
     Intent mResultData;
     @Mock
-    private WifiConfigController mController;
-    @Mock
-    private WifiConfigController2 mController2;
+    WifiConfigController mController;
 
-    private ActivityScenario<WifiDialogActivity> mWifiDialogActivity;
+    WifiDialogActivity mActivity;
 
     @Before
     public void setUp() {
@@ -102,254 +87,190 @@
         when(mController.getConfig()).thenReturn(mWifiConfiguration);
         when(mController.getAccessPoint()).thenReturn(mAccessPoint);
         when(mWifiDialog2.getController()).thenReturn(mWifiConfiguration2);
+        when(mWifiConfiguration2.getWifiEntry()).thenReturn(mWifiEntry);
+        when(mWifiEntry.canConnect()).thenReturn(true);
+        FakeFeatureFactory.setupForTest();
 
-        WifiConfiguration wifiConfig = new WifiConfiguration();
-        wifiConfig.SSID = AP1_SSID;
-        doReturn(wifiConfig).when(mController).getConfig();
-        doReturn(wifiConfig).when(mController2).getConfig();
-    }
-
-    @After
-    public void cleanUp() {
-        if (mWifiDialogActivity != null) {
-            mWifiDialogActivity.close();
-        }
-    }
-
-    private ActivityScenario<WifiDialogActivity> createTargetActivity(Intent activityIntent) {
-        return ActivityScenario.launch(activityIntent);
+        mActivity = spy(Robolectric.setupActivity(WifiDialogActivity.class));
+        when(mActivity.getSystemService(WifiManager.class)).thenReturn(mWifiManager);
     }
 
     @Test
     public void onSubmit_shouldConnectToNetwork() {
-        WifiDialogActivity activity = Robolectric.setupActivity(WifiDialogActivity.class);
-        WifiDialog dialog = (WifiDialog) ShadowAlertDialogCompat.getLatestAlertDialog();
-        assertThat(dialog).isNotNull();
+        mActivity.onSubmit(mWifiDialog);
 
-        ReflectionHelpers.setField(dialog, "mController", mController);
-
-        activity.onSubmit(dialog);
-
-        assertThat(ShadowWifiManager.get().savedWifiConfig.SSID).isEqualTo(AP1_SSID);
+        verify(mWifiManager).connect(any(), any());
     }
 
     @Test
     public void onSubmit_noPermissionForResult_setResultWithoutData() {
-        WifiDialogActivity activity = spy(Robolectric.setupActivity(WifiDialogActivity.class));
-        when(activity.hasPermissionForResult()).thenReturn(false);
-        when(activity.getSystemService(WifiManager.class)).thenReturn(mWifiManager);
+        when(mActivity.hasPermissionForResult()).thenReturn(false);
 
-        activity.onSubmit(mWifiDialog);
+        mActivity.onSubmit(mWifiDialog);
 
-        verify(activity).setResult(RESULT_CONNECTED, null);
+        verify(mActivity).setResult(RESULT_CONNECTED, null);
     }
 
     @Test
     public void onSubmit_hasPermissionForResult_setResultWithData() {
-        WifiDialogActivity activity = spy(Robolectric.setupActivity(WifiDialogActivity.class));
-        when(activity.hasPermissionForResult()).thenReturn(true);
-        when(activity.createResultData(any(), any())).thenReturn(mResultData);
-        when(activity.getSystemService(WifiManager.class)).thenReturn(mWifiManager);
+        when(mActivity.hasPermissionForResult()).thenReturn(true);
+        when(mActivity.createResultData(any(), any())).thenReturn(mResultData);
 
-        activity.onSubmit(mWifiDialog);
+        mActivity.onSubmit(mWifiDialog);
 
-        verify(activity).setResult(RESULT_CONNECTED, mResultData);
+        verify(mActivity).setResult(RESULT_CONNECTED, mResultData);
     }
 
     @Test
     public void onSubmit2_noPermissionForResult_setResultWithoutData() {
-        WifiDialogActivity activity = spy(Robolectric.setupActivity(WifiDialogActivity.class));
-        when(activity.hasPermissionForResult()).thenReturn(false);
-        when(activity.getSystemService(WifiManager.class)).thenReturn(mWifiManager);
+        when(mActivity.hasPermissionForResult()).thenReturn(false);
 
-        activity.onSubmit(mWifiDialog2);
+        mActivity.onSubmit(mWifiDialog2);
 
-        verify(activity).setResult(RESULT_CONNECTED, null);
+        verify(mActivity).setResult(RESULT_CONNECTED, null);
     }
 
     @Test
     public void onSubmit2_hasPermissionForResult_setResultWithData() {
-        WifiDialogActivity activity = spy(Robolectric.setupActivity(WifiDialogActivity.class));
-        when(activity.hasPermissionForResult()).thenReturn(true);
-        when(activity.createResultData(any(), any())).thenReturn(mResultData);
-        when(activity.getSystemService(WifiManager.class)).thenReturn(mWifiManager);
+        when(mActivity.hasPermissionForResult()).thenReturn(true);
+        when(mActivity.createResultData(any(), any())).thenReturn(mResultData);
 
-        activity.onSubmit(mWifiDialog2);
+        mActivity.onSubmit(mWifiDialog2);
 
-        verify(activity).setResult(RESULT_CONNECTED, mResultData);
+        verify(mActivity).setResult(RESULT_CONNECTED, mResultData);
     }
 
     @Test
-    @Ignore
     public void onSubmit2_whenConnectForCallerIsTrue_shouldConnectToNetwork() {
         final Intent intent = new Intent("com.android.settings.WIFI_DIALOG");
         intent.putExtra(WifiDialogActivity.KEY_CHOSEN_WIFIENTRY_KEY, "FAKE_KEY");
         intent.putExtra(WifiDialogActivity.KEY_CONNECT_FOR_CALLER, true);
-        mWifiDialogActivity = createTargetActivity(intent);
+        mActivity = spy(Robolectric.buildActivity(WifiDialogActivity.class, intent).setup().get());
+        when(mActivity.getSystemService(WifiManager.class)).thenReturn(mWifiManager);
 
-        mWifiDialogActivity.moveToState(State.CREATED);
-        mWifiDialogActivity.moveToState(State.STARTED);
+        mActivity.onSubmit(mWifiDialog2);
 
-        WifiDialog2 dialog = (WifiDialog2) ShadowAlertDialogCompat.getLatestAlertDialog();
-        assertThat(dialog).isNotNull();
-
-        ReflectionHelpers.setField(dialog, "mController", mController2);
-
-        mWifiDialogActivity.onActivity(activity -> {
-            activity.onSubmit(dialog);
-            assertThat(ShadowWifiManager.get().savedWifiConfig.SSID).isEqualTo(AP1_SSID);
-        });
+        verify(mWifiEntry).connect(any());
     }
 
     @Test
     public void onSubmit_whenConnectForCallerIsFalse_shouldNotConnectToNetwork() {
-        WifiDialogActivity activity =
-                Robolectric.buildActivity(
-                        WifiDialogActivity.class,
-                        new Intent().putExtra(WifiDialogActivity.KEY_CONNECT_FOR_CALLER, false))
-                        .setup().get();
-        WifiDialog dialog = (WifiDialog) ShadowAlertDialogCompat.getLatestAlertDialog();
+        final Intent intent = new Intent();
+        intent.putExtra(WifiDialogActivity.KEY_CONNECT_FOR_CALLER, false);
+        mActivity = spy(Robolectric.buildActivity(WifiDialogActivity.class, intent).setup().get());
+        when(mActivity.getSystemService(WifiManager.class)).thenReturn(mWifiManager);
 
-        assertThat(dialog).isNotNull();
+        mActivity.onSubmit(mWifiDialog);
 
-        ReflectionHelpers.setField(dialog, "mController", mController);
-
-        activity.onSubmit(dialog);
-
-        assertThat(ShadowWifiManager.get().savedWifiConfig).isNull();
+        verify(mWifiManager, never()).connect(any(), any());
     }
 
     @Test
-    @Ignore
     public void onSubmit2_whenConnectForCallerIsFalse_shouldNotConnectToNetwork() {
         final Intent intent = new Intent("com.android.settings.WIFI_DIALOG");
         intent.putExtra(WifiDialogActivity.KEY_CHOSEN_WIFIENTRY_KEY, "FAKE_KEY");
         intent.putExtra(WifiDialogActivity.KEY_CONNECT_FOR_CALLER, false);
-        mWifiDialogActivity = createTargetActivity(intent);
+        mActivity = spy(Robolectric.buildActivity(WifiDialogActivity.class, intent).setup().get());
+        when(mActivity.getSystemService(WifiManager.class)).thenReturn(mWifiManager);
 
-        mWifiDialogActivity.moveToState(State.CREATED);
-        mWifiDialogActivity.moveToState(State.STARTED);
+        mActivity.onSubmit(mWifiDialog2);
 
-        WifiDialog2 dialog = (WifiDialog2) ShadowAlertDialogCompat.getLatestAlertDialog();
-        assertThat(dialog).isNotNull();
-
-        ReflectionHelpers.setField(dialog, "mController", mController2);
-
-        mWifiDialogActivity.onActivity(activity -> {
-            activity.onSubmit(dialog);
-            assertThat(ShadowWifiManager.get().savedWifiConfig).isEqualTo(null);
-        });
+        verify(mWifiEntry, never()).connect(any());
     }
 
     @Test
-    public void onSubmit_whenLaunchInSetupFlow_shouldBeLightThemeForWifiDialog() {
-        WifiDialogActivity activity =
-                Robolectric.buildActivity(
-                        WifiDialogActivity.class,
-                        new Intent()
-                                .putExtra(WifiDialogActivity.KEY_CONNECT_FOR_CALLER, false)
-                                .putExtra(WizardManagerHelper.EXTRA_IS_FIRST_RUN, true)
-                                .putExtra(WizardManagerHelper.EXTRA_IS_SETUP_FLOW, true))
-                        .setup().get();
-        WifiDialog dialog = (WifiDialog) ShadowAlertDialogCompat.getLatestAlertDialog();
+    public void onStart_whenLaunchInSetupFlow_shouldCreateDialogWithSuwTheme() {
+        final Intent intent = new Intent();
+        intent.putExtra(WifiDialogActivity.KEY_CONNECT_FOR_CALLER, false);
+        intent.putExtra(WizardManagerHelper.EXTRA_IS_FIRST_RUN, true);
+        intent.putExtra(WizardManagerHelper.EXTRA_IS_SETUP_FLOW, true);
+        mActivity = spy(Robolectric.buildActivity(WifiDialogActivity.class, intent).setup().get());
+        when(mActivity.getSystemService(WifiManager.class)).thenReturn(mWifiManager);
+        doNothing().when(mActivity).createDialogWithSuwTheme();
 
-        assertThat(dialog).isNotNull();
+        mActivity.onStart();
 
-        activity.onSubmit(dialog);
-
-        assertThat(dialog.getContext().getThemeResId())
-                .isEqualTo(R.style.SuwAlertDialogThemeCompat_Light);
+        verify(mActivity).createDialogWithSuwTheme();
     }
 
     @Test
     public void onActivityResult_noPermissionForResult_setResultWithoutData() {
-        WifiDialogActivity activity = spy(Robolectric.setupActivity(WifiDialogActivity.class));
-        when(activity.hasPermissionForResult()).thenReturn(false);
-        final Intent data = new Intent();
+        when(mActivity.hasPermissionForResult()).thenReturn(false);
 
-        activity.onActivityResult(REQUEST_CODE_WIFI_DPP_ENROLLEE_QR_CODE_SCANNER, RESULT_OK,
-                data);
+        mActivity.onActivityResult(REQUEST_CODE, RESULT_OK, mResultData);
 
-        verify(activity).setResult(RESULT_CONNECTED);
+        verify(mActivity).setResult(RESULT_CONNECTED);
     }
 
     @Test
     public void onActivityResult_hasPermissionForResult_setResultWithData() {
-        WifiDialogActivity activity = spy(Robolectric.setupActivity(WifiDialogActivity.class));
-        when(activity.hasPermissionForResult()).thenReturn(true);
-        final Intent data = new Intent();
+        when(mActivity.hasPermissionForResult()).thenReturn(true);
 
-        activity.onActivityResult(REQUEST_CODE_WIFI_DPP_ENROLLEE_QR_CODE_SCANNER, RESULT_OK,
-                data);
+        mActivity.onActivityResult(REQUEST_CODE, RESULT_OK, mResultData);
 
-        verify(activity).setResult(RESULT_CONNECTED, data);
+        verify(mActivity).setResult(RESULT_CONNECTED, mResultData);
     }
 
     @Test
     public void hasPermissionForResult_noCallingPackage_returnFalse() {
-        WifiDialogActivity activity = spy(Robolectric.setupActivity(WifiDialogActivity.class));
-        when(activity.getCallingPackage()).thenReturn(null);
+        when(mActivity.getCallingPackage()).thenReturn(null);
 
-        final boolean result = activity.hasPermissionForResult();
+        final boolean result = mActivity.hasPermissionForResult();
 
         assertThat(result).isFalse();
     }
 
     @Test
     public void hasPermissionForResult_noPermission_returnFalse() {
-        WifiDialogActivity activity = spy(Robolectric.setupActivity(WifiDialogActivity.class));
-        when(activity.getCallingPackage()).thenReturn(null);
+        when(mActivity.getCallingPackage()).thenReturn(null);
         when(mPackageManager.checkPermission(ACCESS_COARSE_LOCATION, CALLING_PACKAGE))
                 .thenReturn(PackageManager.PERMISSION_DENIED);
         when(mPackageManager.checkPermission(ACCESS_FINE_LOCATION, CALLING_PACKAGE))
                 .thenReturn(PackageManager.PERMISSION_DENIED);
 
-        final boolean result = activity.hasPermissionForResult();
+        final boolean result = mActivity.hasPermissionForResult();
 
         assertThat(result).isFalse();
     }
 
     @Test
     public void hasPermissionForResult_hasCoarseLocationPermission_returnFalse() {
-        WifiDialogActivity activity = spy(Robolectric.setupActivity(WifiDialogActivity.class));
-        when(activity.getCallingPackage()).thenReturn(CALLING_PACKAGE);
-        when(activity.getPackageManager()).thenReturn(mPackageManager);
+        when(mActivity.getCallingPackage()).thenReturn(CALLING_PACKAGE);
+        when(mActivity.getPackageManager()).thenReturn(mPackageManager);
         when(mPackageManager.checkPermission(ACCESS_COARSE_LOCATION, CALLING_PACKAGE))
                 .thenReturn(PackageManager.PERMISSION_GRANTED);
         when(mPackageManager.checkPermission(ACCESS_FINE_LOCATION, CALLING_PACKAGE))
                 .thenReturn(PackageManager.PERMISSION_DENIED);
 
-        final boolean result = activity.hasPermissionForResult();
+        final boolean result = mActivity.hasPermissionForResult();
 
         assertThat(result).isFalse();
     }
 
     @Test
     public void hasPermissionForResult_hasFineLocationPermission_returnTrue() {
-        WifiDialogActivity activity = spy(Robolectric.setupActivity(WifiDialogActivity.class));
-        when(activity.getCallingPackage()).thenReturn(CALLING_PACKAGE);
-        when(activity.getPackageManager()).thenReturn(mPackageManager);
+        when(mActivity.getCallingPackage()).thenReturn(CALLING_PACKAGE);
+        when(mActivity.getPackageManager()).thenReturn(mPackageManager);
         when(mPackageManager.checkPermission(ACCESS_COARSE_LOCATION, CALLING_PACKAGE))
                 .thenReturn(PackageManager.PERMISSION_DENIED);
         when(mPackageManager.checkPermission(ACCESS_FINE_LOCATION, CALLING_PACKAGE))
                 .thenReturn(PackageManager.PERMISSION_GRANTED);
 
-        final boolean result = activity.hasPermissionForResult();
+        final boolean result = mActivity.hasPermissionForResult();
 
         assertThat(result).isTrue();
     }
 
     @Test
     public void hasPermissionForResult_haveBothLocationPermissions_returnTrue() {
-        WifiDialogActivity activity = spy(Robolectric.setupActivity(WifiDialogActivity.class));
-        when(activity.getCallingPackage()).thenReturn(CALLING_PACKAGE);
-        when(activity.getPackageManager()).thenReturn(mPackageManager);
+        when(mActivity.getCallingPackage()).thenReturn(CALLING_PACKAGE);
+        when(mActivity.getPackageManager()).thenReturn(mPackageManager);
         when(mPackageManager.checkPermission(ACCESS_COARSE_LOCATION, CALLING_PACKAGE))
                 .thenReturn(PackageManager.PERMISSION_GRANTED);
         when(mPackageManager.checkPermission(ACCESS_FINE_LOCATION, CALLING_PACKAGE))
                 .thenReturn(PackageManager.PERMISSION_GRANTED);
 
-        final boolean result = activity.hasPermissionForResult();
+        final boolean result = mActivity.hasPermissionForResult();
 
         assertThat(result).isTrue();
     }