From 466ec0b596e8d9141d94c00e9946bfcc2799796a Mon Sep 17 00:00:00 2001
From: Luke Lashley <conway220@gmail.com>
Date: Fri, 21 Mar 2025 08:31:17 -0400
Subject: [PATCH] Fix failing Roborock test (#141059)

Fix the falky test
---
 tests/components/roborock/conftest.py         | 11 +-------
 .../roborock/snapshots/test_diagnostics.ambr  |  2 +-
 tests/components/roborock/test_init.py        | 26 ++++++++++++++-----
 3 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/tests/components/roborock/conftest.py b/tests/components/roborock/conftest.py
index fcd469ca10a..332a9143c51 100644
--- a/tests/components/roborock/conftest.py
+++ b/tests/components/roborock/conftest.py
@@ -11,7 +11,6 @@ import uuid
 import pytest
 from roborock import RoborockCategory, RoomMapping
 from roborock.code_mappings import DyadError, RoborockDyadStateCode, ZeoError, ZeoState
-from roborock.containers import NetworkInfo
 from roborock.roborock_message import RoborockDyadDataProtocol, RoborockZeoProtocol
 from roborock.version_a01_apis import RoborockMqttClientA01
 
@@ -30,7 +29,6 @@ from .mock_data import (
     MAP_DATA,
     MULTI_MAP_LIST,
     NETWORK_INFO,
-    NETWORK_INFO_2,
     PROP,
     SCENES,
     USER_DATA,
@@ -89,13 +87,6 @@ def bypass_api_client_fixture() -> None:
         yield
 
 
-def cycle_network_info() -> Generator[NetworkInfo]:
-    """Return the appropriate network info for the corresponding device."""
-    while True:
-        yield NETWORK_INFO
-        yield NETWORK_INFO_2
-
-
 @pytest.fixture(name="bypass_api_fixture")
 def bypass_api_fixture(bypass_api_client_fixture: Any) -> None:
     """Skip calls to the API."""
@@ -107,7 +98,7 @@ def bypass_api_fixture(bypass_api_client_fixture: Any) -> None:
         ),
         patch(
             "homeassistant.components.roborock.RoborockMqttClientV1.get_networking",
-            side_effect=cycle_network_info(),
+            return_value=NETWORK_INFO,
         ),
         patch(
             "homeassistant.components.roborock.coordinator.RoborockLocalClientV1.get_prop",
diff --git a/tests/components/roborock/snapshots/test_diagnostics.ambr b/tests/components/roborock/snapshots/test_diagnostics.ambr
index 313824e70ec..26ecb729312 100644
--- a/tests/components/roborock/snapshots/test_diagnostics.ambr
+++ b/tests/components/roborock/snapshots/test_diagnostics.ambr
@@ -357,7 +357,7 @@
           }),
           'network_info': dict({
             'bssid': '**REDACTED**',
-            'ip': '123.232.12.2',
+            'ip': '123.232.12.1',
             'mac': '**REDACTED**',
             'rssi': 90,
             'ssid': 'wifi',
diff --git a/tests/components/roborock/test_init.py b/tests/components/roborock/test_init.py
index 226eea816b9..3d288b6479b 100644
--- a/tests/components/roborock/test_init.py
+++ b/tests/components/roborock/test_init.py
@@ -20,7 +20,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceRegistry
 from homeassistant.setup import async_setup_component
 
-from .mock_data import HOME_DATA, NETWORK_INFO
+from .mock_data import HOME_DATA, NETWORK_INFO, NETWORK_INFO_2
 
 from tests.common import MockConfigEntry
 from tests.typing import ClientSessionGenerator
@@ -305,7 +305,11 @@ async def test_stale_device(
     device_registry: DeviceRegistry,
 ) -> None:
     """Test that we remove a device if it no longer is given by home_data."""
-    await hass.config_entries.async_setup(mock_roborock_entry.entry_id)
+    with patch(
+        "homeassistant.components.roborock.RoborockMqttClientV1.get_networking",
+        side_effect=[NETWORK_INFO, NETWORK_INFO_2],
+    ):
+        await hass.config_entries.async_setup(mock_roborock_entry.entry_id)
     assert mock_roborock_entry.state is ConfigEntryState.LOADED
     existing_devices = device_registry.devices.get_devices_for_config_entry_id(
         mock_roborock_entry.entry_id
@@ -314,9 +318,15 @@ async def test_stale_device(
     hd = deepcopy(HOME_DATA)
     hd.devices = [hd.devices[0]]
 
-    with patch(
-        "homeassistant.components.roborock.RoborockApiClient.get_home_data_v2",
-        return_value=hd,
+    with (
+        patch(
+            "homeassistant.components.roborock.RoborockApiClient.get_home_data_v2",
+            return_value=hd,
+        ),
+        patch(
+            "homeassistant.components.roborock.RoborockMqttClientV1.get_networking",
+            side_effect=[NETWORK_INFO, NETWORK_INFO_2],
+        ),
     ):
         await hass.config_entries.async_reload(mock_roborock_entry.entry_id)
         await hass.async_block_till_done()
@@ -336,7 +346,11 @@ async def test_no_stale_device(
     device_registry: DeviceRegistry,
 ) -> None:
     """Test that we don't remove a device if fails to setup."""
-    await hass.config_entries.async_setup(mock_roborock_entry.entry_id)
+    with patch(
+        "homeassistant.components.roborock.RoborockMqttClientV1.get_networking",
+        side_effect=[NETWORK_INFO, NETWORK_INFO_2],
+    ):
+        await hass.config_entries.async_setup(mock_roborock_entry.entry_id)
     assert mock_roborock_entry.state is ConfigEntryState.LOADED
     existing_devices = device_registry.devices.get_devices_for_config_entry_id(
         mock_roborock_entry.entry_id