SNMP sensor refactor to ManualTriggerSensorEntity (#98630)

* SNMP to ManualTriggerSensorEntity

* Mods
This commit is contained in:
G Johansson 2023-08-24 11:45:14 +02:00 committed by GitHub
parent 3b31c58eba
commit 31a8a62165
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -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