From 467749eb579e8bd7e57ad96c677eee21485d10b8 Mon Sep 17 00:00:00 2001
From: Shay Levy <levyshay1@gmail.com>
Date: Tue, 27 Aug 2024 23:48:13 +0300
Subject: [PATCH] Fix Shelly sleepy RPC setup if device is already awake
 (#124734)

---
 homeassistant/components/shelly/__init__.py   |  5 ++
 .../components/shelly/coordinator.py          | 12 +++--
 tests/components/shelly/test_init.py          | 46 +++++++++++++++++++
 3 files changed, 58 insertions(+), 5 deletions(-)

diff --git a/homeassistant/components/shelly/__init__.py b/homeassistant/components/shelly/__init__.py
index 1d3f67220fa..e0d9d17d55d 100644
--- a/homeassistant/components/shelly/__init__.py
+++ b/homeassistant/components/shelly/__init__.py
@@ -290,6 +290,11 @@ async def _async_setup_rpc_entry(hass: HomeAssistant, entry: ShellyConfigEntry)
         )
         runtime_data.rpc = ShellyRpcCoordinator(hass, entry, device)
         runtime_data.rpc.async_setup(runtime_data.platforms)
+        # Try to connect to the device, if we reached here from config flow
+        # and user woke up the device when adding it, we can continue setup
+        # otherwise we will wait for the device to wake up
+        if sleep_period:
+            await runtime_data.rpc.async_device_online("setup")
     else:
         # Restore sensors for sleeping device
         LOGGER.debug("Setting up offline RPC device %s", entry.title)
diff --git a/homeassistant/components/shelly/coordinator.py b/homeassistant/components/shelly/coordinator.py
index 03dcdedbb6f..918dd920765 100644
--- a/homeassistant/components/shelly/coordinator.py
+++ b/homeassistant/components/shelly/coordinator.py
@@ -173,7 +173,7 @@ class ShellyCoordinatorBase[_DeviceT: BlockDevice | RpcDevice](
             await self.device.initialize()
             update_device_fw_info(self.hass, self.device, self.entry)
         except DeviceConnectionError as err:
-            LOGGER.error(
+            LOGGER.debug(
                 "Error connecting to Shelly device %s, error: %r", self.name, err
             )
             return False
@@ -480,15 +480,17 @@ class ShellyRpcCoordinator(ShellyCoordinatorBase[RpcDevice]):
         self._connect_task: asyncio.Task | None = None
         entry.async_on_unload(entry.add_update_listener(self._async_update_listener))
 
-    async def async_device_online(self) -> None:
+    async def async_device_online(self, source: str) -> None:
         """Handle device going online."""
         if not self.sleep_period:
             await self.async_request_refresh()
         elif not self._came_online_once or not self.device.initialized:
             LOGGER.debug(
-                "Sleepy device %s is online, trying to poll and configure", self.name
+                "Sleepy device %s is online (source: %s), trying to poll and configure",
+                self.name,
+                source,
             )
-            # Zeroconf told us the device is online, try to poll
+            # Source told us the device is online, try to poll
             # the device and if possible, set up the outbound
             # websocket so the device will send us updates
             # instead of relying on polling it fast enough before
@@ -847,7 +849,7 @@ async def async_reconnect_soon(hass: HomeAssistant, entry: ShellyConfigEntry) ->
     ):
         entry.async_create_background_task(
             hass,
-            coordinator.async_device_online(),
+            coordinator.async_device_online("zeroconf"),
             "reconnect soon",
             eager_start=True,
         )
diff --git a/tests/components/shelly/test_init.py b/tests/components/shelly/test_init.py
index 46698c23c0a..b5516485501 100644
--- a/tests/components/shelly/test_init.py
+++ b/tests/components/shelly/test_init.py
@@ -310,6 +310,52 @@ async def test_sleeping_rpc_device_online_new_firmware(
     assert entry.data["sleep_period"] == 1500
 
 
+async def test_sleeping_rpc_device_online_during_setup(
+    hass: HomeAssistant,
+    mock_rpc_device: Mock,
+    monkeypatch: pytest.MonkeyPatch,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test sleeping device Gen2 woke up by user during setup."""
+    monkeypatch.setattr(mock_rpc_device, "connected", False)
+    monkeypatch.setitem(mock_rpc_device.status["sys"], "wakeup_period", 1000)
+    await init_integration(hass, 2, sleep_period=1000)
+    await hass.async_block_till_done(wait_background_tasks=True)
+
+    assert "will resume when device is online" in caplog.text
+    assert "is online (source: setup)" in caplog.text
+    assert hass.states.get("sensor.test_name_temperature") is not None
+
+
+async def test_sleeping_rpc_device_offline_during_setup(
+    hass: HomeAssistant,
+    mock_rpc_device: Mock,
+    monkeypatch: pytest.MonkeyPatch,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test sleeping device Gen2 woke up by user during setup."""
+    monkeypatch.setattr(mock_rpc_device, "connected", False)
+    monkeypatch.setitem(mock_rpc_device.status["sys"], "wakeup_period", 1000)
+    monkeypatch.setattr(
+        mock_rpc_device, "initialize", AsyncMock(side_effect=DeviceConnectionError)
+    )
+
+    # Init integration, should fail since device is offline
+    await init_integration(hass, 2, sleep_period=1000)
+    await hass.async_block_till_done(wait_background_tasks=True)
+
+    assert "will resume when device is online" in caplog.text
+    assert "is online (source: setup)" in caplog.text
+    assert hass.states.get("sensor.test_name_temperature") is None
+
+    # Create an online event and verify that device is init successfully
+    monkeypatch.setattr(mock_rpc_device, "initialize", AsyncMock())
+    mock_rpc_device.mock_online()
+    await hass.async_block_till_done(wait_background_tasks=True)
+
+    assert hass.states.get("sensor.test_name_temperature") is not None
+
+
 @pytest.mark.parametrize(
     ("gen", "entity_id"),
     [