From 14f4f8aeb59481776525663f75ddf4ec0f3a9cd3 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Mon, 16 Dec 2024 15:37:29 +0100
Subject: [PATCH] Update hassio backup agents on mount added or removed
 (#133344)

* Update hassio backup agents on mount added or removed

* Address review comments
---
 homeassistant/components/hassio/backup.py | 34 +++++++++++++
 tests/components/conftest.py              |  3 ++
 tests/components/hassio/test_backup.py    | 62 +++++++++++++++++++++++
 3 files changed, 99 insertions(+)

diff --git a/homeassistant/components/hassio/backup.py b/homeassistant/components/hassio/backup.py
index e544a56a3c8..0353255fe7b 100644
--- a/homeassistant/components/hassio/backup.py
+++ b/homeassistant/components/hassio/backup.py
@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import asyncio
 from collections.abc import AsyncIterator, Callable, Coroutine, Mapping
+import logging
 from pathlib import Path
 from typing import Any, cast
 
@@ -32,6 +33,8 @@ from .const import DOMAIN, EVENT_SUPERVISOR_EVENT
 from .handler import get_supervisor_client
 
 LOCATION_CLOUD_BACKUP = ".cloud_backup"
+MOUNT_JOBS = ("mount_manager_create_mount", "mount_manager_remove_mount")
+_LOGGER = logging.getLogger(__name__)
 
 
 async def async_get_backup_agents(
@@ -49,6 +52,37 @@ async def async_get_backup_agents(
     return agents
 
 
+@callback
+def async_register_backup_agents_listener(
+    hass: HomeAssistant,
+    *,
+    listener: Callable[[], None],
+    **kwargs: Any,
+) -> Callable[[], None]:
+    """Register a listener to be called when agents are added or removed."""
+
+    @callback
+    def unsub() -> None:
+        """Unsubscribe from job events."""
+        unsub_signal()
+
+    @callback
+    def handle_signal(data: Mapping[str, Any]) -> None:
+        """Handle a job signal."""
+        if (
+            data.get("event") != "job"
+            or not (event_data := data.get("data"))
+            or event_data.get("name") not in MOUNT_JOBS
+            or event_data.get("done") is not True
+        ):
+            return
+        _LOGGER.debug("Mount added or removed %s, calling listener", data)
+        listener()
+
+    unsub_signal = async_dispatcher_connect(hass, EVENT_SUPERVISOR_EVENT, handle_signal)
+    return unsub
+
+
 def _backup_details_to_agent_backup(
     details: supervisor_backups.BackupComplete,
 ) -> AgentBackup:
diff --git a/tests/components/conftest.py b/tests/components/conftest.py
index ac30d105299..3828cc5ff37 100644
--- a/tests/components/conftest.py
+++ b/tests/components/conftest.py
@@ -514,11 +514,14 @@ def resolution_suggestions_for_issue_fixture(supervisor_client: AsyncMock) -> As
 @pytest.fixture(name="supervisor_client")
 def supervisor_client() -> Generator[AsyncMock]:
     """Mock the supervisor client."""
+    mounts_info_mock = AsyncMock(spec_set=["mounts"])
+    mounts_info_mock.mounts = []
     supervisor_client = AsyncMock()
     supervisor_client.addons = AsyncMock()
     supervisor_client.discovery = AsyncMock()
     supervisor_client.homeassistant = AsyncMock()
     supervisor_client.host = AsyncMock()
+    supervisor_client.mounts.info.return_value = mounts_info_mock
     supervisor_client.os = AsyncMock()
     supervisor_client.resolution = AsyncMock()
     supervisor_client.supervisor = AsyncMock()
diff --git a/tests/components/hassio/test_backup.py b/tests/components/hassio/test_backup.py
index 660753bd815..3e928bc996b 100644
--- a/tests/components/hassio/test_backup.py
+++ b/tests/components/hassio/test_backup.py
@@ -231,6 +231,68 @@ async def test_agent_delete_backup(
     supervisor_client.backups.remove_backup.assert_called_once_with(backup_id)
 
 
+@pytest.mark.usefixtures("hassio_client")
+@pytest.mark.parametrize(
+    ("event_data", "mount_info_calls"),
+    [
+        (
+            {
+                "event": "job",
+                "data": {"name": "mount_manager_create_mount", "done": True},
+            },
+            1,
+        ),
+        (
+            {
+                "event": "job",
+                "data": {"name": "mount_manager_create_mount", "done": False},
+            },
+            0,
+        ),
+        (
+            {
+                "event": "job",
+                "data": {"name": "mount_manager_remove_mount", "done": True},
+            },
+            1,
+        ),
+        (
+            {
+                "event": "job",
+                "data": {"name": "mount_manager_remove_mount", "done": False},
+            },
+            0,
+        ),
+        ({"event": "job", "data": {"name": "other_job", "done": True}}, 0),
+        (
+            {
+                "event": "other_event",
+                "data": {"name": "mount_manager_remove_mount", "done": True},
+            },
+            0,
+        ),
+    ],
+)
+async def test_agents_notify_on_mount_added_removed(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    supervisor_client: AsyncMock,
+    event_data: dict[str, Any],
+    mount_info_calls: int,
+) -> None:
+    """Test the listener is called when mounts are added or removed."""
+    client = await hass_ws_client(hass)
+    assert supervisor_client.mounts.info.call_count == 1
+    assert supervisor_client.mounts.info.call_args[0] == ()
+    supervisor_client.mounts.info.reset_mock()
+
+    await client.send_json_auto_id({"type": "supervisor/event", "data": event_data})
+    response = await client.receive_json()
+    assert response["success"]
+    await hass.async_block_till_done()
+    assert supervisor_client.mounts.info.call_count == mount_info_calls
+
+
 @pytest.mark.usefixtures("hassio_client")
 async def test_reader_writer_create(
     hass: HomeAssistant,