Migrate SNMP sensor to TemplateEntity (#73324)

homeassistant/components/snmp/sensor.py

@@ -17,12 +17,11 @@ from pysnmp.hlapi.asyncio import (
 )
 import voluptuous as vol
 
-from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity
+from homeassistant.components.sensor import PLATFORM_SCHEMA
 from homeassistant.const import (
     CONF_HOST,
-    CONF_NAME,
     CONF_PORT,
-    CONF_UNIT_OF_MEASUREMENT,
+    CONF_UNIQUE_ID,
     CONF_USERNAME,
     CONF_VALUE_TEMPLATE,
     STATE_UNKNOWN,
@@ -30,6 +29,10 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.template_entity import (
+    TEMPLATE_SENSOR_BASE_SCHEMA,
+    TemplateSensor,
+)
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 
 from .const import (
@@ -66,9 +69,7 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
         vol.Optional(CONF_COMMUNITY, default=DEFAULT_COMMUNITY): cv.string,
         vol.Optional(CONF_DEFAULT_VALUE): cv.string,
         vol.Optional(CONF_HOST, default=DEFAULT_HOST): cv.string,
-        vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string,
         vol.Optional(CONF_PORT, default=DEFAULT_PORT): cv.port,
-        vol.Optional(CONF_UNIT_OF_MEASUREMENT): cv.string,
         vol.Optional(CONF_VALUE_TEMPLATE): cv.template,
         vol.Optional(CONF_VERSION, default=DEFAULT_VERSION): vol.In(SNMP_VERSIONS),
         vol.Optional(CONF_USERNAME): cv.string,
@@ -81,7 +82,7 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
             MAP_PRIV_PROTOCOLS
         ),
     }
-)
+).extend(TEMPLATE_SENSOR_BASE_SCHEMA.schema)
 
 
 async def async_setup_platform(
@@ -91,12 +92,10 @@ async def async_setup_platform(
     discovery_info: DiscoveryInfoType | None = None,
 ) -> None:
     """Set up the SNMP sensor."""
-    name = config.get(CONF_NAME)
     host = config.get(CONF_HOST)
     port = config.get(CONF_PORT)
     community = config.get(CONF_COMMUNITY)
     baseoid = config.get(CONF_BASEOID)
-    unit = config.get(CONF_UNIT_OF_MEASUREMENT)
     version = config[CONF_VERSION]
     username = config.get(CONF_USERNAME)
     authkey = config.get(CONF_AUTH_KEY)
@@ -105,10 +104,7 @@ async def async_setup_platform(
     privproto = config[CONF_PRIV_PROTOCOL]
     accept_errors = config.get(CONF_ACCEPT_ERRORS)
     default_value = config.get(CONF_DEFAULT_VALUE)
-    value_template = config.get(CONF_VALUE_TEMPLATE)
-
-    if value_template is not None:
-        value_template.hass = hass
+    unique_id = config.get(CONF_UNIQUE_ID)
 
     if version == "3":
 
@@ -146,35 +142,30 @@ async def async_setup_platform(
         return
 
     data = SnmpData(request_args, baseoid, accept_errors, default_value)
-    async_add_entities([SnmpSensor(data, name, unit, value_template)], True)
+    async_add_entities([SnmpSensor(hass, data, config, unique_id)], True)
 
 
-class SnmpSensor(SensorEntity):
+class SnmpSensor(TemplateSensor):
     """Representation of a SNMP sensor."""
 
-    def __init__(self, data, name, unit_of_measurement, value_template):
-        """Initialize the sensor."""
-        self.data = data
-        self._name = name
-        self._state = None
-        self._unit_of_measurement = unit_of_measurement
-        self._value_template = value_template
+    _attr_should_poll = True
 
-    @property
-    def name(self):
-        """Return the name of the sensor."""
-        return self._name
+    def __init__(self, hass, data, config, unique_id):
+        """Initialize the sensor."""
+        super().__init__(
+            hass, config=config, unique_id=unique_id, fallback_name=DEFAULT_NAME
+        )
+        self.data = data
+        self._state = None
+        self._value_template = config.get(CONF_VALUE_TEMPLATE)
+        if (value_template := self._value_template) is not None:
+            value_template.hass = hass
 
     @property
     def native_value(self):
         """Return the state of the sensor."""
         return self._state
 
-    @property
-    def native_unit_of_measurement(self):
-        """Return the unit the value is expressed in."""
-        return self._unit_of_measurement
-
     async def async_update(self):
         """Get the latest data and updates the states."""
         await self.data.async_update()

requirements_test_all.txt

@@ -1241,6 +1241,9 @@ pysmartapp==0.3.3
 # homeassistant.components.smartthings
 pysmartthings==0.7.6
 
+# homeassistant.components.snmp
+pysnmplib==5.0.15
+
 # homeassistant.components.soma
 pysoma==0.0.10
 

tests/components/snmp/__init__.py

@@ -0,0 +1 @@
+"""Tests for the SNMP integration."""

tests/components/snmp/test_sensor.py

@@ -0,0 +1,79 @@
+"""SNMP sensor tests."""
+
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+
+from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import entity_registry as er
+from homeassistant.setup import async_setup_component
+
+
+@pytest.fixture(autouse=True)
+def hlapi_mock():
+    """Mock out 3rd party API."""
+    mock_data = MagicMock()
+    mock_data.prettyPrint = Mock(return_value="hello")
+    with patch(
+        "homeassistant.components.snmp.sensor.getCmd",
+        return_value=(None, None, None, [[mock_data]]),
+    ):
+        yield
+
+
+async def test_basic_config(hass: HomeAssistant) -> None:
+    """Test basic entity configuration."""
+
+    config = {
+        SENSOR_DOMAIN: {
+            "platform": "snmp",
+            "host": "192.168.1.32",
+            "baseoid": "1.3.6.1.4.1.2021.10.1.3.1",
+        },
+    }
+
+    assert await async_setup_component(hass, SENSOR_DOMAIN, config)
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.snmp")
+    assert state.state == "hello"
+    assert state.attributes == {"friendly_name": "SNMP"}
+
+
+async def test_entity_config(hass: HomeAssistant) -> None:
+    """Test entity configuration."""
+
+    config = {
+        SENSOR_DOMAIN: {
+            # SNMP configuration
+            "platform": "snmp",
+            "host": "192.168.1.32",
+            "baseoid": "1.3.6.1.4.1.2021.10.1.3.1",
+            # Entity configuration
+            "icon": "{{'mdi:one_two_three'}}",
+            "picture": "{{'blabla.png'}}",
+            "device_class": "temperature",
+            "name": "{{'SNMP' + ' ' + 'Sensor'}}",
+            "state_class": "measurement",
+            "unique_id": "very_unique",
+            "unit_of_measurement": "beardsecond",
+        },
+    }
+
+    assert await async_setup_component(hass, SENSOR_DOMAIN, config)
+    await hass.async_block_till_done()
+
+    entity_registry = er.async_get(hass)
+    assert entity_registry.async_get("sensor.snmp_sensor").unique_id == "very_unique"
+
+    state = hass.states.get("sensor.snmp_sensor")
+    assert state.state == "hello"
+    assert state.attributes == {
+        "device_class": "temperature",
+        "entity_picture": "blabla.png",
+        "friendly_name": "SNMP Sensor",
+        "icon": "mdi:one_two_three",
+        "state_class": "measurement",
+        "unit_of_measurement": "beardsecond",
+    }