From 72b851928c2f457020608cb6cb629d2dc16fac2c Mon Sep 17 00:00:00 2001
From: Motomu Utsumi <motomuman@google.com>
Date: Thu, 27 Mar 2025 16:38:32 -0700
Subject: [PATCH] Wait for Wi-Fi validation after reconnect

RestoreDefaultNetwork has been flaky due to delays in Wi-Fi validation
after reconnecting (sometimes exceeding 30 seconds).
RestoreDefaultNetwork fails at the end of all tests. So connectivity
diagnostics does not contain diagnostics for the network that fails to
be validated.
This CL updates some tests to fail if the Wi-Fi network is not
validated within 30 seconds after the reconnect. This ensures that
connectivity diagnostics include the diagnostics for the problematic
network.

Flag: TEST_ONLY
Bug: 361401328
Test: ConnectivityManagerTest
(cherry picked from https://android-review.googlesource.com/q/commit:cd1bd59a1a18b04e71c3137c4e7d506757691c6c)
Merged-In: Ifeadc8b030f0073133df4fca838e64f1b37e735c

Change-Id: Ifeadc8b030f0073133df4fca838e64f1b37e735c
---
 .../net/cts/ConnectivityManagerTest.java      | 37 ++++++++++++-------
 1 file changed, 23 insertions(+), 14 deletions(-)

diff --git a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
index 207b2e5206..3a3b256fc9 100644
--- a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
@@ -484,7 +484,10 @@ public class ConnectivityManagerTest {
             final TestableNetworkCallback callback =
                     networkCallbackRule.registerDefaultNetworkCallback();
             assertNotNull("Couldn't restore Internet connectivity",
-                    callback.eventuallyExpect(CallbackEntry.AVAILABLE));
+                    callback.eventuallyExpect(CallbackEntry.NETWORK_CAPS_UPDATED,
+                            NETWORK_CALLBACK_TIMEOUT_MS,
+                            entry -> ((CallbackEntry.CapabilitiesChanged) entry)
+                                    .getCaps().hasCapability(NET_CAPABILITY_VALIDATED)));
         });
     }
 
@@ -1172,6 +1175,7 @@ public class ConnectivityManagerTest {
             mCtsNetUtils.restorePrivateDnsSetting();
             // Toggle networks to make sure they are re-validated
             mCtsNetUtils.reconnectWifiIfSupported();
+            ensureWifiIsValidatedIfSupported();
             mCtsNetUtils.reconnectCellIfSupported();
         }
     }
@@ -1418,7 +1422,7 @@ public class ConnectivityManagerTest {
             if (firstIntent != null) mCm.unregisterNetworkCallback(firstIntent);
             if (secondIntent != null) mCm.unregisterNetworkCallback(secondIntent);
             if (receiver != null) mContext.unregisterReceiver(receiver);
-            mCtsNetUtils.ensureWifiConnected();
+            ensureWifiIsValidatedIfSupported();
         }
     }
 
@@ -1516,6 +1520,7 @@ public class ConnectivityManagerTest {
         assumeTrue(mPackageManager.hasSystemFeature(FEATURE_WIFI));
 
         mCtsNetUtils.reconnectWifiAndWaitForConnectivityAction();
+        ensureWifiIsValidatedIfSupported();
     }
 
     /** Verify restricted networks cannot be requested. */
@@ -2407,6 +2412,7 @@ public class ConnectivityManagerTest {
             // Restore the previous state of airplane mode and permissions:
             runShellCommand("cmd connectivity airplane-mode "
                     + (isAirplaneModeEnabled ? "enable" : "disable"));
+            ensureWifiIsValidatedIfSupported();
         }
     }
 
@@ -2836,7 +2842,7 @@ public class ConnectivityManagerTest {
             ConnectivitySettingsManager.setPrivateDnsMode(mContext, curPrivateDnsMode);
             tetherUtils.unregisterTetheringEventCallback(tetherEventCallback);
             tetherUtils.stopAllTethering();
-            mCtsNetUtils.ensureWifiConnected();
+            ensureWifiIsValidatedIfSupported();
         }
     }
 
