nmaitland
diff --git a/‎libraries/WiFi/src/WiFiGeneric.cpp
Lines changed: 25 additions & 16 deletions b/‎libraries/WiFi/src/WiFiGeneric.cpp
Lines changed: 25 additions & 16 deletions
diff --git a/‎libraries/WiFi/src/WiFiGeneric.h
Lines changed: 5 additions & 3 deletions b/‎libraries/WiFi/src/WiFiGeneric.h
Lines changed: 5 additions & 3 deletions
@@ -619,8 +619,13 @@ static std::vector<WiFiEventCbList_t> cbEventList;
 bool WiFiGenericClass::_persistent = true;
 bool WiFiGenericClass::_long_range = false;
 wifi_mode_t WiFiGenericClass::_forceSleepLastMode = WIFI_MODE_NULL;
+#if CONFIG_IDF_TARGET_ESP32S2
+wifi_ps_type_t WiFiGenericClass::_sleepEnabled = WIFI_PS_NONE;
+#else
+wifi_ps_type_t WiFiGenericClass::_sleepEnabled = WIFI_PS_MIN_MODEM;
+#endif
 
-WiFiGenericClass::WiFiGenericClass()
+WiFiGenericClass::WiFiGenericClass() 
 {
 
 }
@@ -789,6 +794,9 @@ esp_err_t WiFiGenericClass::_eventCallback(arduino_event_t *event)
     } else if(event->event_id == ARDUINO_EVENT_WIFI_STA_START) {
         WiFiSTAClass::_setStatus(WL_DISCONNECTED);
         setStatusBits(STA_STARTED_BIT);
+        if(esp_wifi_set_ps(_sleepEnabled) != ESP_OK){
+            log_e("esp_wifi_set_ps failed");
+        }
     } else if(event->event_id == ARDUINO_EVENT_WIFI_STA_STOP) {
         WiFiSTAClass::_setStatus(WL_NO_SHIELD);
         clearStatusBits(STA_STARTED_BIT | STA_CONNECTED_BIT | STA_HAS_IP_BIT | STA_HAS_IP6_BIT);
@@ -1053,30 +1061,31 @@ bool WiFiGenericClass::enableAP(bool enable)
  * @param enable bool
  * @return ok
  */
-bool WiFiGenericClass::setSleep(bool enable)
+bool WiFiGenericClass::setSleep(bool enabled){
+    return setSleep(enabled?WIFI_PS_MIN_MODEM:WIFI_PS_NONE);
+}
+
+bool WiFiGenericClass::setSleep(wifi_ps_type_t sleepType)
 {
-    if((getMode() & WIFI_MODE_STA) == 0){
-        log_w("STA has not been started");
-        return false;
+    if(sleepType != _sleepEnabled){
+        _sleepEnabled = sleepType;
+        if((getMode() & WIFI_MODE_STA) != 0){
+            if(esp_wifi_set_ps(_sleepEnabled) != ESP_OK){
+                log_e("esp_wifi_set_ps failed!");
+            }
+        }
+        return true;
     }
-    return esp_wifi_set_ps(enable?WIFI_PS_MIN_MODEM:WIFI_PS_NONE) == ESP_OK;
+    return false;
 }
 
 /**
  * get modem sleep enabled
  * @return true if modem sleep is enabled
  */
-bool WiFiGenericClass::getSleep()
+wifi_ps_type_t WiFiGenericClass::
8000
getSleep()
 {
-    wifi_ps_type_t ps;
-    if((getMode() & WIFI_MODE_STA) == 0){
-        log_w("STA has not been started");
-        return false;
-    }
-    if(esp_wifi_get_ps(&ps) == ESP_OK){
-        return ps == WIFI_PS_MIN_MODEM;
-    }
-    return false;
+    return _sleepEnabled;
 }
 
 /**
 
@@ -163,8 +163,9 @@ class WiFiGenericClass
     bool enableSTA(bool enable);
     bool enableAP(bool enable);
 
-    bool setSleep(bool enable);
-    bool getSleep();
+    bool setSleep(bool enabled);
+    bool setSleep(wifi_ps_type_t sleepType);
+    wifi_ps_type_t getSleep();
 
     bool setTxPower(wifi_power_t power);
     wifi_power_t getTxPower();
@@ -178,10 +179,11 @@ class WiFiGenericClass
     static bool _persistent;
     static bool _long_range;
     static wifi_mode_t _forceSleepLastMode;
+    static wifi_ps_type_t _sleepEnabled;
 
     static int setStatusBits(int bits);
     static int clearStatusBits(int bits);
-
+    
   public:
     static int hostByName(const char *aHostname, IPAddress &aResult);