Do sanity check EntityPlatform.async_register_entity_service schema (#123058)

* Do a sanity check of schema passed to EntityPlatform.async_register_entity_service

* Only attempt to check schema of Schema

* Handle All/Any wrapped in schema

* Clarify comment

* Apply suggestions from code review

Co-authored-by: Robert Resch <robert@resch.dev>

---------

Co-authored-by: Robert Resch <robert@resch.dev>
pull/124038/head
Erik Montnemery 2024-08-16 10:01:12 +02:00 committed by GitHub
parent 6d9764185b
commit f9ade788eb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 39 additions and 0 deletions

View File

@ -10,6 +10,8 @@ from functools import partial
from logging import Logger, getLogger
from typing import TYPE_CHECKING, Any, Protocol
import voluptuous as vol
from homeassistant import config_entries
from homeassistant.const import (
ATTR_RESTORED,
@ -999,6 +1001,20 @@ class EntityPlatform:
if schema is None or isinstance(schema, dict):
schema = cv.make_entity_service_schema(schema)
# Do a sanity check to check this is a valid entity service schema,
# the check could be extended to require All/Any to have sub schema(s)
# with all entity service fields
elif (
# Don't check All/Any
not isinstance(schema, (vol.All, vol.Any))
# Don't check All/Any wrapped in schema
and not isinstance(schema.schema, (vol.All, vol.Any))
and any(key not in schema.schema for key in cv.ENTITY_SERVICE_FIELDS)
):
raise HomeAssistantError(
"The schema does not include all required keys: "
f"{", ".join(str(key) for key in cv.ENTITY_SERVICE_FIELDS)}"
)
service_func: str | HassJob[..., Any]
service_func = func if isinstance(func, str) else HassJob(func)

View File

@ -8,6 +8,7 @@ from typing import Any
from unittest.mock import ANY, AsyncMock, Mock, patch
import pytest
import voluptuous as vol
from homeassistant.const import EVENT_HOMEASSISTANT_STARTED, PERCENTAGE, EntityCategory
from homeassistant.core import (
@ -1788,6 +1789,28 @@ async def test_register_entity_service_none_schema(
assert entity2 in entities
async def test_register_entity_service_non_entity_service_schema(
hass: HomeAssistant,
) -> None:
"""Test attempting to register a service with an incomplete schema."""
entity_platform = MockEntityPlatform(
hass, domain="mock_integration", platform_name="mock_platform", platform=None
)
with pytest.raises(
HomeAssistantError,
match=(
"The schema does not include all required keys: entity_id, device_id, area_id, "
"floor_id, label_id"
),
):
entity_platform.async_register_entity_service(
"hello",
vol.Schema({"some": str}),
Mock(),
)
@pytest.mark.parametrize("update_before_add", [True, False])
async def test_invalid_entity_id(
hass: HomeAssistant, caplog: pytest.LogCaptureFixture, update_before_add: bool