From 06513e88c62131ebaa0d52662b9c8d7f62f8391c Mon Sep 17 00:00:00 2001
From: Mike Degatano <michael.degatano@gmail.com>
Date: Wed, 17 Apr 2024 02:54:56 -0400
Subject: [PATCH] Allow restarting core in safe mode (#5017)

---
 supervisor/api/const.py           |  1 +
 supervisor/api/homeassistant.py   | 15 +++++++++++++--
 supervisor/homeassistant/const.py |  2 ++
 supervisor/homeassistant/core.py  |  9 ++++++++-
 tests/api/test_homeassistant.py   | 23 +++++++++++++++++++++++
 5 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/supervisor/api/const.py b/supervisor/api/const.py
index 437411b3b..2b7b2cfbf 100644
--- a/supervisor/api/const.py
+++ b/supervisor/api/const.py
@@ -53,6 +53,7 @@ ATTR_PANEL_PATH = "panel_path"
 ATTR_REMOVABLE = "removable"
 ATTR_REMOVE_CONFIG = "remove_config"
 ATTR_REVISION = "revision"
+ATTR_SAFE_MODE = "safe_mode"
 ATTR_SEAT = "seat"
 ATTR_SIGNED = "signed"
 ATTR_STARTUP_TIME = "startup_time"
diff --git a/supervisor/api/homeassistant.py b/supervisor/api/homeassistant.py
index de6cb13e1..1ced92f74 100644
--- a/supervisor/api/homeassistant.py
+++ b/supervisor/api/homeassistant.py
@@ -36,6 +36,7 @@ from ..const import (
 from ..coresys import CoreSysAttributes
 from ..exceptions import APIError
 from ..validate import docker_image, network_port, version_tag
+from .const import ATTR_SAFE_MODE
 from .utils import api_process, api_validate
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
@@ -62,6 +63,12 @@ SCHEMA_UPDATE = vol.Schema(
     }
 )
 
+SCHEMA_RESTART = vol.Schema(
+    {
+        vol.Optional(ATTR_SAFE_MODE, default=False): vol.Boolean(),
+    }
+)
+
 
 class APIHomeAssistant(CoreSysAttributes):
     """Handle RESTful API for Home Assistant functions."""
@@ -166,9 +173,13 @@ class APIHomeAssistant(CoreSysAttributes):
         return asyncio.shield(self.sys_homeassistant.core.start())
 
     @api_process
-    def restart(self, request: web.Request) -> Awaitable[None]:
+    async def restart(self, request: web.Request) -> None:
         """Restart Home Assistant."""
-        return asyncio.shield(self.sys_homeassistant.core.restart())
+        body = await api_validate(SCHEMA_RESTART, request)
+
+        await asyncio.shield(
+            self.sys_homeassistant.core.restart(safe_mode=body[ATTR_SAFE_MODE])
+        )
 
     @api_process
     def rebuild(self, request: web.Request) -> Awaitable[None]:
diff --git a/supervisor/homeassistant/const.py b/supervisor/homeassistant/const.py
index 1a7577470..ee34c5163 100644
--- a/supervisor/homeassistant/const.py
+++ b/supervisor/homeassistant/const.py
@@ -1,6 +1,7 @@
 """Constants for homeassistant."""
 from datetime import timedelta
 from enum import StrEnum
+from pathlib import PurePath
 
 from awesomeversion import AwesomeVersion
 
@@ -12,6 +13,7 @@ WATCHDOG_RETRY_SECONDS = 10
 WATCHDOG_MAX_ATTEMPTS = 5
 WATCHDOG_THROTTLE_PERIOD = timedelta(minutes=30)
 WATCHDOG_THROTTLE_MAX_CALLS = 10
+SAFE_MODE_FILENAME = PurePath("safe-mode")
 
 CLOSING_STATES = [
     CoreState.SHUTDOWN,
diff --git a/supervisor/homeassistant/core.py b/supervisor/homeassistant/core.py
index ee98b51aa..a73856701 100644
--- a/supervisor/homeassistant/core.py
+++ b/supervisor/homeassistant/core.py
@@ -35,6 +35,7 @@ from ..utils import convert_to_ascii
 from ..utils.sentry import capture_exception
 from .const import (
     LANDINGPAGE,
+    SAFE_MODE_FILENAME,
     WATCHDOG_MAX_ATTEMPTS,
     WATCHDOG_RETRY_SECONDS,
     WATCHDOG_THROTTLE_MAX_CALLS,
@@ -362,8 +363,14 @@ class HomeAssistantCore(JobGroup):
         limit=JobExecutionLimit.GROUP_ONCE,
         on_condition=HomeAssistantJobError,
     )
-    async def restart(self) -> None:
+    async def restart(self, *, safe_mode: bool = False) -> None:
         """Restart Home Assistant Docker."""
+        # Create safe mode marker file if necessary
+        if safe_mode:
+            await self.sys_run_in_executor(
+                (self.sys_config.path_homeassistant / SAFE_MODE_FILENAME).touch
+            )
+
         try:
             await self.instance.restart()
         except DockerError as err:
diff --git a/tests/api/test_homeassistant.py b/tests/api/test_homeassistant.py
index 19958864f..78b647c00 100644
--- a/tests/api/test_homeassistant.py
+++ b/tests/api/test_homeassistant.py
@@ -1,11 +1,13 @@
 """Test homeassistant api."""
 
+from pathlib import Path
 from unittest.mock import MagicMock, patch
 
 from aiohttp.test_utils import TestClient
 import pytest
 
 from supervisor.coresys import CoreSys
+from supervisor.homeassistant.core import HomeAssistantCore
 from supervisor.homeassistant.module import HomeAssistant
 
 from tests.api import common_test_api_advanced_logs
@@ -92,3 +94,24 @@ async def test_api_set_image(api_client: TestClient, coresys: CoreSys):
         coresys.homeassistant.image == "ghcr.io/home-assistant/qemux86-64-homeassistant"
     )
     assert coresys.homeassistant.override_image is False
+
+
+async def test_api_restart(
+    api_client: TestClient,
+    container: MagicMock,
+    tmp_supervisor_data: Path,
+):
+    """Test restarting homeassistant."""
+    safe_mode_marker = tmp_supervisor_data / "homeassistant" / "safe-mode"
+
+    with patch.object(HomeAssistantCore, "_block_till_run"):
+        await api_client.post("/homeassistant/restart")
+
+    container.restart.assert_called_once()
+    assert not safe_mode_marker.exists()
+
+    with patch.object(HomeAssistantCore, "_block_till_run"):
+        await api_client.post("/homeassistant/restart", json={"safe_mode": True})
+
+    assert container.restart.call_count == 2
+    assert safe_mode_marker.exists()