From 19ccc479c31d7fe35d0f30d23329a21ac4b805a5 Mon Sep 17 00:00:00 2001
From: lbernstone <lbernstone@gmail.com>
Date: Wed, 30 Sep 2020 06:06:58 -0600
Subject: [PATCH] WIFI_PS_MAX_MODEM feature requested #3896 (#3900)

---
 libraries/WiFi/src/WiFiGeneric.cpp | 14 ++++++++++++++
 libraries/WiFi/src/WiFiGeneric.h   |  1 +
 2 files changed, 15 insertions(+)

diff --git a/libraries/WiFi/src/WiFiGeneric.cpp b/libraries/WiFi/src/WiFiGeneric.cpp
index 049a0c25..a1bdc76f 100644
--- a/libraries/WiFi/src/WiFiGeneric.cpp
+++ b/libraries/WiFi/src/WiFiGeneric.cpp
@@ -639,6 +639,20 @@ bool WiFiGenericClass::setSleep(bool enable)
     return esp_wifi_set_ps(enable?WIFI_PS_MIN_MODEM:WIFI_PS_NONE) == ESP_OK;
 }
 
+/**
+ * control modem sleep when only in STA mode
+ * @param mode wifi_ps_type_t
+ * @return ok
+ */
+bool WiFiGenericClass::setSleep(wifi_ps_type_t mode)
+{
+    if((getMode() & WIFI_MODE_STA) == 0){
+        log_w("STA has not been started");
+        return false;
+    }
+    return esp_wifi_set_ps(mode) == ESP_OK;
+}
+
 /**
  * get modem sleep enabled
  * @return true if modem sleep is enabled
diff --git a/libraries/WiFi/src/WiFiGeneric.h b/libraries/WiFi/src/WiFiGeneric.h
index 364a685d..2cd2db4a 100644
--- a/libraries/WiFi/src/WiFiGeneric.h
+++ b/libraries/WiFi/src/WiFiGeneric.h
@@ -108,6 +108,7 @@ class WiFiGenericClass
     bool enableAP(bool enable);
 
     bool setSleep(bool enable);
+    bool setSleep(wifi_ps_type_t mode);
     bool getSleep();
 
     bool setTxPower(wifi_power_t power);