Add env on core restart due to restore (#5548)

* Add env on core restart due to restore

* Move is_restore to backup manager

supervisor/backups/backup.py

@@ -772,7 +772,7 @@ class Backup(JobGroup):
     @Job(name="backup_restore_homeassistant", cleanup=False)
     async def restore_homeassistant(self) -> Awaitable[None]:
         """Restore Home Assistant Core configuration folder."""
-        await self.sys_homeassistant.core.stop()
+        await self.sys_homeassistant.core.stop(remove_container=True)
 
         # Restore Home Assistant Core config directory
         tar_name = Path(

supervisor/backups/const.py

@@ -39,7 +39,6 @@ class RestoreJobStage(StrEnum):
     ADDONS = "addons"
     AWAIT_ADDON_RESTARTS = "await_addon_restarts"
     AWAIT_HOME_ASSISTANT_RESTART = "await_home_assistant_restart"
-    CHECK_HOME_ASSISTANT = "check_home_assistant"
     DOCKER_CONFIG = "docker_config"
     FOLDERS = "folders"
     HOME_ASSISTANT = "home_assistant"

supervisor/backups/manager.py

@@ -47,6 +47,9 @@ from .validate import ALL_FOLDERS, SCHEMA_BACKUPS_CONFIG
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
+JOB_FULL_RESTORE = "backup_manager_full_restore"
+JOB_PARTIAL_RESTORE = "backup_manager_partial_restore"
+
 
 class BackupManager(FileConfiguration, JobGroup):
     """Manage backups."""
@@ -86,6 +89,16 @@ class BackupManager(FileConfiguration, JobGroup):
             if mount.state == UnitActiveState.ACTIVE
         }
 
+    @property
+    def current_restore(self) -> str | None:
+        """Return id of current restore job if a restore job is in progress."""
+        job = self.sys_jobs.current
+        while job.parent_id:
+            job = self.sys_jobs.get_job(job.parent_id)
+            if job.name in {JOB_FULL_RESTORE, JOB_PARTIAL_RESTORE}:
+                return job.uuid
+        return None
+
     def get(self, slug: str) -> Backup:
         """Return backup object."""
         return self._backups.get(slug)
@@ -619,9 +632,6 @@ class BackupManager(FileConfiguration, JobGroup):
 
                 # Wait for Home Assistant Core update/downgrade
                 if task_hass:
-                    self._change_stage(
-                        RestoreJobStage.AWAIT_HOME_ASSISTANT_RESTART, backup
-                    )
                     await task_hass
         except BackupError:
             raise
@@ -644,7 +654,7 @@ class BackupManager(FileConfiguration, JobGroup):
         finally:
             # Leave Home Assistant alone if it wasn't part of the restore
             if homeassistant:
-                self._change_stage(RestoreJobStage.CHECK_HOME_ASSISTANT, backup)
+                self._change_stage(RestoreJobStage.AWAIT_HOME_ASSISTANT_RESTART, backup)
 
                 # Do we need start Home Assistant Core?
                 if not await self.sys_homeassistant.core.is_running():
@@ -660,7 +670,7 @@ class BackupManager(FileConfiguration, JobGroup):
                     )
 
     @Job(
-        name="backup_manager_full_restore",
+        name=JOB_FULL_RESTORE,
         conditions=[
             JobCondition.FREE_SPACE,
             JobCondition.HEALTHY,
@@ -706,7 +716,7 @@ class BackupManager(FileConfiguration, JobGroup):
 
         try:
             # Stop Home-Assistant / Add-ons
-            await self.sys_core.shutdown()
+            await self.sys_core.shutdown(remove_homeassistant_container=True)
 
             success = await self._do_restore(
                 backup,
@@ -724,7 +734,7 @@ class BackupManager(FileConfiguration, JobGroup):
         return success
 
     @Job(
-        name="backup_manager_partial_restore",
+        name=JOB_PARTIAL_RESTORE,
         conditions=[
             JobCondition.FREE_SPACE,
             JobCondition.HEALTHY,

supervisor/core.py

@@ -333,7 +333,7 @@ class Core(CoreSysAttributes):
         _LOGGER.info("Supervisor is down - %d", self.exit_code)
         self.sys_loop.stop()
 
-    async def shutdown(self):
+    async def shutdown(self, *, remove_homeassistant_container: bool = False):
         """Shutdown all running containers in correct order."""
         # don't process scheduler anymore
         if self.state == CoreState.RUNNING:
@@ -344,7 +344,9 @@ class Core(CoreSysAttributes):
 
         # Close Home Assistant
         with suppress(HassioError):
-            await self.sys_homeassistant.core.stop()
+            await self.sys_homeassistant.core.stop(
+                remove_container=remove_homeassistant_container
+            )
 
         # Shutdown System Add-ons
         await self.sys_addons.shutdown(AddonStartup.SERVICES)

supervisor/docker/homeassistant.py

@@ -35,6 +35,7 @@ _LOGGER: logging.Logger = logging.getLogger(__name__)
 _VERIFY_TRUST: AwesomeVersion = AwesomeVersion("2021.5.0")
 _HASS_DOCKER_NAME: str = "homeassistant"
 ENV_S6_GRACETIME = re.compile(r"^S6_SERVICES_GRACETIME=([0-9]+)$")
+ENV_RESTORE_JOB_ID = "SUPERVISOR_RESTORE_JOB_ID"
 
 
 class DockerHomeAssistant(DockerInterface):
@@ -163,8 +164,17 @@ class DockerHomeAssistant(DockerInterface):
         limit=JobExecutionLimit.GROUP_ONCE,
         on_condition=DockerJobError,
     )
-    async def run(self) -> None:
+    async def run(self, *, restore_job_id: str | None = None) -> None:
         """Run Docker image."""
+        environment = {
+            "SUPERVISOR": self.sys_docker.network.supervisor,
+            "HASSIO": self.sys_docker.network.supervisor,
+            ENV_TIME: self.sys_timezone,
+            ENV_TOKEN: self.sys_homeassistant.supervisor_token,
+            ENV_TOKEN_OLD: self.sys_homeassistant.supervisor_token,
+        }
+        if restore_job_id:
+            environment[ENV_RESTORE_JOB_ID] = restore_job_id
         await self._run(
             tag=(self.sys_homeassistant.version),
             name=self.name,
@@ -180,13 +190,7 @@ class DockerHomeAssistant(DockerInterface):
                 "supervisor": self.sys_docker.network.supervisor,
                 "observer": self.sys_docker.network.observer,
             },
-            environment={
+            environment=environment,
-                "SUPERVISOR": self.sys_docker.network.supervisor,
-                "HASSIO": self.sys_docker.network.supervisor,
-                ENV_TIME: self.sys_timezone,
-                ENV_TOKEN: self.sys_homeassistant.supervisor_token,
-                ENV_TOKEN_OLD: self.sys_homeassistant.supervisor_token,
-            },
             tmpfs={"/tmp": ""},  # noqa: S108
             oom_score_adj=-300,
         )

supervisor/homeassistant/core.py

@@ -345,7 +345,7 @@ class HomeAssistantCore(JobGroup):
             self.sys_homeassistant.write_pulse()
 
             try:
-                await self.instance.run()
+                await self.instance.run(restore_job_id=self.sys_backups.current_restore)
             except DockerError as err:
                 raise HomeAssistantError() from err
 
@@ -356,10 +356,10 @@ class HomeAssistantCore(JobGroup):
         limit=JobExecutionLimit.GROUP_ONCE,
         on_condition=HomeAssistantJobError,
     )
-    async def stop(self) -> None:
+    async def stop(self, *, remove_container: bool = False) -> None:
         """Stop Home Assistant Docker."""
         try:
-            return await self.instance.stop(remove_container=False)
+            return await self.instance.stop(remove_container=remove_container)
         except DockerError as err:
             raise HomeAssistantError() from err
 

supervisor/jobs/__init__.py

@@ -8,7 +8,7 @@ from dataclasses import dataclass
 from datetime import datetime
 import logging
 from typing import Any
-from uuid import UUID, uuid4
+from uuid import uuid4
 
 from attrs import Attribute, define, field
 from attrs.setters import convert as attr_convert, frozen, validate as attr_validate
@@ -28,7 +28,7 @@ from .validate import SCHEMA_JOBS_CONFIG
 # When a new asyncio task is started the current context is copied over.
 # Modifications to it in one task are not visible to others though.
 # This allows us to track what job is currently in progress in each task.
-_CURRENT_JOB: ContextVar[UUID] = ContextVar("current_job")
+_CURRENT_JOB: ContextVar[str] = ContextVar("current_job")
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -85,7 +85,7 @@ class SupervisorJob:
     """Representation of a job running in supervisor."""
 
     created: datetime = field(init=False, factory=utcnow, on_setattr=frozen)
-    uuid: UUID = field(init=False, factory=lambda: uuid4().hex, on_setattr=frozen)
+    uuid: str = field(init=False, factory=lambda: uuid4().hex, on_setattr=frozen)
     name: str | None = field(default=None, validator=[_invalid_if_started])
     reference: str | None = field(default=None, on_setattr=_on_change)
     progress: float = field(
@@ -97,7 +97,7 @@ class SupervisorJob:
     stage: str | None = field(
         default=None, validator=[_invalid_if_done], on_setattr=_on_change
     )
-    parent_id: UUID | None = field(
+    parent_id: str | None = field(
         factory=lambda: _CURRENT_JOB.get(None), on_setattr=frozen
     )
     done: bool | None = field(init=False, default=None, on_setattr=_on_change)
@@ -146,7 +146,7 @@ class SupervisorJob:
             raise JobStartException("Job has a different parent from current job")
 
         self.done = False
-        token: Token[UUID] | None = None
+        token: Token[str] | None = None
         try:
             token = _CURRENT_JOB.set(self.uuid)
             yield self
@@ -237,7 +237,7 @@ class JobManager(FileConfiguration, CoreSysAttributes):
         self._jobs[job.uuid] = job
         return job
 
-    def get_job(self, uuid: UUID) -> SupervisorJob:
+    def get_job(self, uuid: str) -> SupervisorJob:
         """Return a job by uuid. Raises if it does not exist."""
         if uuid not in self._jobs:
             raise JobNotFound(f"No job found with id {uuid}")

tests/api/test_backups.py

@@ -15,9 +15,11 @@ from supervisor.addons.addon import Addon
 from supervisor.backups.backup import Backup
 from supervisor.const import CoreState
 from supervisor.coresys import CoreSys
+from supervisor.docker.manager import DockerAPI
 from supervisor.exceptions import AddonsError, HomeAssistantBackupError
 from supervisor.homeassistant.core import HomeAssistantCore
 from supervisor.homeassistant.module import HomeAssistant
+from supervisor.homeassistant.websocket import HomeAssistantWebSocket
 from supervisor.mounts.mount import Mount
 from supervisor.supervisor import Supervisor
 
@@ -857,3 +859,58 @@ async def test_restore_backup_from_location(
     )
     assert resp.status == 200
     assert test_file.is_file()
+
+
+@pytest.mark.parametrize(
+    ("backup_type", "postbody"), [("partial", {"homeassistant": True}), ("full", {})]
+)
+@pytest.mark.usefixtures("tmp_supervisor_data", "path_extern")
+async def test_restore_homeassistant_adds_env(
+    api_client: TestClient,
+    coresys: CoreSys,
+    docker: DockerAPI,
+    backup_type: str,
+    postbody: dict[str, Any],
+):
+    """Test restoring home assistant from backup adds env to container."""
+    event = asyncio.Event()
+    coresys.core.state = CoreState.RUNNING
+    coresys.hardware.disk.get_disk_free_space = lambda x: 5000
+    coresys.homeassistant.version = AwesomeVersion("2025.1.0")
+    backup = await coresys.backups.do_backup_full()
+
+    async def mock_async_send_message(_, message: dict[str, Any]):
+        """Mock of async send message in ws client."""
+        if (
+            message["data"]["event"] == "job"
+            and message["data"]["data"]["name"]
+            == f"backup_manager_{backup_type}_restore"
+            and message["data"]["data"]["reference"] == backup.slug
+            and message["data"]["data"]["done"]
+        ):
+            event.set()
+
+    with (
+        patch.object(HomeAssistantCore, "_block_till_run"),
+        patch.object(
+            HomeAssistantWebSocket, "async_send_message", new=mock_async_send_message
+        ),
+    ):
+        resp = await api_client.post(
+            f"/backups/{backup.slug}/restore/{backup_type}",
+            json={"background": True} | postbody,
+        )
+        assert resp.status == 200
+        body = await resp.json()
+        job = coresys.jobs.get_job(body["data"]["job_id"])
+
+        if not job.done:
+            await asyncio.wait_for(event.wait(), 5)
+
+    assert docker.containers.create.call_args.kwargs["name"] == "homeassistant"
+    assert (
+        docker.containers.create.call_args.kwargs["environment"][
+            "SUPERVISOR_RESTORE_JOB_ID"
+        ]
+        == job.uuid
+    )

tests/backups/test_manager.py

@@ -1249,11 +1249,6 @@ async def test_restore_progress(
         _make_backup_message_for_assert(
             action="full_restore", reference=full_backup.slug, stage="addons"
         ),
-        _make_backup_message_for_assert(
-            action="full_restore",
-            reference=full_backup.slug,
-            stage="await_home_assistant_restart",
-        ),
         _make_backup_message_for_assert(
             action="full_restore",
             reference=full_backup.slug,
@@ -1262,12 +1257,12 @@ async def test_restore_progress(
         _make_backup_message_for_assert(
             action="full_restore",
             reference=full_backup.slug,
-            stage="check_home_assistant",
+            stage="await_home_assistant_restart",
         ),
         _make_backup_message_for_assert(
             action="full_restore",
             reference=full_backup.slug,
-            stage="check_home_assistant",
+            stage="await_home_assistant_restart",
             done=True,
         ),
     ]