From eef6236bc946c43e7d9e41cd89b308b5709e1b73 Mon Sep 17 00:00:00 2001 From: Rapsssito Date: Tue, 17 Sep 2024 12:31:13 +0200 Subject: [PATCH] Check connection before saving it --- esphome/components/wifi/__init__.py | 27 ++++++++-- esphome/components/wifi/wifi_component.h | 68 +++++++++++++++++++++++- tests/components/wifi/common.yaml | 6 ++- 3 files changed, 93 insertions(+), 8 deletions(-) diff --git a/esphome/components/wifi/__init__.py b/esphome/components/wifi/__init__.py index b834fcefef..d89244395f 100644 --- a/esphome/components/wifi/__init__.py +++ b/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", - WiFiSaveAPSettingsAction, + "wifi.set_sta", + 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 diff --git a/esphome/components/wifi/wifi_component.h b/esphome/components/wifi/wifi_component.h index 8722e83297..77f4b35f37 100644 --- a/esphome/components/wifi/wifi_component.h +++ b/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 class WiFiDisableAction : public Action { void play(Ts... x) override { global_wifi_component->disable(); } }; -template class WiFiSaveAPSettingsAction : public Action { +template class WiFiSetSTAAction : public Action, 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 diff --git a/tests/components/wifi/common.yaml b/tests/components/wifi/common.yaml index 50c49bf4ac..9c4870818b 100644 --- a/tests/components/wifi/common.yaml +++ b/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