Refactor mqtt callbacks for switch (#118127)

This commit is contained in:
Jan Bouwhuis 2024-05-25 23:17:54 +02:00 committed by GitHub
parent 6b1b15ef9b
commit fc9f7aee7e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -3,6 +3,7 @@
from __future__ import annotations from __future__ import annotations
from collections.abc import Callable from collections.abc import Callable
from functools import partial
from typing import Any from typing import Any
import voluptuous as vol import voluptuous as vol
@ -36,12 +37,7 @@ from .const import (
CONF_STATE_TOPIC, CONF_STATE_TOPIC,
PAYLOAD_NONE, PAYLOAD_NONE,
) )
from .debug_info import log_messages from .mixins import MqttEntity, async_setup_entity_entry_helper
from .mixins import (
MqttEntity,
async_setup_entity_entry_helper,
write_state_on_attr_change,
)
from .models import MqttValueTemplate, ReceiveMessage from .models import MqttValueTemplate, ReceiveMessage
from .schemas import MQTT_ENTITY_COMMON_SCHEMA from .schemas import MQTT_ENTITY_COMMON_SCHEMA
@ -118,13 +114,8 @@ class MqttSwitch(MqttEntity, SwitchEntity, RestoreEntity):
self._config.get(CONF_VALUE_TEMPLATE), entity=self self._config.get(CONF_VALUE_TEMPLATE), entity=self
).async_render_with_possible_json_value ).async_render_with_possible_json_value
def _prepare_subscribe_topics(self) -> None:
"""(Re)Subscribe to topics."""
@callback @callback
@log_messages(self.hass, self.entity_id) def _state_message_received(self, msg: ReceiveMessage) -> None:
@write_state_on_attr_change(self, {"_attr_is_on"})
def state_message_received(msg: ReceiveMessage) -> None:
"""Handle new MQTT state messages.""" """Handle new MQTT state messages."""
payload = self._value_template(msg.payload) payload = self._value_template(msg.payload)
if payload == self._state_on: if payload == self._state_on:
@ -134,17 +125,24 @@ class MqttSwitch(MqttEntity, SwitchEntity, RestoreEntity):
elif payload == PAYLOAD_NONE: elif payload == PAYLOAD_NONE:
self._attr_is_on = None self._attr_is_on = None
def _prepare_subscribe_topics(self) -> None:
"""(Re)Subscribe to topics."""
if self._config.get(CONF_STATE_TOPIC) is None: if self._config.get(CONF_STATE_TOPIC) is None:
# Force into optimistic mode. # Force into optimistic mode.
self._optimistic = True self._optimistic = True
else: return
self._sub_state = subscription.async_prepare_subscribe_topics( self._sub_state = subscription.async_prepare_subscribe_topics(
self.hass, self.hass,
self._sub_state, self._sub_state,
{ {
CONF_STATE_TOPIC: { CONF_STATE_TOPIC: {
"topic": self._config.get(CONF_STATE_TOPIC), "topic": self._config.get(CONF_STATE_TOPIC),
"msg_callback": state_message_received, "msg_callback": partial(
self._message_callback,
self._state_message_received,
{"_attr_is_on"},
),
"entity_id": self.entity_id,
"qos": self._config[CONF_QOS], "qos": self._config[CONF_QOS],
"encoding": self._config[CONF_ENCODING] or None, "encoding": self._config[CONF_ENCODING] or None,
} }