From 05c291e5a31ab20539a0e39d845222311665ad4a Mon Sep 17 00:00:00 2001
From: Chalard Jean <jchalard@google.com>
Date: Mon, 15 Nov 2021 15:38:00 +0900
Subject: [PATCH] Fix a possible flake in disconnectFromWiFi

1. There is a theoretical issue where the callback is not yet
   registered when wifi is disabled, but there is no evidence
   of it actually happening
2. 2 minutes timeout makes no sense for these tests that have a
   total 1 minute timeout anyway

Bug: 196387278
Test: CtsNetTestCases
Change-Id: I120af9b312ca34431d0e62dd85233fcdaa1b09b9
---
 .../cts/net/src/android/net/cts/CaptivePortalTest.kt |  4 +---
 .../util/java/android/net/cts/util/CtsNetUtils.java  | 12 +++++++++---
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/tests/cts/net/src/android/net/cts/CaptivePortalTest.kt b/tests/cts/net/src/android/net/cts/CaptivePortalTest.kt
index 9f98e3fab5..a378aa79f8 100644
--- a/tests/cts/net/src/android/net/cts/CaptivePortalTest.kt
+++ b/tests/cts/net/src/android/net/cts/CaptivePortalTest.kt
@@ -40,7 +40,6 @@ import android.net.cts.NetworkValidationTestUtil.setUrlExpirationDeviceConfig
 import android.net.cts.util.CtsNetUtils
 import android.net.util.NetworkStackUtils.TEST_CAPTIVE_PORTAL_HTTPS_URL
 import android.net.util.NetworkStackUtils.TEST_CAPTIVE_PORTAL_HTTP_URL
-import android.net.wifi.WifiManager
 import android.os.Build
 import android.platform.test.annotations.AppModeFull
 import android.provider.DeviceConfig
@@ -77,7 +76,7 @@ private const val TEST_PORTAL_URL_PATH = "/portal_path"
 private const val LOCALHOST_HOSTNAME = "localhost"
 
 // Re-connecting to the AP, obtaining an IP address, revalidating can take a long time
-private const val WIFI_CONNECT_TIMEOUT_MS = 120_000L
+private const val WIFI_CONNECT_TIMEOUT_MS = 40_000L
 private const val TEST_TIMEOUT_MS = 10_000L
 
 private const val TAG = "CaptivePortalTest"
@@ -94,7 +93,6 @@ private fun <T> CompletableFuture<T>.assertGet(timeoutMs: Long, message: String)
 @RunWith(AndroidJUnit4::class)
 class CaptivePortalTest {
     private val context: android.content.Context by lazy { getInstrumentation().context }
-    private val wm by lazy { context.getSystemService(WifiManager::class.java) }
     private val cm by lazy { context.getSystemService(ConnectivityManager::class.java) }
     private val pm by lazy { context.packageManager }
     private val utils by lazy { CtsNetUtils(context) }
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 e17200e6ca..05e40f9106 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
@@ -306,14 +306,18 @@ public final class CtsNetUtils {
         }
 
         try {
+            if (wasWifiConnected) {
+                // Make sure the callback is registered before turning off WiFi.
+                callback.waitForAvailable();
+            }
             SystemUtil.runShellCommand("svc wifi disable");
             if (wasWifiConnected) {
                 // Ensure we get both an onLost callback and a CONNECTIVITY_ACTION.
                 assertNotNull("Did not receive onLost callback after disabling wifi",
                         callback.waitForLost());
-            }
-            if (wasWifiConnected && expectLegacyBroadcast) {
-                assertTrue("Wifi failed to reach DISCONNECTED state.", receiver.waitForState());
+                if (expectLegacyBroadcast) {
+                    assertTrue("Wifi failed to reach DISCONNECTED state.", receiver.waitForState());
+                }
             }
         } catch (InterruptedException ex) {
             fail("disconnectFromWifi was interrupted");
@@ -599,12 +603,14 @@ public final class CtsNetUtils {
 
         @Override
         public void onAvailable(Network network) {
+            Log.i(TAG, "CtsNetUtils TestNetworkCallback onAvailable " + network);
             currentNetwork = network;
             mAvailableCv.open();
         }
 
         @Override
         public void onLost(Network network) {
+            Log.i(TAG, "CtsNetUtils TestNetworkCallback onLost " + network);
             lastLostNetwork = network;
             if (network.equals(currentNetwork)) {
                 mAvailableCv.close();