From 31a8a62165829fc481b0a28452f53c8f100aed25 Mon Sep 17 00:00:00 2001
From: G Johansson <goran.johansson@shiftit.se>
Date: Thu, 24 Aug 2023 11:45:14 +0200
Subject: [PATCH] SNMP sensor refactor to ManualTriggerSensorEntity (#98630)

* SNMP to ManualTriggerSensorEntity

* Mods
---
 homeassistant/components/snmp/sensor.py | 68 ++++++++++++++++++-------
 1 file changed, 50 insertions(+), 18 deletions(-)

diff --git a/homeassistant/components/snmp/sensor.py b/homeassistant/components/snmp/sensor.py
index fc8068fb532..85c69ddf76b 100644
--- a/homeassistant/components/snmp/sensor.py
+++ b/homeassistant/components/snmp/sensor.py
@@ -19,11 +19,15 @@ from pysnmp.hlapi.asyncio import (
 )
 import voluptuous as vol
 
-from homeassistant.components.sensor import PLATFORM_SCHEMA
+from homeassistant.components.sensor import CONF_STATE_CLASS, PLATFORM_SCHEMA
 from homeassistant.const import (
+    CONF_DEVICE_CLASS,
     CONF_HOST,
+    CONF_ICON,
+    CONF_NAME,
     CONF_PORT,
     CONF_UNIQUE_ID,
+    CONF_UNIT_OF_MEASUREMENT,
     CONF_USERNAME,
     CONF_VALUE_TEMPLATE,
     STATE_UNKNOWN,
@@ -31,9 +35,12 @@ 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 import Template
 from homeassistant.helpers.template_entity import (
+    CONF_AVAILABILITY,
+    CONF_PICTURE,
     TEMPLATE_SENSOR_BASE_SCHEMA,
-    TemplateSensor,
+    ManualTriggerSensorEntity,
 )
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 
@@ -64,6 +71,16 @@ _LOGGER = logging.getLogger(__name__)
 
 SCAN_INTERVAL = timedelta(seconds=10)
 
+TRIGGER_ENTITY_OPTIONS = (
+    CONF_AVAILABILITY,
+    CONF_DEVICE_CLASS,
+    CONF_ICON,
+    CONF_PICTURE,
+    CONF_UNIQUE_ID,
+    CONF_STATE_CLASS,
+    CONF_UNIT_OF_MEASUREMENT,
+)
+
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
     {
         vol.Required(CONF_BASEOID): cv.string,
@@ -106,7 +123,6 @@ async def async_setup_platform(
     privproto = config[CONF_PRIV_PROTOCOL]
     accept_errors = config.get(CONF_ACCEPT_ERRORS)
     default_value = config.get(CONF_DEFAULT_VALUE)
-    unique_id = config.get(CONF_UNIQUE_ID)
 
     try:
         # Try IPv4 first.
@@ -151,35 +167,50 @@ async def async_setup_platform(
         _LOGGER.error("Please check the details in the configuration file")
         return
 
+    name = config.get(CONF_NAME, Template(DEFAULT_NAME, hass))
+    trigger_entity_config = {CONF_NAME: name}
+    for key in TRIGGER_ENTITY_OPTIONS:
+        if key not in config:
+            continue
+        trigger_entity_config[key] = config[key]
+
+    value_template: Template | None = config.get(CONF_VALUE_TEMPLATE)
+    if value_template is not None:
+        value_template.hass = hass
+
     data = SnmpData(request_args, baseoid, accept_errors, default_value)
-    async_add_entities([SnmpSensor(hass, data, config, unique_id)], True)
+    async_add_entities([SnmpSensor(hass, data, trigger_entity_config, value_template)])
 
 
-class SnmpSensor(TemplateSensor):
+class SnmpSensor(ManualTriggerSensorEntity):
     """Representation of a SNMP sensor."""
 
     _attr_should_poll = True
 
-    def __init__(self, hass, data, config, unique_id):
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        data: SnmpData,
+        config: ConfigType,
+        value_template: Template | None,
+    ) -> None:
         """Initialize the sensor."""
-        super().__init__(
-            hass, config=config, unique_id=unique_id, fallback_name=DEFAULT_NAME
-        )
+        super().__init__(hass, config)
         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
+        self._value_template = value_template
 
-    @property
-    def native_value(self):
-        """Return the state of the sensor."""
-        return self._state
+    async def async_added_to_hass(self) -> None:
+        """Handle adding to Home Assistant."""
+        await super().async_added_to_hass()
+        await self.async_update()
 
     async def async_update(self) -> None:
         """Get the latest data and updates the states."""
         await self.data.async_update()
 
+        raw_value = self.data.value
+
         if (value := self.data.value) is None:
             value = STATE_UNKNOWN
         elif self._value_template is not None:
@@ -187,13 +218,14 @@ class SnmpSensor(TemplateSensor):
                 value, STATE_UNKNOWN
             )
 
-        self._state = value
+        self._attr_native_value = value
+        self._process_manual_data(raw_value)
 
 
 class SnmpData:
     """Get the latest data and update the states."""
 
-    def __init__(self, request_args, baseoid, accept_errors, default_value):
+    def __init__(self, request_args, baseoid, accept_errors, default_value) -> None:
         """Initialize the data object."""
         self._request_args = request_args
         self._baseoid = baseoid