diff --git a/services/java/com/android/server/WifiService.java b/services/java/com/android/server/WifiService.java
index ab6e49c11cd1b..f3b0d1c459ac3 100644
--- a/services/java/com/android/server/WifiService.java
+++ b/services/java/com/android/server/WifiService.java
@@ -1577,8 +1577,9 @@ public class WifiService extends IWifiManager.Stub {
                  * or plugged in to AC).
                  */
                 if (!shouldWifiStayAwake(stayAwakeConditions, mPluggedType)) {
-                    if (!mWifiStateTracker.hasIpAddress()) {
-                        // do not keep Wifi awake when screen is off if Wifi is not fully active
+                    WifiInfo info = mWifiStateTracker.requestConnectionInfo();
+                    if (info.getSupplicantState() != SupplicantState.COMPLETED) {
+                        // do not keep Wifi awake when screen is off if Wifi is not associated
                         mDeviceIdle = true;
                         updateWifiState();
                     } else {