Refactor goalzero (#72398)

homeassistant/components/goalzero/__init__.py

@@ -1,70 +1,31 @@
 """The Goal Zero Yeti integration."""
 from __future__ import annotations
 
-import logging
-
 from goalzero import Yeti, exceptions
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import ATTR_MODEL, CONF_HOST, CONF_NAME, Platform
+from homeassistant.const import CONF_HOST, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
-from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
-from homeassistant.helpers.entity import DeviceInfo
-from homeassistant.helpers.update_coordinator import (
-    CoordinatorEntity,
-    DataUpdateCoordinator,
-    UpdateFailed,
-)
-
-from .const import (
-    ATTRIBUTION,
-    DATA_KEY_API,
-    DATA_KEY_COORDINATOR,
-    DOMAIN,
-    MANUFACTURER,
-    MIN_TIME_BETWEEN_UPDATES,
-)
-
-_LOGGER = logging.getLogger(__name__)
 
+from .const import DOMAIN
+from .coordinator import GoalZeroDataUpdateCoordinator
 
 PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR, Platform.SWITCH]
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Goal Zero Yeti from a config entry."""
-    name = entry.data[CONF_NAME]
+    api = Yeti(entry.data[CONF_HOST], async_get_clientsession(hass))
-    host = entry.data[CONF_HOST]
-
-    api = Yeti(host, async_get_clientsession(hass))
     try:
         await api.init_connect()
     except exceptions.ConnectError as ex:
         raise ConfigEntryNotReady(f"Failed to connect to device: {ex}") from ex
 
-    async def async_update_data() -> None:
+    coordinator = GoalZeroDataUpdateCoordinator(hass, api)
-        """Fetch data from API endpoint."""
-        try:
-            await api.get_state()
-        except exceptions.ConnectError as err:
-            raise UpdateFailed("Failed to communicate with device") from err
-
-    coordinator = DataUpdateCoordinator(
-        hass,
-        _LOGGER,
-        name=name,
-        update_method=async_update_data,
-        update_interval=MIN_TIME_BETWEEN_UPDATES,
-    )
     await coordinator.async_config_entry_first_refresh()
-    hass.data.setdefault(DOMAIN, {})
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
-    hass.data[DOMAIN][entry.entry_id] = {
-        DATA_KEY_API: api,
-        DATA_KEY_COORDINATOR: coordinator,
-    }
-
     hass.config_entries.async_setup_platforms(entry, PLATFORMS)
 
     return True
@@ -72,38 +33,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-    if unload_ok:
         hass.data[DOMAIN].pop(entry.entry_id)
     return unload_ok
-
-
-class YetiEntity(CoordinatorEntity):
-    """Representation of a Goal Zero Yeti entity."""
-
-    _attr_attribution = ATTRIBUTION
-
-    def __init__(
-        self,
-        api: Yeti,
-        coordinator: DataUpdateCoordinator,
-        name: str,
-        server_unique_id: str,
-    ) -> None:
-        """Initialize a Goal Zero Yeti entity."""
-        super().__init__(coordinator)
-        self.api = api
-        self._name = name
-        self._server_unique_id = server_unique_id
-
-    @property
-    def device_info(self) -> DeviceInfo:
-        """Return the device information of the entity."""
-        return DeviceInfo(
-            connections={(dr.CONNECTION_NETWORK_MAC, self.api.sysdata["macAddress"])},
-            identifiers={(DOMAIN, self._server_unique_id)},
-            manufacturer=MANUFACTURER,
-            model=self.api.sysdata[ATTR_MODEL],
-            name=self._name,
-            sw_version=self.api.data["firmwareVersion"],
-        )

homeassistant/components/goalzero/binary_sensor.py

@@ -3,22 +3,18 @@ from __future__ import annotations
 
 from typing import cast
 
-from goalzero import Yeti
-
 from homeassistant.components.binary_sensor import (
     BinarySensorDeviceClass,
     BinarySensorEntity,
     BinarySensorEntityDescription,
 )
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_NAME
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from . import YetiEntity
+from .const import DOMAIN
-from .const import DATA_KEY_API, DATA_KEY_COORDINATOR, DOMAIN
+from .entity import GoalZeroEntity
 
 PARALLEL_UPDATES = 0
 
@@ -51,38 +47,19 @@ async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up the Goal Zero Yeti sensor."""
-    name = entry.data[CONF_NAME]
-    goalzero_data = hass.data[DOMAIN][entry.entry_id]
     async_add_entities(
-        YetiBinarySensor(
+        GoalZeroBinarySensor(
-            goalzero_data[DATA_KEY_API],
+            hass.data[DOMAIN][entry.entry_id],
-            goalzero_data[DATA_KEY_COORDINATOR],
-            name,
             description,
-            entry.entry_id,
         )
         for description in BINARY_SENSOR_TYPES
     )
 
 
-class YetiBinarySensor(YetiEntity, BinarySensorEntity):
+class GoalZeroBinarySensor(GoalZeroEntity, BinarySensorEntity):
     """Representation of a Goal Zero Yeti sensor."""
 
-    def __init__(
-        self,
-        api: Yeti,
-        coordinator: DataUpdateCoordinator,
-        name: str,
-        description: BinarySensorEntityDescription,
-        server_unique_id: str,
-    ) -> None:
-        """Initialize a Goal Zero Yeti sensor."""
-        super().__init__(api, coordinator, name, server_unique_id)
-        self.entity_description = description
-        self._attr_name = f"{name} {description.name}"
-        self._attr_unique_id = f"{server_unique_id}/{description.key}"
-
     @property
     def is_on(self) -> bool:
         """Return True if the service is on."""
-        return cast(bool, self.api.data[self.entity_description.key] == 1)
+        return cast(bool, self._api.data[self.entity_description.key] == 1)

homeassistant/components/goalzero/const.py

@@ -1,12 +1,12 @@
 """Constants for the Goal Zero Yeti integration."""
-from datetime import timedelta
+import logging
+from typing import Final
 
 ATTRIBUTION = "Data provided by Goal Zero"
 ATTR_DEFAULT_ENABLED = "default_enabled"
 
-DATA_KEY_COORDINATOR = "coordinator"
+DOMAIN: Final = "goalzero"
-DOMAIN = "goalzero"
 DEFAULT_NAME = "Yeti"
-DATA_KEY_API = "api"
 MANUFACTURER = "Goal Zero"
-MIN_TIME_BETWEEN_UPDATES = timedelta(seconds=30)
+
+LOGGER = logging.getLogger(__name__)

homeassistant/components/goalzero/coordinator.py

@@ -0,0 +1,34 @@
+"""Data update coordinator for the Goal zero integration."""
+
+from datetime import timedelta
+
+from goalzero import Yeti, exceptions
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import DOMAIN, LOGGER
+
+
+class GoalZeroDataUpdateCoordinator(DataUpdateCoordinator):
+    """Data update coordinator for the Goal zero integration."""
+
+    config_entry: ConfigEntry
+
+    def __init__(self, hass: HomeAssistant, api: Yeti) -> None:
+        """Initialize the coordinator."""
+        super().__init__(
+            hass=hass,
+            logger=LOGGER,
+            name=DOMAIN,
+            update_interval=timedelta(seconds=30),
+        )
+        self.api = api
+
+    async def _async_update_data(self) -> None:
+        """Fetch data from API endpoint."""
+        try:
+            await self.api.get_state()
+        except exceptions.ConnectError as err:
+            raise UpdateFailed("Failed to communicate with device") from err

homeassistant/components/goalzero/entity.py

@@ -0,0 +1,48 @@
+"""Entity representing a Goal Zero Yeti device."""
+
+from goalzero import Yeti
+
+from homeassistant.const import ATTR_MODEL, CONF_NAME
+from homeassistant.helpers import device_registry as dr
+from homeassistant.helpers.entity import DeviceInfo, EntityDescription
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
+
+from .const import ATTRIBUTION, DOMAIN, MANUFACTURER
+from .coordinator import GoalZeroDataUpdateCoordinator
+
+
+class GoalZeroEntity(CoordinatorEntity[GoalZeroDataUpdateCoordinator]):
+    """Representation of a Goal Zero Yeti entity."""
+
+    _attr_attribution = ATTRIBUTION
+
+    def __init__(
+        self,
+        coordinator: GoalZeroDataUpdateCoordinator,
+        description: EntityDescription,
+    ) -> None:
+        """Initialize a Goal Zero Yeti entity."""
+        super().__init__(coordinator)
+        self.coordinator = coordinator
+        self.entity_description = description
+        self._attr_name = (
+            f"{coordinator.config_entry.data[CONF_NAME]} {description.name}"
+        )
+        self._attr_unique_id = f"{coordinator.config_entry.entry_id}/{description.key}"
+
+    @property
+    def device_info(self) -> DeviceInfo:
+        """Return the device information of the entity."""
+        return DeviceInfo(
+            connections={(dr.CONNECTION_NETWORK_MAC, self._api.sysdata["macAddress"])},
+            identifiers={(DOMAIN, self.coordinator.config_entry.entry_id)},
+            manufacturer=MANUFACTURER,
+            model=self._api.sysdata[ATTR_MODEL],
+            name=self.coordinator.config_entry.data[CONF_NAME],
+            sw_version=self._api.data["firmwareVersion"],
+        )
+
+    @property
+    def _api(self) -> Yeti:
+        """Return api from coordinator."""
+        return self.coordinator.api

homeassistant/components/goalzero/sensor.py

@@ -3,8 +3,6 @@ from __future__ import annotations
 
 from typing import cast
 
-from goalzero import Yeti
-
 from homeassistant.components.sensor import (
     SensorDeviceClass,
     SensorEntity,
@@ -13,7 +11,6 @@ from homeassistant.components.sensor import (
 )
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
-    CONF_NAME,
     ELECTRIC_CURRENT_AMPERE,
     ELECTRIC_POTENTIAL_VOLT,
     ENERGY_WATT_HOUR,
@@ -28,10 +25,9 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from . import YetiEntity
+from .const import DOMAIN
-from .const import DATA_KEY_API, DATA_KEY_COORDINATOR, DOMAIN
+from .entity import GoalZeroEntity
 
 SENSOR_TYPES: tuple[SensorEntityDescription, ...] = (
     SensorEntityDescription(
@@ -139,39 +135,19 @@ async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up the Goal Zero Yeti sensor."""
-    name = entry.data[CONF_NAME]
+    async_add_entities(
-    goalzero_data = hass.data[DOMAIN][entry.entry_id]
+        GoalZeroSensor(
-    sensors = [
+            hass.data[DOMAIN][entry.entry_id],
-        YetiSensor(
-            goalzero_data[DATA_KEY_API],
-            goalzero_data[DATA_KEY_COORDINATOR],
-            name,
             description,
-            entry.entry_id,
         )
         for description in SENSOR_TYPES
-    ]
+    )
-    async_add_entities(sensors, True)
 
 
-class YetiSensor(YetiEntity, SensorEntity):
+class GoalZeroSensor(GoalZeroEntity, SensorEntity):
     """Representation of a Goal Zero Yeti sensor."""
 
-    def __init__(
-        self,
-        api: Yeti,
-        coordinator: DataUpdateCoordinator,
-        name: str,
-        description: SensorEntityDescription,
-        server_unique_id: str,
-    ) -> None:
-        """Initialize a Goal Zero Yeti sensor."""
-        super().__init__(api, coordinator, name, server_unique_id)
-        self._attr_name = f"{name} {description.name}"
-        self.entity_description = description
-        self._attr_unique_id = f"{server_unique_id}/{description.key}"
-
     @property
     def native_value(self) -> StateType:
         """Return the state."""
-        return cast(StateType, self.api.data[self.entity_description.key])
+        return cast(StateType, self._api.data[self.entity_description.key])

homeassistant/components/goalzero/switch.py

@@ -3,17 +3,13 @@ from __future__ import annotations
 
 from typing import Any, cast
 
-from goalzero import Yeti
-
 from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_NAME
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from . import YetiEntity
+from .const import DOMAIN
-from .const import DATA_KEY_API, DATA_KEY_COORDINATOR, DOMAIN
+from .entity import GoalZeroEntity
 
 SWITCH_TYPES: tuple[SwitchEntityDescription, ...] = (
     SwitchEntityDescription(
@@ -35,50 +31,31 @@ async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up the Goal Zero Yeti switch."""
-    name = entry.data[CONF_NAME]
-    goalzero_data = hass.data[DOMAIN][entry.entry_id]
     async_add_entities(
-        YetiSwitch(
+        GoalZeroSwitch(
-            goalzero_data[DATA_KEY_API],
+            hass.data[DOMAIN][entry.entry_id],
-            goalzero_data[DATA_KEY_COORDINATOR],
-            name,
             description,
-            entry.entry_id,
         )
         for description in SWITCH_TYPES
     )
 
 
-class YetiSwitch(YetiEntity, SwitchEntity):
+class GoalZeroSwitch(GoalZeroEntity, SwitchEntity):
     """Representation of a Goal Zero Yeti switch."""
 
-    def __init__(
-        self,
-        api: Yeti,
-        coordinator: DataUpdateCoordinator,
-        name: str,
-        description: SwitchEntityDescription,
-        server_unique_id: str,
-    ) -> None:
-        """Initialize a Goal Zero Yeti switch."""
-        super().__init__(api, coordinator, name, server_unique_id)
-        self.entity_description = description
-        self._attr_name = f"{name} {description.name}"
-        self._attr_unique_id = f"{server_unique_id}/{description.key}"
-
     @property
     def is_on(self) -> bool:
         """Return state of the switch."""
-        return cast(bool, self.api.data[self.entity_description.key] == 1)
+        return cast(bool, self._api.data[self.entity_description.key] == 1)
 
     async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn off the switch."""
         payload = {self.entity_description.key: 0}
-        await self.api.post_state(payload=payload)
+        await self._api.post_state(payload=payload)
         self.coordinator.async_set_updated_data(data=payload)
 
     async def async_turn_on(self, **kwargs: Any) -> None:
         """Turn on the switch."""
         payload = {self.entity_description.key: 1}
-        await self.api.post_state(payload=payload)
+        await self._api.post_state(payload=payload)
         self.coordinator.async_set_updated_data(data=payload)