Check connection before saving it

esphome/components/wifi/__init__.py

@@ -27,6 +27,7 @@ from esphome.const import (
     CONF_NETWORKS,
     CONF_ON_CONNECT,
     CONF_ON_DISCONNECT,
+    CONF_ON_ERROR,
     CONF_PASSWORD,
     CONF_POWER_SAVE_MODE,
     CONF_PRIORITY,
@@ -34,6 +35,7 @@ from esphome.const import (
     CONF_SSID,
     CONF_STATIC_IP,
     CONF_SUBNET,
+    CONF_TIMEOUT,
     CONF_TTLS_PHASE_2,
     CONF_USE_ADDRESS,
     CONF_USERNAME,
@@ -63,8 +65,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)
-WiFiSaveAPSettingsAction = wifi_ns.class_("WiFiSaveAPSettingsAction", automation.Action)
+WiFiSetSTAAction = wifi_ns.class_("WiFiSetSTAAction", automation.Action, cg.Component)
-
 
 def validate_password(value):
     value = cv.string_strict(value)
@@ -487,19 +488,35 @@ async def wifi_disable_to_code(config, action_id, template_arg, args):
 
 
 @automation.register_action(
-    "wifi.save_ap_settings",
+    "wifi.set_sta",
-    WiFiSaveAPSettingsAction,
+    WiFiSetSTAAction,
     cv.Schema(
         {
             cv.Required(CONF_SSID): cv.templatable(cv.ssid),
             cv.Required(CONF_PASSWORD): cv.templatable(validate_password),
+            cv.Optional(CONF_TIMEOUT, default="30000ms"): cv.templatable(
+                cv.positive_time_period_milliseconds
+            ),
+            cv.Optional(CONF_ON_CONNECT): automation.validate_automation(single=True),
+            cv.Optional(CONF_ON_ERROR): automation.validate_automation(single=True),
         }
     ),
 )
-async def wifi_save_settings_to_code(config, action_id, template_arg, args):
+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)
+    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_connection_timeout(timeout))
+    if on_connect_config := config.get(CONF_ON_CONNECT):
+        await automation.build_automation(
+            var.get_connect_trigger(), [], on_connect_config
+        )
+    if on_error_config := config.get(CONF_ON_ERROR):
+        await automation.build_automation(
+            var.get_error_trigger(), [], on_error_config
+        )
+    await cg.register_component(var, config)
     return var

esphome/components/wifi/wifi_component.h

@@ -209,6 +209,7 @@ class WiFiComponent : public Component {
   WiFiComponent();
 
   void set_sta(const WiFiAP &ap);
+  WiFiAP get_sta() { return this->selected_ap_; }
   void add_sta(const WiFiAP &ap);
   void clear_sta();
 
@@ -441,14 +442,77 @@ template<typename... Ts> class WiFiDisableAction : public Action<Ts...> {
   void play(Ts... x) override { global_wifi_component->disable(); }
 };
 
-template<typename... Ts> class WiFiSaveAPSettingsAction : public Action<Ts...> {
+template<typename... Ts> class WiFiSetSTAAction : public Action<Ts...>, public Component {
  public:
   TEMPLATABLE_VALUE(std::string, ssid)
   TEMPLATABLE_VALUE(std::string, password)
+  TEMPLATABLE_VALUE(uint32_t, connection_timeout)
 
   void play(Ts... x) override {
-    global_wifi_component->save_wifi_sta(this->ssid_.value(x...), this->password_.value(x...));
+    auto ssid = this->ssid_.value(x...);
+    auto password = this->password_.value(x...);
+    // Avoid multiple calls
+    if (this->connecting_)
+      return;
+    // If already connected to the same AP, do nothing
+    if (global_wifi_component->wifi_ssid() == ssid) {
+      // Callback to notify the user that the connection was successful
+      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);
+    // Save the current STA
+    this->old_sta_ = global_wifi_component->get_sta();
+    // Disable WiFi
+    global_wifi_component->disable();
+    // 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());
+    // Enable WiFi
+    global_wifi_component->enable();
+    // Set timeout for the connection
+    this->set_timeout("wifi-connect-timeout", this->connection_timeout_.value(x...), [this]() {
+      // 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;
+    });
   }
+
+  Trigger<> *get_connect_trigger() const { return this->connect_trigger_; }
+  Trigger<> *get_error_trigger() const { return this->error_trigger_; }
+
+  void loop() override {
+    if (!this->connecting_)
+      return;
+    if (global_wifi_component->is_connected()) {
+      // The WiFi is connected, stop the timeout and reset the connecting flag
+      this->cancel_timeout("wifi-connect-timeout");
+      this->connecting_ = false;
+      if (global_wifi_component->wifi_ssid() == this->new_sta_.get_ssid()) {
+        // Callback to notify the user that the connection was successful
+        this->connect_trigger_->trigger();
+      } else {
+        // Callback to notify the user that the connection failed
+        this->error_trigger_->trigger();
+      }
+    }
+  }
+ 
+ protected:
+  bool connecting_{false};
+  WiFiAP new_sta_;
+  WiFiAP old_sta_;
+  Trigger<> *connect_trigger_{new Trigger<>()};
+  Trigger<> *error_trigger_{new Trigger<>()};
 };
 
 }  // namespace wifi

tests/components/wifi/common.yaml

@@ -3,9 +3,13 @@ esphome:
     then:
       - wifi.disable
       - wifi.enable
-      - wifi.save_ap_settings:
+      - wifi.set_sta:
           ssid: MySSID
           password: password1
+          on_connect: 
+            - logger.log: "Connected to WiFi!"
+          on_error:
+            - logger.log: "Failed to connect to WiFi!"
 
 wifi:
   ssid: MySSID