diff --git a/esphome/voluptuous_schema.py b/esphome/voluptuous_schema.py index 8fb966e3b2..7220fb307f 100644 --- a/esphome/voluptuous_schema.py +++ b/esphome/voluptuous_schema.py @@ -225,9 +225,10 @@ class _Schema(vol.Schema): return ret schema = schemas[0] + extra_schemas = self._extra_schemas.copy() + if isinstance(schema, _Schema): + extra_schemas.extend(schema._extra_schemas) if isinstance(schema, vol.Schema): schema = schema.schema ret = super().extend(schema, extra=extra) - return _Schema( - ret.schema, extra=ret.extra, extra_schemas=self._extra_schemas.copy() - ) + return _Schema(ret.schema, extra=ret.extra, extra_schemas=extra_schemas) diff --git a/tests/component_tests/config_validation/test_config.py b/tests/component_tests/config_validation/test_config.py new file mode 100644 index 0000000000..1a9b9bc1f3 --- /dev/null +++ b/tests/component_tests/config_validation/test_config.py @@ -0,0 +1,51 @@ +""" +Test schema.extend functionality in esphome.config_validation. +""" + +from typing import Any + +import esphome.config_validation as cv + + +def test_config_extend() -> None: + """Test that schema.extend correctly merges schemas with extras.""" + + def func1(data: dict[str, Any]) -> dict[str, Any]: + data["extra_1"] = "value1" + return data + + def func2(data: dict[str, Any]) -> dict[str, Any]: + data["extra_2"] = "value2" + return data + + schema1 = cv.Schema( + { + cv.Required("key1"): cv.string, + } + ) + schema1.add_extra(func1) + schema2 = cv.Schema( + { + cv.Required("key2"): cv.string, + } + ) + schema2.add_extra(func2) + extended_schema = schema1.extend(schema2) + config = { + "key1": "initial_value1", + "key2": "initial_value2", + } + validated = extended_schema(config) + assert validated["key1"] == "initial_value1" + assert validated["key2"] == "initial_value2" + assert validated["extra_1"] == "value1" + assert validated["extra_2"] == "value2" + + # Check the opposite order of extension + extended_schema = schema2.extend(schema1) + + validated = extended_schema(config) + assert validated["key1"] == "initial_value1" + assert validated["key2"] == "initial_value2" + assert validated["extra_1"] == "value1" + assert validated["extra_2"] == "value2"