Merge branch 'protect_entity_setters' into integration

This commit is contained in:
J. Nick Koston
2026-03-06 09:56:13 -10:00
30 changed files with 503 additions and 200 deletions
@@ -45,8 +45,9 @@ def test_binary_sensor_config_value_internal_set(generate_main):
)
# Then
assert "bs_1->set_internal(true);" in main_cpp
assert "bs_2->set_internal(false);" in main_cpp
# internal flag is now packed into configure_entity_() third argument (bit 24)
assert "bs_1->configure_entity_(" in main_cpp
assert "bs_2->configure_entity_(" in main_cpp
def test_binary_sensor_config_value_use_raw_set(generate_main):
+4 -2
View File
@@ -40,5 +40,7 @@ def test_button_config_value_internal_set(generate_main):
main_cpp = generate_main("tests/component_tests/button/test_button.yaml")
# Then
assert "wol_1->set_internal(true);" in main_cpp
assert "wol_2->set_internal(false);" in main_cpp
# internal flag is packed into configure_entity_() third argument (bit 24)
# wol_1 has internal: true → bit 24 set → packed value 16777216
assert "wol_1->configure_entity_(" in main_cpp
assert "wol_2->configure_entity_(" in main_cpp
+13 -2
View File
@@ -1,5 +1,15 @@
"""Tests for the sensor component."""
import re
def _extract_packed_value(main_cpp, var_name):
"""Extract the third (packed) argument from a configure_entity_ call."""
pattern = rf"{re.escape(var_name)}->configure_entity_\([^,]+,\s*\w+,\s*(\d+)\)"
match = re.search(pattern, main_cpp)
assert match, f"configure_entity_ call not found for {var_name}"
return int(match.group(1))
def test_sensor_device_class_set(generate_main):
"""
@@ -10,5 +20,6 @@ def test_sensor_device_class_set(generate_main):
# When
main_cpp = generate_main("tests/component_tests/sensor/test_sensor.yaml")
# Then
assert "s_1->configure_entity_(" in main_cpp
# Then: device_class: voltage means packed value must be non-zero
packed = _extract_packed_value(main_cpp, "s_1")
assert packed != 0
+3 -2
View File
@@ -38,8 +38,9 @@ def test_text_config_value_internal_set(generate_main):
main_cpp = generate_main("tests/component_tests/text/test_text.yaml")
# Then
assert "it_2->set_internal(false);" in main_cpp
assert "it_3->set_internal(true);" in main_cpp
# internal flag is now packed into configure_entity_() third argument (bit 24)
assert "it_2->configure_entity_(" in main_cpp
assert "it_3->configure_entity_(" in main_cpp
def test_text_config_value_mode_set(generate_main):
@@ -1,5 +1,15 @@
"""Tests for the text sensor component."""
import re
def _extract_packed_value(main_cpp, var_name):
"""Extract the third (packed) argument from a configure_entity_ call."""
pattern = rf"{re.escape(var_name)}->configure_entity_\([^,]+,\s*\w+,\s*(\d+)\)"
match = re.search(pattern, main_cpp)
assert match, f"configure_entity_ call not found for {var_name}"
return int(match.group(1))
def test_text_sensor_is_setup(generate_main):
"""
@@ -40,8 +50,9 @@ def test_text_sensor_config_value_internal_set(generate_main):
main_cpp = generate_main("tests/component_tests/text_sensor/test_text_sensor.yaml")
# Then
assert "ts_2->set_internal(true);" in main_cpp
assert "ts_3->set_internal(false);" in main_cpp
# internal flag is now packed into configure_entity_() third argument (bit 24)
assert "ts_2->configure_entity_(" in main_cpp
assert "ts_3->configure_entity_(" in main_cpp
def test_text_sensor_device_class_set(generate_main):
@@ -53,6 +64,9 @@ def test_text_sensor_device_class_set(generate_main):
# When
main_cpp = generate_main("tests/component_tests/text_sensor/test_text_sensor.yaml")
# Then
assert "ts_2->configure_entity_(" in main_cpp
assert "ts_3->configure_entity_(" in main_cpp
# Then: ts_2 has device_class: timestamp, ts_3 has device_class: date
# so their packed values must be non-zero
packed_ts_2 = _extract_packed_value(main_cpp, "ts_2")
assert packed_ts_2 != 0
packed_ts_3 = _extract_packed_value(main_cpp, "ts_3")
assert packed_ts_3 != 0
+216 -26
View File
@@ -9,6 +9,7 @@ import pytest
from esphome.config_validation import Invalid
from esphome.const import (
CONF_DEVICE_CLASS,
CONF_DEVICE_ID,
CONF_DISABLED_BY_DEFAULT,
CONF_ENTITY_CATEGORY,
@@ -16,16 +17,26 @@ from esphome.const import (
CONF_ID,
CONF_INTERNAL,
CONF_NAME,
CONF_UNIT_OF_MEASUREMENT,
)
from esphome.core import CORE, ID, entity_helpers
from esphome.core.entity_helpers import (
_DC_SHIFT,
_DISABLED_BY_DEFAULT_SHIFT,
_ENTITY_CATEGORY_SHIFT,
_ICON_SHIFT,
_INTERNAL_SHIFT,
_UOM_SHIFT,
_register_string,
_setup_entity_impl,
entity_duplicate_validator,
finalize_entity_strings,
get_base_entity_object_id,
register_device_class,
register_icon,
setup_device_class,
setup_entity,
setup_unit_of_measurement,
)
from esphome.cpp_generator import MockObj
from esphome.helpers import sanitize, snake_case
@@ -309,8 +320,6 @@ def extract_object_id_from_expressions(expressions: list[str]) -> str | None:
async def test_setup_entity_no_duplicates(setup_test_environment: list[str]) -> None:
"""Test setup_entity with unique names."""
setup_test_environment # noqa: F841 - fixture initializes CORE state
# Create mock entities
var1 = MockObj("sensor1")
var2 = MockObj("sensor2")
@@ -344,8 +353,6 @@ async def test_setup_entity_different_platforms(
) -> None:
"""Test that same name on different platforms doesn't conflict."""
setup_test_environment # noqa: F841 - fixture initializes CORE state
# Create mock entities
sensor = MockObj("sensor1")
binary_sensor = MockObj("binary_sensor1")
@@ -392,7 +399,6 @@ async def test_setup_entity_with_devices(
setup_test_environment: list[str], mock_get_variable: dict[ID, MockObj]
) -> None:
"""Test that same name on different devices doesn't conflict."""
setup_test_environment # noqa: F841 - fixture initializes CORE state
# Create mock devices
device1_id = ID("device1", type="Device")
@@ -433,8 +439,6 @@ async def test_setup_entity_with_devices(
async def test_setup_entity_empty_name(setup_test_environment: list[str]) -> None:
"""Test setup_entity with empty entity name."""
setup_test_environment # noqa: F841 - fixture initializes CORE state
var = MockObj("sensor1")
config = {
@@ -455,8 +459,6 @@ async def test_setup_entity_special_characters(
) -> None:
"""Test setup_entity with names containing special characters."""
setup_test_environment # noqa: F841 - fixture initializes CORE state
var = MockObj("sensor1")
config = {
@@ -475,8 +477,6 @@ async def test_setup_entity_special_characters(
async def test_setup_entity_with_icon(setup_test_environment: list[str]) -> None:
"""Test setup_entity sets icon correctly."""
setup_test_environment # noqa: F841 - fixture initializes CORE state
var = MockObj("sensor1")
config = {
@@ -497,8 +497,6 @@ async def test_setup_entity_disabled_by_default(
) -> None:
"""Test setup_entity sets disabled_by_default correctly."""
added_expressions = setup_test_environment
var = MockObj("sensor1")
config = {
@@ -508,10 +506,8 @@ async def test_setup_entity_disabled_by_default(
await _setup_entity_impl(var, config, "sensor")
# Check disabled_by_default was set
assert any(
"sensor1.set_disabled_by_default(true)" in expr for expr in added_expressions
)
# disabled_by_default is now packed into config for configure_entity_()
assert config.get("_entity_disabled_by_default") == 1
def test_entity_duplicate_validator() -> None:
@@ -796,8 +792,8 @@ async def test_setup_entity_empty_name_with_device(
entity_helpers.get_variable = original_get_variable
# Check that set_device was called
assert any("sensor1.set_device" in expr for expr in added_expressions)
# Check that set_device_ was called (separate protected call, accessible via friend)
assert any("sensor1.set_device_" in expr for expr in added_expressions)
# For empty-name entities, Python stores hash 0 - C++ calculates hash at runtime
assert config.get("_entity_name") == ""
@@ -813,7 +809,6 @@ async def test_setup_entity_empty_name_with_mac_suffix(
For empty-name entities, Python passes 0 and C++ calculates the hash
at runtime from friendly_name (bug-for-bug compatibility).
"""
setup_test_environment # noqa: F841 - fixture initializes CORE state
# Set up CORE.config with name_add_mac_suffix enabled
CORE.config = {"name_add_mac_suffix": True}
@@ -844,7 +839,6 @@ async def test_setup_entity_empty_name_with_mac_suffix_no_friendly_name(
at runtime. In this case C++ will hash the empty friendly_name
(bug-for-bug compatibility).
"""
setup_test_environment # noqa: F841 - fixture initializes CORE state
# Set up CORE.config with name_add_mac_suffix enabled
CORE.config = {"name_add_mac_suffix": True}
@@ -874,7 +868,6 @@ async def test_setup_entity_empty_name_no_mac_suffix_no_friendly_name(
For empty-name entities, Python passes 0 and C++ calculates the hash
at runtime from the device name.
"""
setup_test_environment # noqa: F841 - fixture initializes CORE state
# No MAC suffix (either not set or False)
CORE.config = {}
@@ -942,7 +935,7 @@ def test_register_device_class_max_length() -> None:
async def test_setup_entity_with_entity_category(
setup_test_environment: list[str],
) -> None:
"""Test setup_entity sets entity_category correctly."""
"""Test entity_category is packed correctly through the full setup flow."""
added_expressions = setup_test_environment
var = MockObj("sensor1")
config = {
@@ -951,9 +944,9 @@ async def test_setup_entity_with_entity_category(
CONF_ENTITY_CATEGORY: "diagnostic",
}
await _setup_entity_impl(var, config, "sensor")
assert any(
'set_entity_category("diagnostic")' in expr for expr in added_expressions
)
finalize_entity_strings(var, config)
packed = _extract_packed_value(added_expressions)
assert (packed >> _ENTITY_CATEGORY_SHIFT) & 0x3 == 2
@pytest.mark.asyncio
@@ -1002,3 +995,200 @@ async def test_setup_entity_decorator_mode(setup_test_environment: list[str]) ->
assert body_called
object_id = extract_object_id_from_expressions(added_expressions)
assert object_id == "temperature"
# Tests for finalize_entity_strings packing
def _extract_packed_value(expressions: list[str]) -> int:
"""Extract the third argument (packed value) from a configure_entity_() call."""
for expr in expressions:
if "configure_entity_" in expr:
# Match the last integer argument before the closing ");"
match = re.search(r",\s*(\d+)\s*\)", expr)
if match:
return int(match.group(1))
raise AssertionError("No configure_entity_ call found")
@pytest.mark.asyncio
async def test_finalize_no_flags(setup_test_environment: list[str]) -> None:
"""Test entity with no special flags — packed value is 0, no comment."""
added_expressions = setup_test_environment
var = MockObj("sensor1")
config = {
CONF_NAME: "Test",
CONF_DISABLED_BY_DEFAULT: False,
}
await _setup_entity_impl(var, config, "sensor")
finalize_entity_strings(var, config)
packed = _extract_packed_value(added_expressions)
assert packed == 0
assert "//" not in added_expressions[0]
@pytest.mark.asyncio
async def test_finalize_internal(setup_test_environment: list[str]) -> None:
"""Test entity with internal=True packs the internal bit."""
added_expressions = setup_test_environment
var = MockObj("sensor1")
config = {
CONF_NAME: "Test",
CONF_DISABLED_BY_DEFAULT: False,
CONF_INTERNAL: True,
}
await _setup_entity_impl(var, config, "sensor")
finalize_entity_strings(var, config)
packed = _extract_packed_value(added_expressions)
assert packed & (1 << _INTERNAL_SHIFT) != 0
assert packed == (1 << _INTERNAL_SHIFT)
@pytest.mark.asyncio
async def test_finalize_disabled_by_default(
setup_test_environment: list[str],
) -> None:
"""Test entity with disabled_by_default=True packs the bit."""
added_expressions = setup_test_environment
var = MockObj("sensor1")
config = {
CONF_NAME: "Test",
CONF_DISABLED_BY_DEFAULT: True,
}
await _setup_entity_impl(var, config, "sensor")
finalize_entity_strings(var, config)
packed = _extract_packed_value(added_expressions)
assert packed & (1 << _DISABLED_BY_DEFAULT_SHIFT) != 0
assert packed == (1 << _DISABLED_BY_DEFAULT_SHIFT)
@pytest.mark.asyncio
async def test_finalize_entity_category(
setup_test_environment: list[str],
) -> None:
"""Test entity_category values (diagnostic=2, config=1) are packed."""
added_expressions = setup_test_environment
var = MockObj("sensor1")
# Test diagnostic (value 2)
config = {
CONF_NAME: "Test",
CONF_DISABLED_BY_DEFAULT: False,
CONF_ENTITY_CATEGORY: "diagnostic",
}
await _setup_entity_impl(var, config, "sensor")
finalize_entity_strings(var, config)
packed = _extract_packed_value(added_expressions)
assert (packed >> _ENTITY_CATEGORY_SHIFT) & 0x3 == 2
# Test config (value 1)
added_expressions.clear()
config2 = {
CONF_NAME: "Test2",
CONF_DISABLED_BY_DEFAULT: False,
CONF_ENTITY_CATEGORY: "config",
}
await _setup_entity_impl(var, config2, "sensor")
finalize_entity_strings(var, config2)
packed = _extract_packed_value(added_expressions)
assert (packed >> _ENTITY_CATEGORY_SHIFT) & 0x3 == 1
@pytest.mark.asyncio
async def test_finalize_string_indices(
setup_test_environment: list[str],
) -> None:
"""Test device_class, unit_of_measurement, and icon are packed as indices."""
added_expressions = setup_test_environment
var = MockObj("sensor1")
config = {
CONF_NAME: "Test",
CONF_DISABLED_BY_DEFAULT: False,
CONF_DEVICE_CLASS: "temperature",
CONF_UNIT_OF_MEASUREMENT: "°C",
CONF_ICON: "mdi:thermometer",
}
await _setup_entity_impl(var, config, "sensor")
setup_device_class(config)
setup_unit_of_measurement(config)
finalize_entity_strings(var, config)
packed = _extract_packed_value(added_expressions)
# All three string indices should be non-zero
assert (packed >> _DC_SHIFT) & 0xFF != 0
assert (packed >> _UOM_SHIFT) & 0xFF != 0
assert (packed >> _ICON_SHIFT) & 0xFF != 0
# No flags set
assert (packed >> _INTERNAL_SHIFT) & 1 == 0
assert (packed >> _DISABLED_BY_DEFAULT_SHIFT) & 1 == 0
assert (packed >> _ENTITY_CATEGORY_SHIFT) & 0x3 == 0
@pytest.mark.asyncio
async def test_finalize_all_fields(
setup_test_environment: list[str],
) -> None:
"""Test all fields set: flags, string indices, and comment."""
added_expressions = setup_test_environment
var = MockObj("sensor1")
config = {
CONF_NAME: "Test",
CONF_DISABLED_BY_DEFAULT: True,
CONF_INTERNAL: True,
CONF_ENTITY_CATEGORY: "diagnostic",
CONF_DEVICE_CLASS: "temperature",
CONF_UNIT_OF_MEASUREMENT: "°C",
CONF_ICON: "mdi:thermometer",
}
await _setup_entity_impl(var, config, "sensor")
setup_device_class(config)
setup_unit_of_measurement(config)
finalize_entity_strings(var, config)
packed = _extract_packed_value(added_expressions)
# Verify flags
assert (packed >> _INTERNAL_SHIFT) & 1 == 1
assert (packed >> _DISABLED_BY_DEFAULT_SHIFT) & 1 == 1
assert (packed >> _ENTITY_CATEGORY_SHIFT) & 0x3 == 2
# Verify string indices are non-zero
assert (packed >> _DC_SHIFT) & 0xFF != 0
assert (packed >> _UOM_SHIFT) & 0xFF != 0
assert (packed >> _ICON_SHIFT) & 0xFF != 0
# Verify comment contains all flags with actual string values
comment_line = added_expressions[0]
assert (
"// internal, disabled_by_default, category:diagnostic,"
" dc:temperature, uom:°C, icon:mdi:thermometer" in comment_line
)
@pytest.mark.asyncio
async def test_finalize_comment_sanitization(
setup_test_environment: list[str],
) -> None:
"""Test that user strings in comments are sanitized against injection."""
added_expressions = setup_test_environment
var = MockObj("sensor1")
config = {
CONF_NAME: "Test",
CONF_DISABLED_BY_DEFAULT: False,
# Backslash at end would cause line splice eating next code line
CONF_ICON: "mdi:evil\\",
}
await _setup_entity_impl(var, config, "sensor")
finalize_entity_strings(var, config)
comment_line = added_expressions[0]
# Backslash must be replaced to prevent line splice
assert "\\" not in comment_line
assert "mdi:evil/" in comment_line
added_expressions.clear()
config2 = {
CONF_NAME: "Test2",
CONF_DISABLED_BY_DEFAULT: False,
CONF_ICON: "mdi:evil\nINJECTED_CODE();",
}
await _setup_entity_impl(var, config2, "sensor")
finalize_entity_strings(var, config2)
comment_line = added_expressions[0]
# Newline must be replaced to prevent breaking out of comment
assert "\n" not in comment_line
assert "INJECTED_CODE" in comment_line # still visible but safe in comment