From 1604cef916e8e4a5daff1ef8e61f694bcd483f12 Mon Sep 17 00:00:00 2001
From: Lou Tianhao <loutianhao@espressif.com>
Date: Wed, 17 Apr 2024 20:09:55 +0800
Subject: [PATCH] chanage(example/power_save): add wifi band (2.4 GHz and 5
 GHz) menuconfig option for esp32c5

---
 examples/wifi/power_save/main/Kconfig.projbuild | 15 +++++++++++++++
 examples/wifi/power_save/main/power_save.c      | 12 ++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/examples/wifi/power_save/main/Kconfig.projbuild b/examples/wifi/power_save/main/Kconfig.projbuild
index 21a9f4c49c..f0faf6cea7 100644
--- a/examples/wifi/power_save/main/Kconfig.projbuild
+++ b/examples/wifi/power_save/main/Kconfig.projbuild
@@ -37,6 +37,21 @@ menu "Example Configuration"
             For Station, If the station does not receive a beacon frame from the connected SoftAP during the
             inactive time, disconnect from SoftAP. Default 6s.
 
+    choice EXAMPLE_WIFI_BAND_MODE
+        prompt "wifi band"
+        default EXAMPLE_WIFI_BAND_MODE_2G
+        depends on SOC_WIFI_HE_SUPPORT_5G
+        help
+            WiFi band for the example to use.
+
+        config EXAMPLE_WIFI_BAND_MODE_2G
+            bool "2G"
+        config EXAMPLE_WIFI_BAND_MODE_5G
+            bool "5G"
+        config EXAMPLE_WIFI_BAND_MODE_AUTO
+            bool "Auto"
+    endchoice
+
     choice EXAMPLE_POWER_SAVE_MODE
         prompt "power save mode"
         default EXAMPLE_POWER_SAVE_MIN_MODEM
diff --git a/examples/wifi/power_save/main/power_save.c b/examples/wifi/power_save/main/power_save.c
index d804974d65..bb35a5c362 100644
--- a/examples/wifi/power_save/main/power_save.c
+++ b/examples/wifi/power_save/main/power_save.c
@@ -36,6 +36,15 @@
 #define DEFAULT_PS_MODE WIFI_PS_NONE
 #endif /*CONFIG_POWER_SAVE_MODEM*/
 
+#if SOC_WIFI_HE_SUPPORT_5G
+#if CONFIG_EXAMPLE_WIFI_BAND_MODE_2G
+#define DEFAULT_WIFI_BAND_MODE  WIFI_BAND_MODE_2G_ONLY
+#elif CONFIG_EXAMPLE_WIFI_BAND_MODE_5G
+#define DEFAULT_WIFI_BAND_MODE  WIFI_BAND_MODE_5G_ONLY
+#else
+#define DEFAULT_WIFI_BAND_MODE  WIFI_BAND_MODE_AUTO
+#endif
+#endif
 
 static const char *TAG = "power_save";
 
@@ -80,6 +89,9 @@ static void wifi_power_save(void)
     ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
     ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &wifi_config));
     ESP_ERROR_CHECK(esp_wifi_start());
+#if SOC_WIFI_HE_SUPPORT_5G
+    ESP_ERROR_CHECK(esp_wifi_set_band_mode(DEFAULT_WIFI_BAND_MODE));
+#endif
     ESP_ERROR_CHECK(esp_wifi_set_inactive_time(WIFI_IF_STA, DEFAULT_BEACON_TIMEOUT));
 
     ESP_LOGI(TAG, "esp_wifi_set_ps().");