Merge "Fix WifiInfoTest flakiness" into jb-dev
diff --git a/tests/cts/net/Android.mk b/tests/cts/net/Android.mk
index 3e48cd2..5c70ad4 100644
--- a/tests/cts/net/Android.mk
+++ b/tests/cts/net/Android.mk
@@ -28,7 +28,7 @@
 
 LOCAL_PACKAGE_NAME := CtsNetTestCases
 
-LOCAL_STATIC_JAVA_LIBRARIES := ctstestserver
+LOCAL_STATIC_JAVA_LIBRARIES := ctstestserver ctsutil
 
 # uncomment when dalvik.annotation.Test* are removed or part of SDK
 #LOCAL_SDK_VERSION := current
diff --git a/tests/cts/net/src/android/net/wifi/cts/WifiInfoTest.java b/tests/cts/net/src/android/net/wifi/cts/WifiInfoTest.java
index df1323b..16dc57d 100644
--- a/tests/cts/net/src/android/net/wifi/cts/WifiInfoTest.java
+++ b/tests/cts/net/src/android/net/wifi/cts/WifiInfoTest.java
@@ -21,12 +21,15 @@
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
+import android.cts.util.PollingCheck;
 import android.net.wifi.SupplicantState;
 import android.net.wifi.WifiInfo;
 import android.net.wifi.WifiManager;
 import android.net.wifi.WifiManager.WifiLock;
 import android.test.AndroidTestCase;
 
+import java.util.concurrent.Callable;
+
 public class WifiInfoTest extends AndroidTestCase {
     private static class MySync {
         int expectedState = STATE_NULL;
@@ -127,10 +130,21 @@
         wifiInfo.getHiddenSSID();
         wifiInfo.getMacAddress();
         setWifiEnabled(false);
-        Thread.sleep(DURATION);
-        wifiInfo = mWifiManager.getConnectionInfo();
-        assertEquals(-1, wifiInfo.getNetworkId());
-        assertEquals(WifiManager.WIFI_STATE_DISABLED, mWifiManager.getWifiState());
+
+        PollingCheck.check("getNetworkId not -1", 20000, new Callable<Boolean>() {
+            @Override
+            public Boolean call() throws Exception {
+                WifiInfo wifiInfo = mWifiManager.getConnectionInfo();
+                return wifiInfo.getNetworkId() == -1;
+            }
+        });
+
+        PollingCheck.check("getWifiState not disabled", 20000, new Callable<Boolean>() {
+           @Override
+            public Boolean call() throws Exception {
+               return mWifiManager.getWifiState() == WifiManager.WIFI_STATE_DISABLED;
+            }
+        });
     }
 
 }