From edeaded161c77776c01ed77f94347bef5066bc60 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Thu, 13 Aug 2026 09:20:56 -0500 Subject: [PATCH] Consolidate zigbee schema factory and shared integration setup --- esphome/components/binary_sensor/__init__.py | 12 +---- esphome/core/entity_helpers.py | 53 ++++++++++++++++---- tests/unit_tests/core/test_entity_helpers.py | 51 +++++++++---------- 3 files changed, 68 insertions(+), 48 deletions(-) diff --git a/esphome/components/binary_sensor/__init__.py b/esphome/components/binary_sensor/__init__.py index ac23e99350..73546b85c9 100644 --- a/esphome/components/binary_sensor/__init__.py +++ b/esphome/components/binary_sensor/__init__.py @@ -27,7 +27,6 @@ from esphome.const import ( CONF_STATE, CONF_TIMING, CONF_TRIGGER_ID, - CONF_WEB_SERVER, DEVICE_CLASS_BATTERY, DEVICE_CLASS_BATTERY_CHARGING, DEVICE_CLASS_CARBON_MONOXIDE, @@ -608,16 +607,7 @@ async def setup_binary_sensor_core_(var, config): CORE.add_job(_build_binary_sensor_automations, var, config) - if mqtt_id := config.get(CONF_MQTT_ID): - from esphome.components import mqtt - - mqtt_ = cg.new_Pvariable(mqtt_id, var) - await mqtt.register_mqtt_component(mqtt_, config) - - if web_server_config := config.get(CONF_WEB_SERVER): - from esphome.components import web_server - - await web_server.add_entity_config(var, web_server_config) + await entity_helpers.setup_entity_integrations(var, config) if "zigbee" in CORE.loaded_integrations: from esphome.components import zigbee diff --git a/esphome/core/entity_helpers.py b/esphome/core/entity_helpers.py index 455b72196c..40f9dbd9a3 100644 --- a/esphome/core/entity_helpers.py +++ b/esphome/core/entity_helpers.py @@ -16,6 +16,7 @@ from esphome.const import ( CONF_ICON, CONF_ID, CONF_INTERNAL, + CONF_MQTT_ID, CONF_NAME, CONF_REPORT, CONF_SORTING_GROUP_ID, @@ -700,7 +701,6 @@ ZIGBEE_MAX_EP_NUMBER = 239 _zigbee_ns = cg.esphome_ns.namespace("zigbee") _ZigbeeComponent = _zigbee_ns.class_("ZigbeeComponent", cg.Component) -_ZigbeeBinarySensor = _zigbee_ns.class_("ZigbeeBinarySensor", cg.Component) _zigbee_report = _zigbee_ns.enum("ZigbeeReportT") ZIGBEE_REPORT = { "coordinator": _zigbee_report.ZIGBEE_REPORT_COORDINATOR, @@ -739,14 +739,28 @@ ZIGBEE_BASE_ENTITY_SCHEMA = cv.Schema( } ) -ZIGBEE_BINARY_SENSOR_SCHEMA = ZIGBEE_BASE_ENTITY_SCHEMA.extend( - { - cv.OnlyWith(CONF_ZIGBEE_ID, ["nrf52", "zigbee"]): cv.use_id(_ZigbeeComponent), - cv.OnlyWith(CONF_ZIGBEE_BINARY_SENSOR, ["nrf52", "zigbee"]): cv.declare_id( - _ZigbeeBinarySensor - ), - } -) +# Entity platform -> (config key, zigbee C++ class). A unit test checks each +# class against the owning declaration in zigbee_zephyr. +_ZIGBEE_ENTITY_CLASSES = { + "binary_sensor": (CONF_ZIGBEE_BINARY_SENSOR, "ZigbeeBinarySensor"), +} + + +def _zigbee_entity_schema(platform: str) -> cv.Schema: + conf_key, class_name = _ZIGBEE_ENTITY_CLASSES[platform] + return ZIGBEE_BASE_ENTITY_SCHEMA.extend( + { + cv.OnlyWith(CONF_ZIGBEE_ID, ["nrf52", "zigbee"]): cv.use_id( + _ZigbeeComponent + ), + cv.OnlyWith(conf_key, ["nrf52", "zigbee"]): cv.declare_id( + _zigbee_ns.class_(class_name, cg.Component) + ), + } + ) + + +ZIGBEE_BINARY_SENSOR_SCHEMA = _zigbee_entity_schema("binary_sensor") def lazy_load_validator( @@ -765,3 +779,24 @@ def lazy_load_validator( return delegate(config) return validator + + +async def setup_entity_integrations(var: MockObj, config: ConfigType) -> MockObj | None: + """Register the mqtt companion and web_server entry for an entity. + + Imports the integrations lazily; returns the mqtt companion (or None) + so callers can apply integration specific options to it. + """ + mqtt_ = None + if (mqtt_id := config.get(CONF_MQTT_ID)) is not None: + from esphome.components import mqtt + + mqtt_ = cg.new_Pvariable(mqtt_id, var) + await mqtt.register_mqtt_component(mqtt_, config) + + if web_server_config := config.get(CONF_WEB_SERVER): + from esphome.components import web_server + + await web_server.add_entity_config(var, web_server_config) + + return mqtt_ diff --git a/tests/unit_tests/core/test_entity_helpers.py b/tests/unit_tests/core/test_entity_helpers.py index 5a76761b3d..abcc6db7e8 100644 --- a/tests/unit_tests/core/test_entity_helpers.py +++ b/tests/unit_tests/core/test_entity_helpers.py @@ -1266,21 +1266,19 @@ def test_lazy_load_validator_defers_import() -> None: validator = lazy_load_validator("zigbee", "validate_binary_sensor") config = {CONF_NAME: "test"} - original = CORE.loaded_integrations - CORE.loaded_integrations = set() - try: - with patch("esphome.core.entity_helpers.import_module") as import_mock: - assert validator(config) is config - import_mock.assert_not_called() + with ( + patch.object(CORE, "loaded_integrations", set()), + patch("esphome.core.entity_helpers.import_module") as import_mock, + ): + assert validator(config) is config + import_mock.assert_not_called() - CORE.loaded_integrations.add("zigbee") - delegate = import_mock.return_value.validate_binary_sensor - delegate.return_value = {CONF_NAME: "validated"} - assert validator(config) == {CONF_NAME: "validated"} - import_mock.assert_called_once_with("esphome.components.zigbee") - delegate.assert_called_once_with(config) - finally: - CORE.loaded_integrations = original + CORE.loaded_integrations.add("zigbee") + delegate = import_mock.return_value.validate_binary_sensor + delegate.return_value = {CONF_NAME: "validated"} + assert validator(config) == {CONF_NAME: "validated"} + import_mock.assert_called_once_with("esphome.components.zigbee") + delegate.assert_called_once_with(config) def test_lazy_load_validator_rejects_unknown_component() -> None: @@ -1293,26 +1291,21 @@ def test_lazy_load_validator_names_missing_hook() -> None: """A missing hook raises a clear error naming the component and hook.""" validator = lazy_load_validator("zigbee", "no_such_hook") - original = CORE.loaded_integrations - CORE.loaded_integrations = {"zigbee"} - try: - with ( - patch("esphome.core.entity_helpers.import_module") as import_mock, - pytest.raises(ValueError, match="no_such_hook"), - ): - del import_mock.return_value.no_such_hook - validator({}) - finally: - CORE.loaded_integrations = original + with ( + patch.object(CORE, "loaded_integrations", {"zigbee"}), + patch("esphome.core.entity_helpers.import_module") as import_mock, + pytest.raises(ValueError, match="no_such_hook"), + ): + del import_mock.return_value.no_such_hook + validator({}) def test_integration_class_handles_match_owning_definitions() -> None: """The cheap class handles must stay string-equal to the integrations' own declarations, or use_id/declare_id resolution silently drifts.""" from esphome.components import mqtt, web_server + from esphome.components.zigbee import zigbee_zephyr from esphome.components.zigbee.const import ZigbeeComponent - from esphome.components.zigbee.zigbee_zephyr import ZigbeeBinarySensor - from esphome.core import entity_helpers mqtt_handle = mqtt_component_class("MQTTBinarySensorComponent") assert str(mqtt_handle) == str(mqtt.MQTTBinarySensorComponent) @@ -1324,4 +1317,6 @@ def test_integration_class_handles_match_owning_definitions() -> None: assert web_server.WebServer.inherits_from(entity_helpers._WebServer) assert str(entity_helpers._ZigbeeComponent) == str(ZigbeeComponent) - assert str(entity_helpers._ZigbeeBinarySensor) == str(ZigbeeBinarySensor) + for _conf_key, class_name in entity_helpers._ZIGBEE_ENTITY_CLASSES.values(): + owning = getattr(zigbee_zephyr, class_name) + assert str(owning) == f"zigbee::{class_name}"