Merge "Merge mainline-release 6664920 to master - DO NOT MERGE"
diff --git a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
index 0790acb..455da26 100644
--- a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
@@ -97,7 +97,9 @@
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.internal.util.ArrayUtils;
+import com.android.testutils.RecorderCallback.CallbackEntry;
 import com.android.testutils.SkipPresubmit;
+import com.android.testutils.TestableNetworkCallback;
 
 import libcore.io.Streams;
 
@@ -155,7 +157,7 @@
     private static final int MIN_NUM_NETWORK_TYPES = 1;
 
     // Airplane Mode BroadcastReceiver Timeout
-    private static final int AIRPLANE_MODE_CHANGE_TIMEOUT_MS = 5000;
+    private static final long AIRPLANE_MODE_CHANGE_TIMEOUT_MS = 10_000L;
 
     // Minimum supported keepalive counts for wifi and cellular.
     public static final int MIN_SUPPORTED_CELLULAR_KEEPALIVE_COUNT = 1;
@@ -473,6 +475,12 @@
                 .build();
     }
 
+    private NetworkRequest makeCellNetworkRequest() {
+        return new NetworkRequest.Builder()
+                .addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR)
+                .build();
+    }
+
     /**
      * Exercises both registerNetworkCallback and unregisterNetworkCallback. This checks to
      * see if we get a callback for the TRANSPORT_WIFI transport type being available.
@@ -1387,11 +1395,19 @@
     @AppModeFull(reason = "NETWORK_AIRPLANE_MODE permission can't be granted to instant apps")
     @Test
     public void testSetAirplaneMode() throws Exception{
+        final boolean supportWifi = mPackageManager.hasSystemFeature(FEATURE_WIFI);
+        final boolean supportTelephony = mPackageManager.hasSystemFeature(FEATURE_TELEPHONY);
         // store the current state of airplane mode
         final boolean isAirplaneModeEnabled = isAirplaneModeEnabled();
-
+        final TestableNetworkCallback wifiCb = new TestableNetworkCallback();
+        final TestableNetworkCallback telephonyCb = new TestableNetworkCallback();
         // disable airplane mode to reach a known state
         runShellCommand("cmd connectivity airplane-mode disable");
+        // Verify that networks are available as expected if wifi or cell is supported. Continue the
+        // test if none of them are supported since test should still able to verify the permission
+        // mechanism.
+        if (supportWifi) requestAndWaitForAvailable(makeWifiNetworkRequest(), wifiCb);
+        if (supportTelephony) requestAndWaitForAvailable(makeCellNetworkRequest(), telephonyCb);
 
         try {
             // Verify we cannot set Airplane Mode without correct permission:
@@ -1414,6 +1430,11 @@
                 fail("SecurityException should not have been thrown when setAirplaneMode(true) was"
                         + "called whilst holding the NETWORK_AIRPLANE_MODE permission.");
             }
+            // Verify that the enabling airplane mode takes effect as expected to prevent flakiness
+            // caused by fast airplane mode switches. Ensure network lost before turning off
+            // airplane mode.
+            if (supportWifi) waitForLost(wifiCb);
+            if (supportTelephony) waitForLost(telephonyCb);
 
             // Verify we can disable Airplane Mode with correct permission:
             try {
@@ -1422,8 +1443,12 @@
                 fail("SecurityException should not have been thrown when setAirplaneMode(false) was"
                         + "called whilst holding the NETWORK_AIRPLANE_MODE permission.");
             }
-
+            // Verify that turning airplane mode off takes effect as expected.
+            if (supportWifi) waitForAvailable(wifiCb);
+            if (supportTelephony) waitForAvailable(telephonyCb);
         } finally {
+            if (supportWifi) mCm.unregisterNetworkCallback(wifiCb);
+            if (supportTelephony) mCm.unregisterNetworkCallback(telephonyCb);
             // Restore the previous state of airplane mode and permissions:
             runShellCommand("cmd connectivity airplane-mode "
                     + (isAirplaneModeEnabled ? "enable" : "disable"));
@@ -1431,6 +1456,22 @@
         }
     }
 
+    private void requestAndWaitForAvailable(@NonNull final NetworkRequest request,
+            @NonNull final TestableNetworkCallback cb) {
+        mCm.registerNetworkCallback(request, cb);
+        waitForAvailable(cb);
+    }
+
+    private void waitForAvailable(@NonNull final TestableNetworkCallback cb) {
+        cb.eventuallyExpect(CallbackEntry.AVAILABLE, AIRPLANE_MODE_CHANGE_TIMEOUT_MS,
+                c -> c instanceof CallbackEntry.Available);
+    }
+
+    private void waitForLost(@NonNull final TestableNetworkCallback cb) {
+        cb.eventuallyExpect(CallbackEntry.LOST, AIRPLANE_MODE_CHANGE_TIMEOUT_MS,
+                c -> c instanceof CallbackEntry.Lost);
+    }
+
     private void setAndVerifyAirplaneMode(Boolean expectedResult)
             throws Exception {
         final CompletableFuture<Boolean> actualResult = new CompletableFuture();