@@ -3081,6 +3087,12 @@ public class ConnectivityManagerTest {
         }
     }
 
+    private void ensureWifiIsValidatedIfSupported() {
+        if (mPackageManager.hasSystemFeature(FEATURE_WIFI)) {
+            new ConnectUtil(mContext).ensureWifiValidated();
+        }
+    }
+
     @AppModeFull(reason = "WRITE_DEVICE_CONFIG permission can't be granted to instant apps")
     @Test
     public void testAcceptPartialConnectivity_validatedNetwork() throws Exception {
@@ -3102,7 +3114,7 @@ public class ConnectivityManagerTest {
             expectNetworkHasCapability(network, NET_CAPABILITY_VALIDATED, WIFI_CONNECT_TIMEOUT_MS);
         } finally {
             mHttpServer.stop();
-            mTestValidationConfigRule.runAfterNextCleanup(this::reconnectWifi);
+            mTestValidationConfigRule.runAfterNextCleanup(this::reconnectWifiAndEnsureValidated);
         }
     }
 
@@ -3132,7 +3144,7 @@ public class ConnectivityManagerTest {
             // apply here. Thus, turn off wifi first and restart to restore.
             mTestValidationConfigRule.runAfterNextCleanup(() -> {
                 mCtsNetUtils.disableWifi();
-                mCtsNetUtils.ensureWifiConnected();
+                ensureWifiIsValidatedIfSupported();
             });
         }
     }
@@ -3178,7 +3190,7 @@ public class ConnectivityManagerTest {
             // apply here. Thus, turn off wifi first and restart to restore.
             mTestValidationConfigRule.runAfterNextCleanup(() -> {
                 mCtsNetUtils.disableWifi();
-                mCtsNetUtils.ensureWifiConnected();
+                ensureWifiIsValidatedIfSupported();
             });
         }
     }
@@ -3242,7 +3254,7 @@ public class ConnectivityManagerTest {
         } finally {
             resetAvoidBadWifi(previousAvoidBadWifi);
             mHttpServer.stop();
-            mTestValidationConfigRule.runAfterNextCleanup(this::reconnectWifi);
+            mTestValidationConfigRule.runAfterNextCleanup(this::reconnectWifiAndEnsureValidated);
         }
     }
 
@@ -3291,9 +3303,9 @@ public class ConnectivityManagerTest {
         mHttpServer.start();
     }
 
-    private Network reconnectWifi() {
+    private Network reconnectWifiAndEnsureValidated() {
         mCtsNetUtils.ensureWifiDisconnected(null /* wifiNetworkToCheck */);
-        return mCtsNetUtils.ensureWifiConnected();
+        return new ConnectUtil(mContext).ensureWifiValidated();
     }
 
     private Network prepareValidatedNetwork() throws Exception {
@@ -3307,10 +3319,7 @@ public class ConnectivityManagerTest {
         prepareHttpServer();
         configTestServer(Status.NO_CONTENT, Status.NO_CONTENT);
         // Disconnect wifi first then start wifi network with configuration.
-        final Network wifiNetwork = reconnectWifi();
-
-        return expectNetworkHasCapability(wifiNetwork, NET_CAPABILITY_VALIDATED,
-                WIFI_CONNECT_TIMEOUT_MS);
+        return reconnectWifiAndEnsureValidated();
     }
 
     private Network preparePartialConnectivity() throws Exception {
@@ -3387,7 +3396,7 @@ public class ConnectivityManagerTest {
         testAndCleanup(() -> {
             // New default network connected will trigger a network activity notification.
             if (supportWifi) {
-                mCtsNetUtils.ensureWifiConnected();
+                ensureWifiIsValidatedIfSupported();
             } else {
                 networkCallbackRule.requestCell();
             }