Add reproduce state for Number (#43870)

homeassistant/components/number/reproduce_state.py

@@ -0,0 +1,65 @@
+"""Reproduce a Number entity state."""
+import asyncio
+import logging
+from typing import Any, Dict, Iterable, Optional
+
+from homeassistant.const import ATTR_ENTITY_ID
+from homeassistant.core import Context, State
+from homeassistant.helpers.typing import HomeAssistantType
+
+from . import ATTR_VALUE, DOMAIN, SERVICE_SET_VALUE
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def _async_reproduce_state(
+    hass: HomeAssistantType,
+    state: State,
+    *,
+    context: Optional[Context] = None,
+    reproduce_options: Optional[Dict[str, Any]] = None,
+) -> None:
+    """Reproduce a single state."""
+    cur_state = hass.states.get(state.entity_id)
+
+    if cur_state is None:
+        _LOGGER.warning("Unable to find entity %s", state.entity_id)
+        return
+
+    try:
+        float(state.state)
+    except ValueError:
+        _LOGGER.warning(
+            "Invalid state specified for %s: %s", state.entity_id, state.state
+        )
+        return
+
+    # Return if we are already at the right state.
+    if cur_state.state == state.state:
+        return
+
+    service = SERVICE_SET_VALUE
+    service_data = {ATTR_ENTITY_ID: state.entity_id, ATTR_VALUE: state.state}
+
+    await hass.services.async_call(
+        DOMAIN, service, service_data, context=context, blocking=True
+    )
+
+
+async def async_reproduce_states(
+    hass: HomeAssistantType,
+    states: Iterable[State],
+    *,
+    context: Optional[Context] = None,
+    reproduce_options: Optional[Dict[str, Any]] = None,
+) -> None:
+    """Reproduce multiple Number states."""
+    # Reproduce states in parallel.
+    await asyncio.gather(
+        *(
+            _async_reproduce_state(
+                hass, state, context=context, reproduce_options=reproduce_options
+            )
+            for state in states
+        )
+    )

tests/components/number/test_reproduce_state.py

@@ -0,0 +1,53 @@
+"""Test reproduce state for Number entities."""
+from homeassistant.components.number.const import (
+    ATTR_MAX,
+    ATTR_MIN,
+    DOMAIN,
+    SERVICE_SET_VALUE,
+)
+from homeassistant.core import State
+
+from tests.common import async_mock_service
+
+VALID_NUMBER1 = "19.0"
+VALID_NUMBER2 = "99.9"
+
+
+async def test_reproducing_states(hass, caplog):
+    """Test reproducing Number states."""
+
+    hass.states.async_set(
+        "number.test_number", VALID_NUMBER1, {ATTR_MIN: 5, ATTR_MAX: 100}
+    )
+
+    # These calls should do nothing as entities already in desired state
+    await hass.helpers.state.async_reproduce_state(
+        [
+            State("number.test_number", VALID_NUMBER1),
+            # Should not raise
+            State("number.non_existing", "234"),
+        ],
+    )
+
+    assert hass.states.get("number.test_number").state == VALID_NUMBER1
+
+    # Test reproducing with different state
+    calls = async_mock_service(hass, DOMAIN, SERVICE_SET_VALUE)
+    await hass.helpers.state.async_reproduce_state(
+        [
+            State("number.test_number", VALID_NUMBER2),
+            # Should not raise
+            State("number.non_existing", "234"),
+        ],
+    )
+
+    assert len(calls) == 1
+    assert calls[0].domain == DOMAIN
+    assert calls[0].data == {"entity_id": "number.test_number", "value": VALID_NUMBER2}
+
+    # Test invalid state
+    await hass.helpers.state.async_reproduce_state(
+        [State("number.test_number", "invalid_state")]
+    )
+
+    assert len(calls) == 1