Merge "Add reconnectCell util, refactor reconnectWifi" into main
diff --git a/Tethering/tests/integration/base/android/net/EthernetTetheringTestBase.java b/Tethering/tests/integration/base/android/net/EthernetTetheringTestBase.java
index c232697..2933a44 100644
--- a/Tethering/tests/integration/base/android/net/EthernetTetheringTestBase.java
+++ b/Tethering/tests/integration/base/android/net/EthernetTetheringTestBase.java
@@ -865,12 +865,12 @@
         // trigger the reselection, the total test time may over test suite 1 minmute timeout.
         // Probably need to disable/restore all internet networks in a common place of test
         // process. Currently, EthernetTetheringTest is part of CTS test which needs wifi network
-        // connection if device has wifi feature. CtsNetUtils#toggleWifi() checks wifi connection
-        // during the toggling process.
-        // See Tethering#chooseUpstreamType, CtsNetUtils#toggleWifi.
+        // connection if device has wifi feature.
+        // See Tethering#chooseUpstreamType
         // TODO: toggle cellular network if the device has no WIFI feature.
         Log.d(TAG, "Toggle WIFI to retry upstream selection");
-        sCtsNetUtils.toggleWifi();
+        sCtsNetUtils.disableWifi();
+        sCtsNetUtils.ensureWifiConnected();
 
         // Wait for expected upstream.
         final CompletableFuture<Network> future = new CompletableFuture<>();
diff --git a/tests/cts/net/api23Test/src/android/net/cts/api23test/ConnectivityManagerApi23Test.java b/tests/cts/net/api23Test/src/android/net/cts/api23test/ConnectivityManagerApi23Test.java
index 8d68c5f..af1af43 100644
--- a/tests/cts/net/api23Test/src/android/net/cts/api23test/ConnectivityManagerApi23Test.java
+++ b/tests/cts/net/api23Test/src/android/net/cts/api23test/ConnectivityManagerApi23Test.java
@@ -65,7 +65,7 @@
         }
         ConnectivityReceiver.prepare();
 
-        mCtsNetUtils.toggleWifi();
+        mCtsNetUtils.reconnectWifiAndWaitForConnectivityAction();
 
         // The connectivity broadcast has been sent; push through a terminal broadcast
         // to wait for in the receive to confirm it didn't see the connectivity change.
@@ -88,7 +88,7 @@
                 .addFlags(Intent.FLAG_ACTIVITY_NEW_TASK));
         Thread.sleep(200);
 
-        mCtsNetUtils.toggleWifi();
+        mCtsNetUtils.reconnectWifiAndWaitForConnectivityAction();
 
         Intent getConnectivityCount = new Intent(GET_WIFI_CONNECTIVITY_ACTION_COUNT);
         assertEquals(2, sendOrderedBroadcastAndReturnResultCode(
@@ -106,7 +106,7 @@
         filter.addAction(ConnectivityManager.CONNECTIVITY_ACTION);
         mContext.registerReceiver(receiver, filter);
 
-        mCtsNetUtils.toggleWifi();
+        mCtsNetUtils.reconnectWifiAndWaitForConnectivityAction();
         Intent finalIntent = new Intent(ConnectivityReceiver.FINAL_ACTION);
         finalIntent.setClass(mContext, ConnectivityReceiver.class);
         mContext.sendBroadcast(finalIntent);
diff --git a/tests/cts/net/src/android/net/cts/BatteryStatsManagerTest.java b/tests/cts/net/src/android/net/cts/BatteryStatsManagerTest.java
index 466514c..3e5d0ba 100644
--- a/tests/cts/net/src/android/net/cts/BatteryStatsManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/BatteryStatsManagerTest.java
@@ -199,7 +199,8 @@
         Log.d(TAG, "Generate traffic on wifi network.");
         generateNetworkTraffic(wifiNetwork, url);
         // Wifi battery stats are updated when wifi on.
-        mCtsNetUtils.toggleWifi();
+        mCtsNetUtils.disableWifi();
+        mCtsNetUtils.ensureWifiConnected();
 
         // Check wifi battery stats are updated.
         runAsShell(UPDATE_DEVICE_STATS,
diff --git a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
index 544f300..2646b60 100644
--- a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
@@ -1353,9 +1353,7 @@
     public void testToggleWifiConnectivityAction() throws Exception {
         assumeTrue(mPackageManager.hasSystemFeature(FEATURE_WIFI));
 
-        // toggleWifi calls connectToWifi and disconnectFromWifi, which both wait for
-        // CONNECTIVITY_ACTION broadcasts.
-        mCtsNetUtils.toggleWifi();
+        mCtsNetUtils.reconnectWifiAndWaitForConnectivityAction();
     }
 
     /** Verify restricted networks cannot be requested. */
diff --git a/tests/cts/net/util/java/android/net/cts/util/CtsNetUtils.java b/tests/cts/net/util/java/android/net/cts/util/CtsNetUtils.java
index 96330e2..3d828a4 100644
--- a/tests/cts/net/util/java/android/net/cts/util/CtsNetUtils.java
+++ b/tests/cts/net/util/java/android/net/cts/util/CtsNetUtils.java
@@ -173,21 +173,39 @@
         return cb;
     }
 
-    // Toggle WiFi twice, leaving it in the state it started in
-    public void toggleWifi() throws Exception {
-        if (mWifiManager.isWifiEnabled()) {
-            Network wifiNetwork = getWifiNetwork();
-            // Ensure system default network is WIFI because it's expected in disconnectFromWifi()
-            expectNetworkIsSystemDefault(wifiNetwork);
-            disconnectFromWifi(wifiNetwork);
-            connectToWifi();
-        } else {
-            connectToWifi();
-            Network wifiNetwork = getWifiNetwork();
-            // Ensure system default network is WIFI because it's expected in disconnectFromWifi()
-            expectNetworkIsSystemDefault(wifiNetwork);
-            disconnectFromWifi(wifiNetwork);
+    /**
+     * Toggle Wi-Fi off and on, waiting for the {@link ConnectivityManager#CONNECTIVITY_ACTION}
+     * broadcast in both cases.
+     */
+    public void reconnectWifiAndWaitForConnectivityAction() throws Exception {
+        assertTrue(mWifiManager.isWifiEnabled());
+        Network wifiNetwork = getWifiNetwork();
+        // Ensure system default network is WIFI because it's expected in disconnectFromWifi()
+        expectNetworkIsSystemDefault(wifiNetwork);
+        disconnectFromWifi(wifiNetwork, true /* expectLegacyBroadcast */);
+        connectToWifi(true /* expectLegacyBroadcast */);
+    }
+
+    /**
+     * Turn Wi-Fi off, then back on and make sure it connects, if it is supported.
+     */
+    public void reconnectWifiIfSupported() throws Exception {
+        if (!mContext.getPackageManager().hasSystemFeature(PackageManager.FEATURE_WIFI)) {
+            return;
         }
+        disableWifi();
+        ensureWifiConnected();
+    }
+
+    /**
+     * Turn cell data off, then back on and make sure it connects, if it is supported.
+     */
+    public void reconnectCellIfSupported() throws Exception {
+        if (!mContext.getPackageManager().hasSystemFeature(PackageManager.FEATURE_TELEPHONY)) {
+            return;
+        }
+        setMobileDataEnabled(false);
+        setMobileDataEnabled(true);
     }
 
     public Network expectNetworkIsSystemDefault(Network network)