diff --git a/tests/cts/net/src/android/net/cts/CaptivePortalTest.kt b/tests/cts/net/src/android/net/cts/CaptivePortalTest.kt
index 9f98e3f..a378aa7 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.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 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 @@
 @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 e17200e..05e40f9 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 @@
         }
 
         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 @@
 
         @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();