diff --git a/src/com/android/settings/wifi/WifiDialogActivity.java b/src/com/android/settings/wifi/WifiDialogActivity.java
index 4f9c7a6..e2e3d01 100644
--- a/src/com/android/settings/wifi/WifiDialogActivity.java
+++ b/src/com/android/settings/wifi/WifiDialogActivity.java
@@ -22,7 +22,9 @@
 import android.net.NetworkInfo;
 import android.net.wifi.WifiConfiguration;
 import android.net.wifi.WifiManager;
+import android.net.wifi.WifiManager.ActionListener;
 import android.os.Bundle;
+import android.support.annotation.VisibleForTesting;
 import android.util.Log;
 
 import com.android.settings.SetupWizardUtils;
@@ -40,6 +42,15 @@
     private static final String KEY_ACCESS_POINT_STATE = "access_point_state";
     private static final String KEY_WIFI_CONFIGURATION = "wifi_configuration";
 
+    /**
+     * Boolean extra indicating whether this activity should connect to an access point on the
+     * caller's behalf. If this is set to false, the caller should check
+     * {@link #KEY_WIFI_CONFIGURATION} in the result data and save that using
+     * {@link WifiManager#connect(WifiConfiguration, ActionListener)}. Default is true.
+     */
+    @VisibleForTesting
+    static final String KEY_CONNECT_FOR_CALLER = "connect_for_caller";
+
     @Override
     protected void onCreate(Bundle savedInstanceState) {
         final Intent intent = getIntent();
@@ -55,7 +66,8 @@
             accessPoint = new AccessPoint(this, accessPointState);
         }
 
-        WifiDialog dialog = WifiDialog.createModal(this, this, accessPoint, WifiConfigUiBase.MODE_CONNECT);
+        WifiDialog dialog = WifiDialog.createModal(
+                this, this, accessPoint, WifiConfigUiBase.MODE_CONNECT);
         dialog.show();
         dialog.setOnDismissListener(this);
     }
@@ -102,17 +114,19 @@
         final AccessPoint accessPoint = dialog.getController().getAccessPoint();
         final WifiManager wifiManager = getSystemService(WifiManager.class);
 
-        if (config == null) {
-            if (accessPoint != null && accessPoint.isSaved()) {
-                wifiManager.connect(accessPoint.getConfig(), null /* listener */);
-            }
-        } else {
-            wifiManager.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 */);
+        if (getIntent().getBooleanExtra(KEY_CONNECT_FOR_CALLER, true)) {
+            if (config == null) {
+                if (accessPoint != null && accessPoint.isSaved()) {
+                    wifiManager.connect(accessPoint.getConfig(), null /* listener */);
+                }
+            } else {
+                wifiManager.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 */);
+                    }
                 }
             }
         }
diff --git a/tests/robotests/src/com/android/settings/TestConfig.java b/tests/robotests/src/com/android/settings/TestConfig.java
index 3b2d2e4..d2d65a9 100644
--- a/tests/robotests/src/com/android/settings/TestConfig.java
+++ b/tests/robotests/src/com/android/settings/TestConfig.java
@@ -16,10 +16,18 @@
 
 package com.android.settings;
 
+import android.os.Build;
+
 /**
  * Constants for Robolectric config
  */
 public class TestConfig {
+
+    /**
+     * @deprecated New tests should use {@link #SDK_VERSION_O}
+     */
+    @Deprecated
     public static final int SDK_VERSION = 23;
+    public static final int SDK_VERSION_O = Build.VERSION_CODES.O;
     public static final String MANIFEST_PATH = "packages/apps/Settings/AndroidManifest.xml";
 }
