Consolidate zigbee schema factory and shared integration setup

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

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_

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}"