Merge "Wait for network to resume in ConcurrencyTest" am: 3ba439e1b2 am: 52697c65f3
am: d2fbf987ce

* commit 'd2fbf987ce788f87aeacf02e04b250f006ecb21c':
  Wait for network to resume in ConcurrencyTest
diff --git a/tests/cts/net/src/android/net/wifi/cts/ConcurrencyTest.java b/tests/cts/net/src/android/net/wifi/cts/ConcurrencyTest.java
index 343c1e6..a066ba8 100644
--- a/tests/cts/net/src/android/net/wifi/cts/ConcurrencyTest.java
+++ b/tests/cts/net/src/android/net/wifi/cts/ConcurrencyTest.java
@@ -20,12 +20,20 @@
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
+import android.net.ConnectivityManager;
+import android.net.ConnectivityManager.NetworkCallback;
+import android.net.Network;
+import android.net.NetworkCapabilities;
+import android.net.NetworkRequest;
 import android.net.wifi.WifiManager;
 import android.net.wifi.p2p.WifiP2pManager;
 import static android.net.wifi.p2p.WifiP2pManager.WIFI_P2P_STATE_DISABLED;
 import static android.net.wifi.p2p.WifiP2pManager.WIFI_P2P_STATE_ENABLED;
 import android.test.AndroidTestCase;
 
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
 public class ConcurrencyTest extends AndroidTestCase {
     private class MySync {
         int expectedWifiState;
@@ -94,10 +102,7 @@
         }
         mContext.unregisterReceiver(mReceiver);
 
-        if (!mWifiManager.isWifiEnabled()) {
-            assertTrue(mWifiManager.setWifiEnabled(true));
-            Thread.sleep(DURATION);
-        }
+        enableWifi();
         super.tearDown();
     }
 
@@ -114,6 +119,33 @@
         }
     }
 
+    /*
+     * Enables Wifi and block until connection is established.
+     */
+    private void enableWifi() throws InterruptedException {
+        if (!mWifiManager.isWifiEnabled()) {
+            assertTrue(mWifiManager.setWifiEnabled(true));
+        }
+
+        ConnectivityManager cm =
+            (ConnectivityManager) getContext().getSystemService(Context.CONNECTIVITY_SERVICE);
+        NetworkRequest request =
+            new NetworkRequest.Builder().addTransportType(NetworkCapabilities.TRANSPORT_WIFI)
+                                        .addCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET)
+                                        .build();
+        final CountDownLatch latch = new CountDownLatch(1);
+        NetworkCallback networkCallback = new NetworkCallback() {
+            @Override
+            public void onAvailable(Network network) {
+                latch.countDown();
+            }
+        };
+        cm.registerNetworkCallback(request, networkCallback);
+        latch.await(DURATION, TimeUnit.MILLISECONDS);
+
+        cm.unregisterNetworkCallback(networkCallback);
+    }
+
     public void testConcurrency() {
         // Cannot support p2p alone
         if (!WifiFeature.isWifiSupported(getContext())) {