diff --git a/tests/robotests/src/com/android/settings/testutils/shadow/ShadowWifiManager.java b/tests/robotests/src/com/android/settings/testutils/shadow/ShadowWifiManager.java
new file mode 100644
index 0000000..e5304da
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/testutils/shadow/ShadowWifiManager.java
@@ -0,0 +1,49 @@
+/*
+ * Copyright (C) 2017 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.settings.testutils.shadow;
+
+import static org.robolectric.RuntimeEnvironment.application;
+
+import android.net.wifi.WifiConfiguration;
+import android.net.wifi.WifiManager;
+
+import org.robolectric.annotation.HiddenApi;
+import org.robolectric.annotation.Implementation;
+import org.robolectric.annotation.Implements;
+import org.robolectric.shadow.api.Shadow;
+
+@Implements(WifiManager.class)
+public class ShadowWifiManager extends org.robolectric.shadows.ShadowWifiManager {
+
+    public WifiConfiguration savedWifiConfig;
+
+    @HiddenApi // @SystemApi
+    @Implementation
+    public void connect(WifiConfiguration config, WifiManager.ActionListener listener) {
+        savedWifiConfig = config;
+    }
+
+    @HiddenApi
+    @Implementation
+    public void save(WifiConfiguration config, WifiManager.ActionListener listener) {
+        savedWifiConfig = config;
+    }
+
+    public static ShadowWifiManager get() {
+        return Shadow.extract(application.getSystemService(WifiManager.class));
+    }
+}
diff --git a/tests/robotests/src/com/android/settings/wifi/WifiDialogActivityTest.java b/tests/robotests/src/com/android/settings/wifi/WifiDialogActivityTest.java
new file mode 100644
index 0000000..40c1478
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/wifi/WifiDialogActivityTest.java
@@ -0,0 +1,97 @@
+/*
+ * Copyright (C) 2017 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.settings.wifi;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.Mockito.doReturn;
+
+import android.content.Intent;
+import android.net.wifi.WifiConfiguration;
+
+import com.android.settings.TestConfig;
+import com.android.settings.testutils.SettingsRobolectricTestRunner;
+import com.android.settings.testutils.shadow.SettingsShadowResources;
+import com.android.settings.testutils.shadow.ShadowConnectivityManager;
+import com.android.settings.testutils.shadow.ShadowWifiManager;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.Robolectric;
+import org.robolectric.annotation.Config;
+import org.robolectric.shadows.ShadowAlertDialog;
+import org.robolectric.util.ReflectionHelpers;
+
+@RunWith(SettingsRobolectricTestRunner.class)
+@Config(
+        manifest = TestConfig.MANIFEST_PATH,
+        sdk = TestConfig.SDK_VERSION_O,
+        shadows = {
+                SettingsShadowResources.class,
+                SettingsShadowResources.SettingsShadowTheme.class,
+                ShadowConnectivityManager.class,
+                ShadowWifiManager.class
+        }
+)
+public class WifiDialogActivityTest {
+
+    public static final String AP1_SSID = "\"ap1\"";
+    @Mock
+    private WifiConfigController mController;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+
+        WifiConfiguration wifiConfig = new WifiConfiguration();
+        wifiConfig.SSID = AP1_SSID;
+        doReturn(wifiConfig).when(mController).getConfig();
+    }
+
+    @Test
+    public void onSubmit_shouldConnectToNetwork() {
+        WifiDialogActivity activity = Robolectric.setupActivity(WifiDialogActivity.class);
+        WifiDialog dialog = (WifiDialog) ShadowAlertDialog.getLatestAlertDialog();
+        assertThat(dialog).isNotNull();
+
+        ReflectionHelpers.setField(dialog, "mController", mController);
+
+        activity.onSubmit(dialog);
+
+        assertThat(ShadowWifiManager.get().savedWifiConfig.SSID).isEqualTo(AP1_SSID);
+    }
+
+    @Test
+    public void onSubmit_shouldNotConnectToNetwork_whenConnectForCallerIsFalse() {
+        WifiDialogActivity activity =
+                Robolectric.buildActivity(
+                        WifiDialogActivity.class,
+                        new Intent().putExtra(WifiDialogActivity.KEY_CONNECT_FOR_CALLER, false))
+                .setup().get();
+        WifiDialog dialog = (WifiDialog) ShadowAlertDialog.getLatestAlertDialog();
+        assertThat(dialog).isNotNull();
+
+        ReflectionHelpers.setField(dialog, "mController", mController);
+
+        activity.onSubmit(dialog);
+
+        assertThat(ShadowWifiManager.get().savedWifiConfig).isNull();
+    }
+}