From 087eb86e3749dd2a206bc8c1e3f330f66e235c60 Mon Sep 17 00:00:00 2001
From: Martin Hjelmare <marhje52@gmail.com>
Date: Fri, 22 Dec 2023 19:48:07 +0100
Subject: [PATCH] Improve cloud binary sensor tests (#106238)

* Clean up cloud binary sensor test

* Test remove entity
---
 tests/components/cloud/test_binary_sensor.py | 76 +++++++++++++-------
 1 file changed, 50 insertions(+), 26 deletions(-)

diff --git a/tests/components/cloud/test_binary_sensor.py b/tests/components/cloud/test_binary_sensor.py
index 7b090bd5eca..6505be1fe10 100644
--- a/tests/components/cloud/test_binary_sensor.py
+++ b/tests/components/cloud/test_binary_sensor.py
@@ -1,44 +1,68 @@
 """Tests for the cloud binary sensor."""
-from unittest.mock import Mock, patch
+from collections.abc import Generator
+from unittest.mock import MagicMock, patch
+
+from hass_nabucasa.const import DISPATCH_REMOTE_CONNECT, DISPATCH_REMOTE_DISCONNECT
+import pytest
 
-from homeassistant.components.cloud.const import DISPATCHER_REMOTE_UPDATE
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.discovery import async_load_platform
-from homeassistant.helpers.dispatcher import async_dispatcher_send
+from homeassistant.helpers.entity_registry import EntityRegistry
 from homeassistant.setup import async_setup_component
 
 
-async def test_remote_connection_sensor(hass: HomeAssistant) -> None:
+@pytest.fixture(autouse=True)
+def mock_wait_until() -> Generator[None, None, None]:
+    """Mock WAIT_UNTIL_CHANGE to execute callback immediately."""
+    with patch("homeassistant.components.cloud.binary_sensor.WAIT_UNTIL_CHANGE", 0):
+        yield
+
+
+async def test_remote_connection_sensor(
+    hass: HomeAssistant,
+    cloud: MagicMock,
+    entity_registry: EntityRegistry,
+) -> None:
     """Test the remote connection sensor."""
+    entity_id = "binary_sensor.remote_ui"
+    cloud.remote.certificate = None
+
     assert await async_setup_component(hass, "cloud", {"cloud": {}})
     await hass.async_block_till_done()
 
-    assert hass.states.get("binary_sensor.remote_ui") is None
+    assert hass.states.get(entity_id) is None
 
-    # Fake connection/discovery
-    await async_load_platform(hass, "binary_sensor", "cloud", {}, {"cloud": {}})
+    on_start_callback = cloud.register_on_start.call_args[0][0]
+    await on_start_callback()
 
-    # Mock test env
-    cloud = hass.data["cloud"] = Mock()
-    cloud.remote.certificate = None
-    await hass.async_block_till_done()
-
-    state = hass.states.get("binary_sensor.remote_ui")
+    state = hass.states.get(entity_id)
     assert state is not None
     assert state.state == "unavailable"
 
-    with patch("homeassistant.components.cloud.binary_sensor.WAIT_UNTIL_CHANGE", 0):
-        cloud.remote.is_connected = False
-        cloud.remote.certificate = object()
-        async_dispatcher_send(hass, DISPATCHER_REMOTE_UPDATE, {})
-        await hass.async_block_till_done()
+    cloud.remote.is_connected = False
+    cloud.remote.certificate = object()
+    cloud.client.dispatcher_message(DISPATCH_REMOTE_DISCONNECT)
+    await hass.async_block_till_done()
 
-        state = hass.states.get("binary_sensor.remote_ui")
-        assert state.state == "off"
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == "off"
 
-        cloud.remote.is_connected = True
-        async_dispatcher_send(hass, DISPATCHER_REMOTE_UPDATE, {})
-        await hass.async_block_till_done()
+    cloud.remote.is_connected = True
+    cloud.client.dispatcher_message(DISPATCH_REMOTE_CONNECT)
+    await hass.async_block_till_done()
 
-        state = hass.states.get("binary_sensor.remote_ui")
-        assert state.state == "on"
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == "on"
+
+    # Test that a state is not set if the entity is removed.
+    entity_registry.async_remove(entity_id)
+    await hass.async_block_till_done()
+
+    assert hass.states.get(entity_id) is None
+
+    cloud.remote.is_connected = False
+    cloud.client.dispatcher_message(DISPATCH_REMOTE_DISCONNECT)
+    await hass.async_block_till_done()
+
+    assert hass.states.get(entity_id) is None