Make reproduce state use platform instead of rely on function (#25856)

* Make reproduce state use platform instead of rely on function

* Fix types

* address comment Martin.

homeassistant/components/climate/__init__.py

@@ -70,7 +70,6 @@ from .const import (
     SUPPORT_TARGET_TEMPERATURE_RANGE,
     SUPPORT_TARGET_TEMPERATURE,
 )
-from .reproduce_state import async_reproduce_states  # noqa
 
 DEFAULT_MIN_TEMP = 7
 DEFAULT_MAX_TEMP = 35

homeassistant/components/climate/reproduce_state.py

@@ -5,7 +5,6 @@ from typing import Iterable, Optional
 from homeassistant.const import ATTR_TEMPERATURE
 from homeassistant.core import Context, State
 from homeassistant.helpers.typing import HomeAssistantType
-from homeassistant.loader import bind_hass
 
 from .const import (
     ATTR_AUX_HEAT,
@@ -69,7 +68,6 @@ async def _async_reproduce_states(
         await call_service(SERVICE_SET_HUMIDITY, [ATTR_HUMIDITY])
 
 
-@bind_hass
 async def async_reproduce_states(
     hass: HomeAssistantType, states: Iterable[State], context: Optional[Context] = None
 ) -> None:

homeassistant/components/group/__init__.py

@@ -34,7 +34,6 @@ import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.config_validation import ENTITY_SERVICE_SCHEMA
 from homeassistant.util.async_ import run_coroutine_threadsafe
 
-from .reproduce_state import async_reproduce_states  # noqa
 
 DOMAIN = "group"
 

homeassistant/components/group/reproduce_state.py

@@ -3,10 +3,8 @@ from typing import Iterable, Optional
 
 from homeassistant.core import Context, State
 from homeassistant.helpers.typing import HomeAssistantType
-from homeassistant.loader import bind_hass
 
 
-@bind_hass
 async def async_reproduce_states(
     hass: HomeAssistantType, states: Iterable[State], context: Optional[Context] = None
 ) -> None:

homeassistant/components/media_player/__init__.py

@@ -96,7 +96,6 @@ from .const import (
     SUPPORT_VOLUME_SET,
     SUPPORT_VOLUME_STEP,
 )
-from .reproduce_state import async_reproduce_states  # noqa
 
 _LOGGER = logging.getLogger(__name__)
 _RND = SystemRandom()

homeassistant/components/media_player/reproduce_state.py

@@ -19,7 +19,6 @@ from homeassistant.const import (
 )
 from homeassistant.core import Context, State
 from homeassistant.helpers.typing import HomeAssistantType
-from homeassistant.loader import bind_hass
 
 from .const import (
     ATTR_MEDIA_VOLUME_LEVEL,
@@ -89,7 +88,6 @@ async def _async_reproduce_states(
         )
 
 
-@bind_hass
 async def async_reproduce_states(
     hass: HomeAssistantType, states: Iterable[State], context: Optional[Context] = None
 ) -> None:

homeassistant/helpers/state.py

@@ -4,7 +4,7 @@ import datetime as dt
 import json
 import logging
 from collections import defaultdict
-from types import TracebackType
+from types import ModuleType, TracebackType
 from typing import (  # noqa: F401 pylint: disable=unused-import
     Awaitable,
     Dict,
@@ -16,7 +16,7 @@ from typing import (  # noqa: F401 pylint: disable=unused-import
     Union,
 )
 
-from homeassistant.loader import bind_hass
+from homeassistant.loader import bind_hass, async_get_integration, IntegrationNotFound
 import homeassistant.util.dt as dt_util
 from homeassistant.components.notify import ATTR_MESSAGE, SERVICE_NOTIFY
 from homeassistant.components.sun import STATE_ABOVE_HORIZON, STATE_BELOW_HORIZON
@@ -152,13 +152,27 @@ async def async_reproduce_state(
     for state in states:
         to_call[state.domain].append(state)
 
-    async def worker(domain: str, data: List[State]) -> None:
-        component = getattr(hass.components, domain)
-        if hasattr(component, "async_reproduce_states"):
-            await component.async_reproduce_states(data, context=context)
+    async def worker(domain: str, states_by_domain: List[State]) -> None:
+        try:
+            integration = await async_get_integration(hass, domain)
+        except IntegrationNotFound:
+            _LOGGER.warning(
+                "Trying to reproduce state for unknown integration: %s", domain
+            )
+            return
+
+        try:
+            platform: Optional[ModuleType] = integration.get_platform("reproduce_state")
+        except ImportError:
+            platform = None
+
+        if platform:
+            await platform.async_reproduce_states(  # type: ignore
+                hass, states_by_domain, context=context
+            )
         else:
             await async_reproduce_state_legacy(
-                hass, domain, data, blocking=blocking, context=context
+                hass, domain, states_by_domain, blocking=blocking, context=context
             )
 
     if to_call:

tests/components/climate/test_reproduce_state.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from homeassistant.components.climate import async_reproduce_states
+from homeassistant.components.climate.reproduce_state import async_reproduce_states
 from homeassistant.components.climate.const import (
     ATTR_AUX_HEAT,
     ATTR_HUMIDITY,

tests/components/group/test_reproduce_state.py

@@ -2,7 +2,7 @@
 
 from asyncio import Future
 from unittest.mock import patch
-from homeassistant.components.group import async_reproduce_states
+from homeassistant.components.group.reproduce_state import async_reproduce_states
 from homeassistant.core import Context, State
 
 

tests/components/media_player/test_reproduce_state.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from homeassistant.components.media_player import async_reproduce_states
+from homeassistant.components.media_player.reproduce_state import async_reproduce_states
 from homeassistant.components.media_player.const import (
     ATTR_INPUT_SOURCE,
     ATTR_MEDIA_CONTENT_ID,

tests/helpers/test_state.py

@@ -53,13 +53,13 @@ def test_async_track_states(hass):
 def test_call_to_component(hass):
     """Test calls to components state reproduction functions."""
     with patch(
-        ("homeassistant.components.media_player." "async_reproduce_states")
+        ("homeassistant.components.media_player.reproduce_state.async_reproduce_states")
     ) as media_player_fun:
         media_player_fun.return_value = asyncio.Future()
         media_player_fun.return_value.set_result(None)
 
         with patch(
-            ("homeassistant.components.climate." "async_reproduce_states")
+            ("homeassistant.components.climate.reproduce_state.async_reproduce_states")
         ) as climate_fun:
             climate_fun.return_value = asyncio.Future()
             climate_fun.return_value.set_result(None)