Version sensor entity cleanup (#53915)

Co-authored-by: Franck Nijhof <git@frenck.dev>

homeassistant/components/version/sensor.py

@@ -2,11 +2,19 @@
 from datetime import timedelta
 import logging
 
-from pyhaversion import HaVersion, HaVersionChannel, HaVersionSource
+from pyhaversion import (
-from pyhaversion.exceptions import HaVersionFetchException, HaVersionParseException
+    HaVersion,
+    HaVersionChannel,
+    HaVersionSource,
+    exceptions as pyhaversionexceptions,
+)
 import voluptuous as vol
 
-from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity
+from homeassistant.components.sensor import (
+    PLATFORM_SCHEMA,
+    SensorEntity,
+    SensorEntityDescription,
+)
 from homeassistant.const import CONF_NAME, CONF_SOURCE
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 import homeassistant.helpers.config_validation as cv
@@ -30,12 +38,10 @@ ALL_IMAGES = [
     "raspberrypi4",
     "tinker",
 ]
-ALL_SOURCES = [
+
-    "container",
+HA_VERSION_SOURCES = [source.value for source in HaVersionSource]
-    "haio",
+
-    "local",
+ALL_SOURCES = HA_VERSION_SOURCES + [
-    "pypi",
-    "supervisor",
     "hassio",  # Kept to not break existing configurations
     "docker",  # Kept to not break existing configurations
 ]
@@ -48,8 +54,6 @@ DEFAULT_NAME_LATEST = "Latest Version"
 DEFAULT_NAME_LOCAL = "Current Version"
 DEFAULT_SOURCE = "local"
 
-ICON = "mdi:package-up"
-
 TIME_BETWEEN_UPDATES = timedelta(minutes=5)
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
@@ -72,40 +76,42 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
     name = config.get(CONF_NAME)
     source = config.get(CONF_SOURCE)
 
+    channel = HaVersionChannel.BETA if beta else HaVersionChannel.STABLE
     session = async_get_clientsession(hass)
 
-    channel = HaVersionChannel.BETA if beta else HaVersionChannel.STABLE
+    if source in HA_VERSION_SOURCES:
+        source = HaVersionSource(source)
+    elif source == "hassio":
+        source = HaVersionSource.SUPERVISOR
+    elif source == "docker":
+        source = HaVersionSource.CONTAINER
 
-    if source == "pypi":
+    if (
-        haversion = VersionData(
+        source in (HaVersionSource.SUPERVISOR, HaVersionSource.CONTAINER)
-            HaVersion(session, source=HaVersionSource.PYPI, channel=channel)
+        and image is not None
-        )
+        and image != DEFAULT_IMAGE
-    elif source in ["hassio", "supervisor"]:
+    ):
-        haversion = VersionData(
+        image = f"{image}-homeassistant"
-            HaVersion(
-                session, source=HaVersionSource.SUPERVISOR, channel=channel, image=image
-            )
-        )
-    elif source in ["docker", "container"]:
-        if image is not None and image != DEFAULT_IMAGE:
-            image = f"{image}-homeassistant"
-        haversion = VersionData(
-            HaVersion(
-                session, source=HaVersionSource.CONTAINER, channel=channel, image=image
-            )
-        )
-    elif source == "haio":
-        haversion = VersionData(HaVersion(session, source=HaVersionSource.HAIO))
-    else:
-        haversion = VersionData(HaVersion(session, source=HaVersionSource.LOCAL))
 
-    if not name:
+    if not (name := config.get(CONF_NAME)):
-        if source == DEFAULT_SOURCE:
+        if source == HaVersionSource.LOCAL:
             name = DEFAULT_NAME_LOCAL
         else:
             name = DEFAULT_NAME_LATEST
 
-    async_add_entities([VersionSensor(haversion, name)], True)
+    async_add_entities(
+        [
+            VersionSensor(
+                VersionData(
+                    HaVersion(
+                        session=session, source=source, image=image, channel=channel
+                    )
+                ),
+                SensorEntityDescription(key=source, name=name),
+            )
+        ],
+        True,
+    )
 
 
 class VersionData:
@@ -120,9 +126,9 @@ class VersionData:
         """Get the latest version information."""
         try:
             await self.api.get_version()
-        except HaVersionFetchException as exception:
+        except pyhaversionexceptions.HaVersionFetchException as exception:
             _LOGGER.warning(exception)
-        except HaVersionParseException as exception:
+        except pyhaversionexceptions.HaVersionParseException as exception:
             _LOGGER.warning(
                 "Could not parse data received for %s - %s", self.api.source, exception
             )
@@ -131,32 +137,19 @@ class VersionData:
 class VersionSensor(SensorEntity):
     """Representation of a Home Assistant version sensor."""
 
-    def __init__(self, data: VersionData, name: str) -> None:
+    _attr_icon = "mdi:package-up"
+
+    def __init__(
+        self,
+        data: VersionData,
+        description: SensorEntityDescription,
+    ) -> None:
         """Initialize the Version sensor."""
         self.data = data
-        self._name = name
+        self.entity_description = description
-        self._state = None
 
     async def async_update(self):
         """Get the latest version information."""
         await self.data.async_update()
-
+        self._attr_state = self.data.api.version
-    @property
+        self._attr_extra_state_attributes = self.data.api.version_data
-    def name(self):
-        """Return the name of the sensor."""
-        return self._name
-
-    @property
-    def state(self):
-        """Return the state of the sensor."""
-        return self.data.api.version
-
-    @property
-    def extra_state_attributes(self):
-        """Return attributes for the sensor."""
-        return self.data.api.version_data
-
-    @property
-    def icon(self):
-        """Return the icon to use in the frontend, if any."""
-        return ICON

tests/components/version/test_sensor.py

@@ -1,26 +1,56 @@
 """The test for the version sensor platform."""
 from unittest.mock import patch
 
+from pyhaversion import HaVersionSource, exceptions as pyhaversionexceptions
+import pytest
+
+from homeassistant.components.version.sensor import ALL_SOURCES
 from homeassistant.setup import async_setup_component
 
 MOCK_VERSION = "10.0"
 
 
-async def test_version_sensor(hass):
+@pytest.mark.parametrize(
-    """Test the Version sensor."""
+    "source",
-    config = {"sensor": {"platform": "version"}}
+    ALL_SOURCES,
+)
+async def test_version_source(hass, source):
+    """Test the Version sensor with different sources."""
+    config = {
+        "sensor": {"platform": "version", "source": source, "image": "qemux86-64"}
+    }
 
-    assert await async_setup_component(hass, "sensor", config)
+    with patch("pyhaversion.version.HaVersion.version", MOCK_VERSION):
-
-
-async def test_version(hass):
-    """Test the Version sensor."""
-    config = {"sensor": {"platform": "version", "name": "test"}}
-
-    with patch("homeassistant.const.__version__", MOCK_VERSION):
         assert await async_setup_component(hass, "sensor", config)
         await hass.async_block_till_done()
 
-    state = hass.states.get("sensor.test")
+    name = "current_version" if source == HaVersionSource.LOCAL else "latest_version"
+    state = hass.states.get(f"sensor.{name}")
 
-    assert state.state == "10.0"
+    assert state.state == MOCK_VERSION
+
+
+async def test_version_fetch_exception(hass, caplog):
+    """Test fetch exception thrown during updates."""
+    config = {"sensor": {"platform": "version"}}
+    with patch(
+        "pyhaversion.version.HaVersion.get_version",
+        side_effect=pyhaversionexceptions.HaVersionFetchException(
+            "Fetch exception from pyhaversion"
+        ),
+    ):
+        assert await async_setup_component(hass, "sensor", config)
+        await hass.async_block_till_done()
+        assert "Fetch exception from pyhaversion" in caplog.text
+
+
+async def test_version_parse_exception(hass, caplog):
+    """Test parse exception thrown during updates."""
+    config = {"sensor": {"platform": "version"}}
+    with patch(
+        "pyhaversion.version.HaVersion.get_version",
+        side_effect=pyhaversionexceptions.HaVersionParseException,
+    ):
+        assert await async_setup_component(hass, "sensor", config)
+        await hass.async_block_till_done()
+        assert "Could not parse data received for HaVersionSource.LOCAL" in caplog.text