From ff2458900a043fadd919d4957b3596c079581b82 Mon Sep 17 00:00:00 2001
From: Rapsssito <contact@rodrigomartin.dev>
Date: Sat, 23 Nov 2024 16:15:52 +0100
Subject: [PATCH] Rename to configure

---
 esphome/components/wifi/__init__.py      | 10 +++++++---
 esphome/components/wifi/wifi_component.h | 15 ++++++++++-----
 tests/components/wifi/common.yaml        |  2 +-
 3 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/esphome/components/wifi/__init__.py b/esphome/components/wifi/__init__.py
index dc10e07060..257d814226 100644
--- a/esphome/components/wifi/__init__.py
+++ b/esphome/components/wifi/__init__.py
@@ -48,6 +48,7 @@ from . import wpa2_eap
 AUTO_LOAD = ["network"]
 
 NO_WIFI_VARIANTS = [const.VARIANT_ESP32H2]
+CONF_SAVE = "save"
 
 wifi_ns = cg.esphome_ns.namespace("wifi")
 EAPAuth = wifi_ns.struct("EAPAuth")
@@ -65,7 +66,7 @@ WiFiConnectedCondition = wifi_ns.class_("WiFiConnectedCondition", Condition)
 WiFiEnabledCondition = wifi_ns.class_("WiFiEnabledCondition", Condition)
 WiFiEnableAction = wifi_ns.class_("WiFiEnableAction", automation.Action)
 WiFiDisableAction = wifi_ns.class_("WiFiDisableAction", automation.Action)
-WiFiSetSTAAction = wifi_ns.class_("WiFiSetSTAAction", automation.Action, cg.Component)
+WiFiConfigureAction = wifi_ns.class_("WiFiConfigureAction", automation.Action, cg.Component)
 
 
 def validate_password(value):
@@ -489,12 +490,13 @@ async def wifi_disable_to_code(config, action_id, template_arg, args):
 
 
 @automation.register_action(
-    "wifi.set_sta",
-    WiFiSetSTAAction,
+    "wifi.configure",
+    WiFiConfigureAction,
     cv.Schema(
         {
             cv.Required(CONF_SSID): cv.templatable(cv.ssid),
             cv.Required(CONF_PASSWORD): cv.templatable(validate_password),
+            cv.Optional(CONF_SAVE, default=True): cv.templatable(cv.boolean),
             cv.Optional(CONF_TIMEOUT, default="30000ms"): cv.templatable(
                 cv.positive_time_period_milliseconds
             ),
@@ -507,9 +509,11 @@ async def wifi_set_sta_to_code(config, action_id, template_arg, args):
     var = cg.new_Pvariable(action_id, template_arg)
     ssid = await cg.templatable(config[CONF_SSID], args, cg.std_string)
     password = await cg.templatable(config[CONF_PASSWORD], args, cg.std_string)
+    save = await cg.templatable(config[CONF_SAVE], args, cg.bool_)
     timeout = await cg.templatable(config.get(CONF_TIMEOUT), args, cg.uint32)
     cg.add(var.set_ssid(ssid))
     cg.add(var.set_password(password))
+    cg.add(var.set_save(save))
     cg.add(var.set_connection_timeout(timeout))
     if on_connect_config := config.get(CONF_ON_CONNECT):
         await automation.build_automation(
diff --git a/esphome/components/wifi/wifi_component.h b/esphome/components/wifi/wifi_component.h
index 22c2db5dbc..d7821a04cf 100644
--- a/esphome/components/wifi/wifi_component.h
+++ b/esphome/components/wifi/wifi_component.h
@@ -442,10 +442,11 @@ template<typename... Ts> class WiFiDisableAction : public Action<Ts...> {
   void play(Ts... x) override { global_wifi_component->disable(); }
 };
 
-template<typename... Ts> class WiFiSetSTAAction : public Action<Ts...>, public Component {
+template<typename... Ts> class WiFiConfigureAction : public Action<Ts...>, public Component {
  public:
   TEMPLATABLE_VALUE(std::string, ssid)
   TEMPLATABLE_VALUE(std::string, password)
+  TEMPLATABLE_VALUE(bool, save)
   TEMPLATABLE_VALUE(uint32_t, connection_timeout)
 
   void play(Ts... x) override {
@@ -460,8 +461,6 @@ template<typename... Ts> class WiFiSetSTAAction : public Action<Ts...>, public C
       this->connect_trigger_->trigger();
       return;
     }
-    // Set the state to connecting
-    this->connecting_ = true;
     // Create a new WiFiAP object with the new SSID and password
     this->new_sta_.set_ssid(ssid);
     this->new_sta_.set_password(password);
@@ -469,21 +468,27 @@ template<typename... Ts> class WiFiSetSTAAction : public Action<Ts...>, public C
     this->old_sta_ = global_wifi_component->get_sta();
     // Disable WiFi
     global_wifi_component->disable();
+    // Set the state to connecting
+    this->connecting_ = true;
     // Store the new STA so once the WiFi is enabled, it will connect to it
     // This is necessary because the WiFiComponent will raise an error and fallback to the saved STA
     // if trying to connect to a new STA while already connected to another one
-    global_wifi_component->save_wifi_sta(new_sta_.get_ssid(), new_sta_.get_password());
+    if (this->save_.value(x...)) {
+      global_wifi_component->save_wifi_sta(new_sta_.get_ssid(), new_sta_.get_password());
+    } else {
+      global_wifi_component->set_sta(new_sta_);
+    }
     // Enable WiFi
     global_wifi_component->enable();
     // Set timeout for the connection
     this->set_timeout("wifi-connect-timeout", this->connection_timeout_.value(x...), [this]() {
+      this->connecting_ = false;
       // If the timeout is reached, stop connecting and revert to the old AP
       global_wifi_component->disable();
       global_wifi_component->save_wifi_sta(old_sta_.get_ssid(), old_sta_.get_password());
       global_wifi_component->enable();
       // Callback to notify the user that the connection failed
       this->error_trigger_->trigger();
-      this->connecting_ = false;
     });
   }
 
diff --git a/tests/components/wifi/common.yaml b/tests/components/wifi/common.yaml
index b1346bd534..343d44b177 100644
--- a/tests/components/wifi/common.yaml
+++ b/tests/components/wifi/common.yaml
@@ -3,7 +3,7 @@ esphome:
     then:
       - wifi.disable
       - wifi.enable
-      - wifi.set_sta:
+      - wifi.configure:
           ssid: MySSID
           password: password1
           on_connect: