From 16ac303994234580fb63582a3fe91e1de0d2a86c Mon Sep 17 00:00:00 2001 From: epenet <6771947+epenet@users.noreply.github.com> Date: Wed, 18 Sep 2024 11:03:32 +0200 Subject: [PATCH] Move tcp base entity to separate module (#126181) --- homeassistant/components/tcp/binary_sensor.py | 3 +- homeassistant/components/tcp/common.py | 112 --------------- homeassistant/components/tcp/entity.py | 130 ++++++++++++++++++ homeassistant/components/tcp/sensor.py | 3 +- tests/components/tcp/test_binary_sensor.py | 4 +- tests/components/tcp/test_sensor.py | 6 +- 6 files changed, 139 insertions(+), 119 deletions(-) create mode 100644 homeassistant/components/tcp/entity.py diff --git a/homeassistant/components/tcp/binary_sensor.py b/homeassistant/components/tcp/binary_sensor.py index 638dfd53de5..13fd0787b5d 100644 --- a/homeassistant/components/tcp/binary_sensor.py +++ b/homeassistant/components/tcp/binary_sensor.py @@ -12,8 +12,9 @@ from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType -from .common import TCP_PLATFORM_SCHEMA, TcpEntity +from .common import TCP_PLATFORM_SCHEMA from .const import CONF_VALUE_ON +from .entity import TcpEntity PLATFORM_SCHEMA: Final = BINARY_SENSOR_PLATFORM_SCHEMA.extend(TCP_PLATFORM_SCHEMA) diff --git a/homeassistant/components/tcp/common.py b/homeassistant/components/tcp/common.py index 263fc416026..a89cd999ddd 100644 --- a/homeassistant/components/tcp/common.py +++ b/homeassistant/components/tcp/common.py @@ -2,10 +2,6 @@ from __future__ import annotations -import logging -import select -import socket -import ssl from typing import Any, Final import voluptuous as vol @@ -21,11 +17,7 @@ from homeassistant.const import ( CONF_VALUE_TEMPLATE, CONF_VERIFY_SSL, ) -from homeassistant.core import HomeAssistant -from homeassistant.exceptions import TemplateError import homeassistant.helpers.config_validation as cv -from homeassistant.helpers.entity import Entity -from homeassistant.helpers.typing import ConfigType from .const import ( CONF_BUFFER_SIZE, @@ -36,10 +28,6 @@ from .const import ( DEFAULT_TIMEOUT, DEFAULT_VERIFY_SSL, ) -from .model import TcpSensorConfig - -_LOGGER: Final = logging.getLogger(__name__) - TCP_PLATFORM_SCHEMA: Final[dict[vol.Marker, Any]] = { vol.Required(CONF_HOST): cv.string, @@ -54,103 +42,3 @@ TCP_PLATFORM_SCHEMA: Final[dict[vol.Marker, Any]] = { vol.Optional(CONF_SSL, default=DEFAULT_SSL): cv.boolean, vol.Optional(CONF_VERIFY_SSL, default=DEFAULT_VERIFY_SSL): cv.boolean, } - - -class TcpEntity(Entity): - """Base entity class for TCP platform.""" - - def __init__(self, hass: HomeAssistant, config: ConfigType) -> None: - """Set all the config values if they exist and get initial state.""" - - self._hass = hass - self._config: TcpSensorConfig = { - CONF_NAME: config[CONF_NAME], - CONF_HOST: config[CONF_HOST], - CONF_PORT: config[CONF_PORT], - CONF_TIMEOUT: config[CONF_TIMEOUT], - CONF_PAYLOAD: config[CONF_PAYLOAD], - CONF_UNIT_OF_MEASUREMENT: config.get(CONF_UNIT_OF_MEASUREMENT), - CONF_VALUE_TEMPLATE: config.get(CONF_VALUE_TEMPLATE), - CONF_VALUE_ON: config.get(CONF_VALUE_ON), - CONF_BUFFER_SIZE: config[CONF_BUFFER_SIZE], - CONF_SSL: config[CONF_SSL], - CONF_VERIFY_SSL: config[CONF_VERIFY_SSL], - } - - self._ssl_context: ssl.SSLContext | None = None - if self._config[CONF_SSL]: - self._ssl_context = ssl.create_default_context() - if not self._config[CONF_VERIFY_SSL]: - self._ssl_context.check_hostname = False - self._ssl_context.verify_mode = ssl.CERT_NONE - - self._state: str | None = None - self.update() - - @property - def name(self) -> str: - """Return the name of this sensor.""" - return self._config[CONF_NAME] - - def update(self) -> None: - """Get the latest value for this sensor.""" - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.settimeout(self._config[CONF_TIMEOUT]) - try: - sock.connect((self._config[CONF_HOST], self._config[CONF_PORT])) - except OSError as err: - _LOGGER.error( - "Unable to connect to %s on port %s: %s", - self._config[CONF_HOST], - self._config[CONF_PORT], - err, - ) - return - - if self._ssl_context is not None: - sock = self._ssl_context.wrap_socket( - sock, server_hostname=self._config[CONF_HOST] - ) - - try: - sock.send(self._config[CONF_PAYLOAD].encode()) - except OSError as err: - _LOGGER.error( - "Unable to send payload %r to %s on port %s: %s", - self._config[CONF_PAYLOAD], - self._config[CONF_HOST], - self._config[CONF_PORT], - err, - ) - return - - readable, _, _ = select.select([sock], [], [], self._config[CONF_TIMEOUT]) - if not readable: - _LOGGER.warning( - ( - "Timeout (%s second(s)) waiting for a response after " - "sending %r to %s on port %s" - ), - self._config[CONF_TIMEOUT], - self._config[CONF_PAYLOAD], - self._config[CONF_HOST], - self._config[CONF_PORT], - ) - return - - value = sock.recv(self._config[CONF_BUFFER_SIZE]).decode() - - value_template = self._config[CONF_VALUE_TEMPLATE] - if value_template is not None: - try: - self._state = value_template.render(parse_result=False, value=value) - except TemplateError: - _LOGGER.error( - "Unable to render template of %r with value: %r", - self._config[CONF_VALUE_TEMPLATE], - value, - ) - return - return - - self._state = value diff --git a/homeassistant/components/tcp/entity.py b/homeassistant/components/tcp/entity.py new file mode 100644 index 00000000000..eaf5cb6963e --- /dev/null +++ b/homeassistant/components/tcp/entity.py @@ -0,0 +1,130 @@ +"""Common code for TCP component.""" + +from __future__ import annotations + +import logging +import select +import socket +import ssl +from typing import Final + +from homeassistant.const import ( + CONF_HOST, + CONF_NAME, + CONF_PAYLOAD, + CONF_PORT, + CONF_SSL, + CONF_TIMEOUT, + CONF_UNIT_OF_MEASUREMENT, + CONF_VALUE_TEMPLATE, + CONF_VERIFY_SSL, +) +from homeassistant.core import HomeAssistant +from homeassistant.exceptions import TemplateError +from homeassistant.helpers.entity import Entity +from homeassistant.helpers.typing import ConfigType + +from .const import CONF_BUFFER_SIZE, CONF_VALUE_ON +from .model import TcpSensorConfig + +_LOGGER: Final = logging.getLogger(__name__) + + +class TcpEntity(Entity): + """Base entity class for TCP platform.""" + + def __init__(self, hass: HomeAssistant, config: ConfigType) -> None: + """Set all the config values if they exist and get initial state.""" + + self._hass = hass + self._config: TcpSensorConfig = { + CONF_NAME: config[CONF_NAME], + CONF_HOST: config[CONF_HOST], + CONF_PORT: config[CONF_PORT], + CONF_TIMEOUT: config[CONF_TIMEOUT], + CONF_PAYLOAD: config[CONF_PAYLOAD], + CONF_UNIT_OF_MEASUREMENT: config.get(CONF_UNIT_OF_MEASUREMENT), + CONF_VALUE_TEMPLATE: config.get(CONF_VALUE_TEMPLATE), + CONF_VALUE_ON: config.get(CONF_VALUE_ON), + CONF_BUFFER_SIZE: config[CONF_BUFFER_SIZE], + CONF_SSL: config[CONF_SSL], + CONF_VERIFY_SSL: config[CONF_VERIFY_SSL], + } + + self._ssl_context: ssl.SSLContext | None = None + if self._config[CONF_SSL]: + self._ssl_context = ssl.create_default_context() + if not self._config[CONF_VERIFY_SSL]: + self._ssl_context.check_hostname = False + self._ssl_context.verify_mode = ssl.CERT_NONE + + self._state: str | None = None + self.update() + + @property + def name(self) -> str: + """Return the name of this sensor.""" + return self._config[CONF_NAME] + + def update(self) -> None: + """Get the latest value for this sensor.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(self._config[CONF_TIMEOUT]) + try: + sock.connect((self._config[CONF_HOST], self._config[CONF_PORT])) + except OSError as err: + _LOGGER.error( + "Unable to connect to %s on port %s: %s", + self._config[CONF_HOST], + self._config[CONF_PORT], + err, + ) + return + + if self._ssl_context is not None: + sock = self._ssl_context.wrap_socket( + sock, server_hostname=self._config[CONF_HOST] + ) + + try: + sock.send(self._config[CONF_PAYLOAD].encode()) + except OSError as err: + _LOGGER.error( + "Unable to send payload %r to %s on port %s: %s", + self._config[CONF_PAYLOAD], + self._config[CONF_HOST], + self._config[CONF_PORT], + err, + ) + return + + readable, _, _ = select.select([sock], [], [], self._config[CONF_TIMEOUT]) + if not readable: + _LOGGER.warning( + ( + "Timeout (%s second(s)) waiting for a response after " + "sending %r to %s on port %s" + ), + self._config[CONF_TIMEOUT], + self._config[CONF_PAYLOAD], + self._config[CONF_HOST], + self._config[CONF_PORT], + ) + return + + value = sock.recv(self._config[CONF_BUFFER_SIZE]).decode() + + value_template = self._config[CONF_VALUE_TEMPLATE] + if value_template is not None: + try: + self._state = value_template.render(parse_result=False, value=value) + except TemplateError: + _LOGGER.error( + "Unable to render template of %r with value: %r", + self._config[CONF_VALUE_TEMPLATE], + value, + ) + return + return + + self._state = value diff --git a/homeassistant/components/tcp/sensor.py b/homeassistant/components/tcp/sensor.py index a3bd4b2c619..1d53b21bc2e 100644 --- a/homeassistant/components/tcp/sensor.py +++ b/homeassistant/components/tcp/sensor.py @@ -13,7 +13,8 @@ from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType, StateType -from .common import TCP_PLATFORM_SCHEMA, TcpEntity +from .common import TCP_PLATFORM_SCHEMA +from .entity import TcpEntity PLATFORM_SCHEMA: Final = SENSOR_PLATFORM_SCHEMA.extend(TCP_PLATFORM_SCHEMA) diff --git a/tests/components/tcp/test_binary_sensor.py b/tests/components/tcp/test_binary_sensor.py index 05aa2a471db..c84a36016ad 100644 --- a/tests/components/tcp/test_binary_sensor.py +++ b/tests/components/tcp/test_binary_sensor.py @@ -23,9 +23,9 @@ TEST_ENTITY = "binary_sensor.test_name" def mock_socket_fixture(): """Mock the socket.""" with ( - patch("homeassistant.components.tcp.common.socket.socket") as mock_socket, + patch("homeassistant.components.tcp.entity.socket.socket") as mock_socket, patch( - "homeassistant.components.tcp.common.select.select", + "homeassistant.components.tcp.entity.select.select", return_value=(True, False, False), ), ): diff --git a/tests/components/tcp/test_sensor.py b/tests/components/tcp/test_sensor.py index 04fbb2c667e..27003df46cd 100644 --- a/tests/components/tcp/test_sensor.py +++ b/tests/components/tcp/test_sensor.py @@ -43,7 +43,7 @@ socket_test_value = "123" @pytest.fixture(name="mock_socket") def mock_socket_fixture(mock_select): """Mock socket.""" - with patch("homeassistant.components.tcp.common.socket.socket") as mock_socket: + with patch("homeassistant.components.tcp.entity.socket.socket") as mock_socket: socket_instance = mock_socket.return_value.__enter__.return_value socket_instance.recv.return_value = socket_test_value.encode() yield socket_instance @@ -53,7 +53,7 @@ def mock_socket_fixture(mock_select): def mock_select_fixture(): """Mock select.""" with patch( - "homeassistant.components.tcp.common.select.select", + "homeassistant.components.tcp.entity.select.select", return_value=(True, False, False), ) as mock_select: yield mock_select @@ -63,7 +63,7 @@ def mock_select_fixture(): def mock_ssl_context_fixture(): """Mock select.""" with patch( - "homeassistant.components.tcp.common.ssl.create_default_context", + "homeassistant.components.tcp.entity.ssl.create_default_context", ) as mock_ssl_context: mock_ssl_context.return_value.wrap_socket.return_value.recv.return_value = ( socket_test_value + "567"