Refactor mqtt callbacks for vacuum (#118137)

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

View File

@ -8,6 +8,7 @@
from __future__ import annotations from __future__ import annotations
from collections.abc import Callable from collections.abc import Callable
from functools import partial
import logging import logging
from typing import Any, cast from typing import Any, cast
@ -49,12 +50,7 @@ from .const import (
CONF_STATE_TOPIC, CONF_STATE_TOPIC,
DOMAIN, DOMAIN,
) )
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 ReceiveMessage from .models import ReceiveMessage
from .schemas import MQTT_ENTITY_COMMON_SCHEMA from .schemas import MQTT_ENTITY_COMMON_SCHEMA
from .util import valid_publish_topic from .util import valid_publish_topic
@ -322,31 +318,32 @@ class MqttStateVacuum(MqttEntity, StateVacuumEntity):
self._attr_fan_speed = self._state_attrs.get(FAN_SPEED, 0) self._attr_fan_speed = self._state_attrs.get(FAN_SPEED, 0)
self._attr_battery_level = max(0, min(100, self._state_attrs.get(BATTERY, 0))) self._attr_battery_level = max(0, min(100, self._state_attrs.get(BATTERY, 0)))
@callback
def _state_message_received(self, msg: ReceiveMessage) -> None:
"""Handle state MQTT message."""
payload = json_loads_object(msg.payload)
if STATE in payload and (
(state := payload[STATE]) in POSSIBLE_STATES or state is None
):
self._attr_state = (
POSSIBLE_STATES[cast(str, state)] if payload[STATE] else None
)
del payload[STATE]
self._update_state_attributes(payload)
def _prepare_subscribe_topics(self) -> None: def _prepare_subscribe_topics(self) -> None:
"""(Re)Subscribe to topics.""" """(Re)Subscribe to topics."""
topics: dict[str, Any] = {} topics: dict[str, Any] = {}
@callback
@log_messages(self.hass, self.entity_id)
@write_state_on_attr_change(
self, {"_attr_battery_level", "_attr_fan_speed", "_attr_state"}
)
def state_message_received(msg: ReceiveMessage) -> None:
"""Handle state MQTT message."""
payload = json_loads_object(msg.payload)
if STATE in payload and (
(state := payload[STATE]) in POSSIBLE_STATES or state is None
):
self._attr_state = (
POSSIBLE_STATES[cast(str, state)] if payload[STATE] else None
)
del payload[STATE]
self._update_state_attributes(payload)
if state_topic := self._config.get(CONF_STATE_TOPIC): if state_topic := self._config.get(CONF_STATE_TOPIC):
topics["state_position_topic"] = { topics["state_position_topic"] = {
"topic": state_topic, "topic": state_topic,
"msg_callback": state_message_received, "msg_callback": partial(
self._message_callback,
self._state_message_received,
{"_attr_battery_level", "_attr_fan_speed", "_attr_state"},
),
"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,
} }