Refactor MQTT to replace get_mqtt_data with HassKey (#117899)

This commit is contained in:
J. Nick Koston 2024-05-21 23:21:51 -10:00 committed by GitHub
parent b4d0562063
commit 4e3c4400a7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 74 additions and 74 deletions

View file

@ -16,8 +16,7 @@ from homeassistant.helpers.typing import DiscoveryInfoType
from homeassistant.util import dt as dt_util
from .const import ATTR_DISCOVERY_PAYLOAD, ATTR_DISCOVERY_TOPIC
from .models import MessageCallbackType, PublishPayloadType
from .util import get_mqtt_data
from .models import DATA_MQTT, MessageCallbackType, PublishPayloadType
STORED_MESSAGES = 10
@ -27,7 +26,7 @@ def log_messages(
) -> Callable[[MessageCallbackType], MessageCallbackType]:
"""Wrap an MQTT message callback to support message logging."""
debug_info_entities = get_mqtt_data(hass).debug_info_entities
debug_info_entities = hass.data[DATA_MQTT].debug_info_entities
def _log_message(msg: Any) -> None:
"""Log message."""
@ -70,7 +69,7 @@ def log_message(
retain: bool,
) -> None:
"""Log an outgoing MQTT message."""
entity_info = get_mqtt_data(hass).debug_info_entities.setdefault(
entity_info = hass.data[DATA_MQTT].debug_info_entities.setdefault(
entity_id, {"subscriptions": {}, "discovery_data": {}, "transmitted": {}}
)
if topic not in entity_info["transmitted"]:
@ -90,7 +89,7 @@ def add_subscription(
) -> None:
"""Prepare debug data for subscription."""
if entity_id := getattr(message_callback, "__entity_id", None):
entity_info = get_mqtt_data(hass).debug_info_entities.setdefault(
entity_info = hass.data[DATA_MQTT].debug_info_entities.setdefault(
entity_id, {"subscriptions": {}, "discovery_data": {}, "transmitted": {}}
)
if subscription not in entity_info["subscriptions"]:
@ -108,7 +107,7 @@ def remove_subscription(
) -> None:
"""Remove debug data for subscription if it exists."""
if (entity_id := getattr(message_callback, "__entity_id", None)) and entity_id in (
debug_info_entities := get_mqtt_data(hass).debug_info_entities
debug_info_entities := hass.data[DATA_MQTT].debug_info_entities
):
debug_info_entities[entity_id]["subscriptions"][subscription]["count"] -= 1
if not debug_info_entities[entity_id]["subscriptions"][subscription]["count"]:
@ -119,7 +118,7 @@ def add_entity_discovery_data(
hass: HomeAssistant, discovery_data: DiscoveryInfoType, entity_id: str
) -> None:
"""Add discovery data."""
entity_info = get_mqtt_data(hass).debug_info_entities.setdefault(
entity_info = hass.data[DATA_MQTT].debug_info_entities.setdefault(
entity_id, {"subscriptions": {}, "discovery_data": {}, "transmitted": {}}
)
entity_info["discovery_data"] = discovery_data
@ -129,7 +128,7 @@ def update_entity_discovery_data(
hass: HomeAssistant, discovery_payload: DiscoveryInfoType, entity_id: str
) -> None:
"""Update discovery data."""
discovery_data = get_mqtt_data(hass).debug_info_entities[entity_id][
discovery_data = hass.data[DATA_MQTT].debug_info_entities[entity_id][
"discovery_data"
]
if TYPE_CHECKING:
@ -139,7 +138,7 @@ def update_entity_discovery_data(
def remove_entity_data(hass: HomeAssistant, entity_id: str) -> None:
"""Remove discovery data."""
if entity_id in (debug_info_entities := get_mqtt_data(hass).debug_info_entities):
if entity_id in (debug_info_entities := hass.data[DATA_MQTT].debug_info_entities):
debug_info_entities.pop(entity_id)
@ -150,7 +149,7 @@ def add_trigger_discovery_data(
device_id: str,
) -> None:
"""Add discovery data."""
get_mqtt_data(hass).debug_info_triggers[discovery_hash] = {
hass.data[DATA_MQTT].debug_info_triggers[discovery_hash] = {
"device_id": device_id,
"discovery_data": discovery_data,
}
@ -162,7 +161,7 @@ def update_trigger_discovery_data(
discovery_payload: DiscoveryInfoType,
) -> None:
"""Update discovery data."""
get_mqtt_data(hass).debug_info_triggers[discovery_hash]["discovery_data"][
hass.data[DATA_MQTT].debug_info_triggers[discovery_hash]["discovery_data"][
ATTR_DISCOVERY_PAYLOAD
] = discovery_payload
@ -171,11 +170,11 @@ def remove_trigger_discovery_data(
hass: HomeAssistant, discovery_hash: tuple[str, str]
) -> None:
"""Remove discovery data."""
get_mqtt_data(hass).debug_info_triggers.pop(discovery_hash)
hass.data[DATA_MQTT].debug_info_triggers.pop(discovery_hash)
def _info_for_entity(hass: HomeAssistant, entity_id: str) -> dict[str, Any]:
entity_info = get_mqtt_data(hass).debug_info_entities[entity_id]
entity_info = hass.data[DATA_MQTT].debug_info_entities[entity_id]
monotonic_time_diff = time.time() - time.monotonic()
subscriptions = [
{
@ -231,7 +230,7 @@ def _info_for_entity(hass: HomeAssistant, entity_id: str) -> dict[str, Any]:
def _info_for_trigger(
hass: HomeAssistant, trigger_key: tuple[str, str]
) -> dict[str, Any]:
trigger = get_mqtt_data(hass).debug_info_triggers[trigger_key]
trigger = hass.data[DATA_MQTT].debug_info_triggers[trigger_key]
discovery_data = None
if trigger["discovery_data"] is not None:
discovery_data = {
@ -244,7 +243,7 @@ def _info_for_trigger(
def info_for_config_entry(hass: HomeAssistant) -> dict[str, list[Any]]:
"""Get debug info for all entities and triggers."""
mqtt_data = get_mqtt_data(hass)
mqtt_data = hass.data[DATA_MQTT]
mqtt_info: dict[str, list[Any]] = {"entities": [], "triggers": []}
mqtt_info["entities"].extend(
@ -262,7 +261,7 @@ def info_for_config_entry(hass: HomeAssistant) -> dict[str, list[Any]]:
def info_for_device(hass: HomeAssistant, device_id: str) -> dict[str, list[Any]]:
"""Get debug info for a device."""
mqtt_data = get_mqtt_data(hass)
mqtt_data = hass.data[DATA_MQTT]
mqtt_info: dict[str, list[Any]] = {"entities": [], "triggers": []}
entity_registry = er.async_get(hass)