Move setup tasks

custom_components/hacs/base.py

@@ -1,4 +1,6 @@
 """Base HACS class."""
+from __future__ import annotations
+
 import logging
 from typing import List, Optional, TYPE_CHECKING
 import pathlib
@@ -17,6 +19,7 @@
 
 if TYPE_CHECKING:
     from .helpers.classes.repository import HacsRepository
+    from .task.manager import HacsTaskManager
 
 
 class HacsCommon:
@@ -42,92 +45,44 @@ class HacsStatus:
 class HacsBaseAttributes:
     """Base HACS class."""
 
-    _default: Optional[AIOGitHubAPIRepository]
-    _github: Optional[AIOGitHubAPI]
-    _hass: Optional[HomeAssistant]
-    _configuration: Optional[Configuration]
-    _repository: Optional[AIOGitHubAPIRepository]
-    _stage: HacsStage = HacsStage.SETUP
-    _common: Optional[HacsCommon]
-
-    core: HacsCore = attr.ib(HacsCore)
-    common: HacsCommon = attr.ib(HacsCommon)
-    status: HacsStatus = attr.ib(HacsStatus)
-    frontend: HacsFrontend = attr.ib(HacsFrontend)
-    log: logging.Logger = getLogger()
-    system: HacsSystem = attr.ib(HacsSystem)
+    default: Optional[AIOGitHubAPIRepository]
+    github: Optional[AIOGitHubAPI]
+
     repositories: List["HacsRepository"] = []
 
+    task: HacsTaskManager | None = None
 
-@attr.s
-class HacsBase(HacsBaseAttributes):
+
+class Hacs:
     """Base HACS class."""
 
+    def __init__(self) -> None:
+        """Initialize."""
+        self.common: HacsCommon = HacsCommon()
+        self.configuration: Configuration = Configuration()
+        self.core: HacsCore = HacsCore()
+        self.default: AIOGitHubAPIRepository | None = None
+        self.frontend: HacsFrontend = HacsFrontend()
+        self.github: AIOGitHubAPI | None = None
+        self.hass: HomeAssistant | None = None
+        self.integration_dir = pathlib.Path(__file__).parent
+        self.log: logging.Logger = getLogger()
+        self.repository: AIOGitHubAPIRepository | None = None
+        self.status: HacsStatus = HacsStatus()
+        self.system: HacsSystem = HacsSystem()
+
+        self._stage: HacsStage = HacsStage.SETUP
+
     @property
     def stage(self) -> HacsStage:
-        """Returns a HacsStage object."""
+        """Returns the current stage."""
         return self._stage
 
     @stage.setter
-    def stage(self, value: HacsStage) -> None:
-        """Set the value for the stage property."""
+    def stage(self, value: Configuration) -> None:
+        """Set the current stage."""
         self._stage = value
 
-    @property
-    def github(self) -> Optional[AIOGitHubAPI]:
-        """Returns a AIOGitHubAPI object."""
-        return self._github
-
-    @github.setter
-    def github(self, value: AIOGitHubAPI) -> None:
-        """Set the value for the github property."""
-        self._github = value
-
-    @property
-    def repository(self) -> Optional[AIOGitHubAPIRepository]:
-        """Returns a AIOGitHubAPIRepository object representing hacs/integration."""
-        return self._repository
-
-    @repository.setter
-    def repository(self, value: AIOGitHubAPIRepository) -> None:
-        """Set the value for the repository property."""
-        self._repository = value
-
-    @property
-    def default(self) -> Optional[AIOGitHubAPIRepository]:
-        """Returns a AIOGitHubAPIRepository object representing hacs/default."""
-        return self._default
-
-    @default.setter
-    def default(self, value: AIOGitHubAPIRepository) -> None:
-        """Set the value for the default property."""
-        self._default = value
-
-    @property
-    def hass(self) -> Optional[HomeAssistant]:
-        """Returns a HomeAssistant object."""
-        return self._hass
-
-    @hass.setter
-    def hass(self, value: HomeAssistant) -> None:
-        """Set the value for the default property."""
-        self._hass = value
-
-    @property
-    def configuration(self) -> Optional[Configuration]:
-        """Returns a Configuration object."""
-        return self._configuration
-
-    @configuration.setter
-    def configuration(self, value: Configuration) -> None:
-        """Set the value for the default property."""
-        self._configuration = value
-
-    @property
-    def integration_dir(self) -> pathlib.Path:
-        """Return the HACS integration dir."""
-        return pathlib.Path(__file__).parent
-
     def disable(self, reason: HacsDisabledReason) -> None:
         """Disable HACS."""
         self.system.disabled = True
@@ -139,3 +94,23 @@ def enable(self) -> None:
         self.system.disabled = False
         self.system.disabled_reason = None
         self.log.info("HACS is enabled")
+
+
+class HacsBase(Hacs):
+    """Hacs Baseclass"""
+
+    __hacs: Hacs
+
+    common: HacsCommon = __hacs.common
+    configuration: Configuration = Configuration()
+    core: HacsCore = HacsCore()
+    default: AIOGitHubAPIRepository | None = None
+    frontend: HacsFrontend = __hacs.frontend
+    github: AIOGitHubAPI | None = __hacs.github
+    hass: HomeAssistant | None = __hacs.hass
+    integration_dir = pathlib.Path(__file__).parent
+    log: logging.Logger = __hacs.log
+    repository: AIOGitHubAPIRepository = __hacs.repository
+    stage: HacsStage = __hacs.stage
+    status: HacsStatus = HacsStatus()
+    system: HacsSystem = HacsSystem()

custom_components/hacs/enums.py

@@ -3,6 +3,11 @@
 from enum import Enum
 
 
+class ConfigurationMode(str, Enum):
+    YAML = "yaml"
+    UI = "flow"
+
+
 class HacsCategory(str, Enum):
     APPDAEMON = "appdaemon"
     INTEGRATION = "integration"
@@ -35,6 +40,8 @@ class HacsSetupTask(str, Enum):
     WEBSOCKET = "WebSocket API"
     FRONTEND = "Frontend"
     SENSOR = "Sensor"
+    RESTORE = "Restore"
+    SCHEDULE = "Scheduled startup tasks"
     HACS_REPO = "Hacs Repository"
     CATEGORIES = "Additional categories"
     CLEAR_STORAGE = "Clear storage"

custom_components/hacs/operational/setup.py

@@ -1,4 +1,5 @@
 """Setup HACS."""
+from custom_components.hacs.task.manager import HacsTaskManager
 from datetime import datetime
 from aiogithubapi import AIOGitHubAPIException, GitHub
 from homeassistant.const import EVENT_HOMEASSISTANT_STARTED
@@ -7,14 +8,22 @@
 from homeassistant.exceptions import ConfigEntryNotReady, HomeAssistantError
 from homeassistant.helpers.aiohttp_client import async_create_clientsession
 from homeassistant.helpers.event import async_call_later
+from homeassistant.helpers import discovery
 
 from custom_components.hacs.const import (
     DOMAIN,
     HACS_GITHUB_API_HEADERS,
     INTEGRATION_VERSION,
     STARTUP,
 )
-from custom_components.hacs.enums import HacsDisabledReason, HacsStage, LovelaceMode
+from custom_components.hacs.enums import (
+    HacsDisabledReason,
+    HacsSetupTask,
+    HacsStage,
+    LovelaceMode,
+    ConfigurationMode,
+)
+from custom_components.hacs.base import HacsBase
 from custom_components.hacs.hacsbase.configuration import Configuration
 from custom_components.hacs.hacsbase.data import HacsData
 from custom_components.hacs.helpers.functions.constrains import check_constrains
@@ -32,16 +41,12 @@
 from custom_components.hacs.operational.setup_actions.load_hacs_repository import (
     async_load_hacs_repository,
 )
-from custom_components.hacs.operational.setup_actions.sensor import async_add_sensor
 from custom_components.hacs.operational.setup_actions.websocket_api import (
     async_setup_hacs_websockt_api,
 )
 from custom_components.hacs.share import get_hacs
 
-try:
-    from homeassistant.components.lovelace import system_health_info
-except ImportError:
-    from homeassistant.components.lovelace.system_health import system_health_info
+from homeassistant.components.lovelace.system_health import system_health_info
 
 
 async def _async_common_setup(hass):
@@ -50,6 +55,12 @@ async def _async_common_setup(hass):
     hacs.hass = hass
     hacs.system.running = True
     hacs.session = async_create_clientsession(hass)
+    hacs.tasks = HacsTaskManager()
+
+    await hacs.tasks.async_load()
+
+    test = hacs.tasks.get("hello_world")
+    await test.execute()
 
 
 async def async_setup_entry(hass, config_entry):
@@ -68,7 +79,7 @@ async def async_setup_entry(hass, config_entry):
     hacs.configuration = Configuration.from_dict(
         config_entry.data, config_entry.options
     )
-    hacs.configuration.config_type = "flow"
+    hacs.configuration.config_type = ConfigurationMode.UI
     hacs.configuration.config_entry = config_entry
 
     return await async_startup_wrapper_for_config_entry()
@@ -79,13 +90,13 @@ async def async_setup(hass, config):
     hacs = get_hacs()
     if DOMAIN not in config:
         return True
-    if hacs.configuration and hacs.configuration.config_type == "flow":
+    if hacs.configuration and hacs.configuration.config_type == ConfigurationMode.UI:
         return True
 
     await _async_common_setup(hass)
 
     hacs.configuration = Configuration.from_dict(config[DOMAIN])
-    hacs.configuration.config_type = "yaml"
+    hacs.configuration.config_type = ConfigurationMode.YAML
     await async_startup_wrapper_for_yaml()
     return True
 
@@ -129,15 +140,17 @@ async def async_hacs_startup():
         lovelace_info = await system_health_info(hacs.hass)
     except (TypeError, KeyError, HomeAssistantError):
         # If this happens, the users YAML is not valid, we assume YAML mode
-        lovelace_info = {"mode": "yaml"}
+        lovelace_info = {"mode": LovelaceMode.YAML}
     hacs.log.debug(f"Configuration type: {hacs.configuration.config_type}")
     hacs.version = INTEGRATION_VERSION
     hacs.log.info(STARTUP)
     hacs.core.config_path = hacs.hass.config.path()
     hacs.system.ha_version = HAVERSION
 
-    hacs.system.lovelace_mode = lovelace_info.get("mode", "yaml")
-    hacs.core.lovelace_mode = LovelaceMode(lovelace_info.get("mode", "yaml"))
+    hacs.system.lovelace_mode = LovelaceMode(
+        lovelace_info.get("mode", LovelaceMode.YAML)
+    )
+    hacs.core.lovelace_mode = LovelaceMode(lovelace_info.get("mode", LovelaceMode.YAML))
 
     # Setup websocket API
     await async_setup_hacs_websockt_api()
@@ -177,42 +190,80 @@ async def async_hacs_startup():
 
     # Check HACS Constrains
     if not await hacs.hass.async_add_executor_job(check_constrains):
-        if hacs.configuration.config_type == "flow":
+        if hacs.configuration.config_type == ConfigurationMode.UI:
             if hacs.configuration.config_entry is not None:
                 await async_remove_entry(hacs.hass, hacs.configuration.config_entry)
         hacs.disable(HacsDisabledReason.CONSTRAINS)
         return False
 
-    # Load HACS
+    for task, name in (
+        (async_load_hacs, HacsSetupTask.HACS_REPO),
+        (async_storrage_restore, HacsSetupTask.RESTORE),
+        (async_schedule_startup_tasks, HacsSetupTask.SCHEDULE),
+        (async_setup_sensor, HacsSetupTask.SENSOR),
+    ):
+        hacs.log.info("Setup task '%s' started", name)
+        start_time = datetime.now()
+        result = await task(hacs)
+        time_elapsed = datetime.now() - start_time
+        hacs.log.debug(
+            "Setup task '%s' completed with result '%s' in '%f' seconds",
+            name,
+            result,
+            time_elapsed.total_seconds(),
+        )
+        if result is False:
+            return False
+
+    # Mischief managed!
+    await hacs.async_set_stage(HacsStage.WAITING)
+    hacs.log.info(
+        "Setup complete, waiting for Home Assistant before startup tasks starts"
+    )
+    return True
+
+
+async def async_schedule_startup_tasks(hacs: HacsBase):
+    """Schedule startup tasks."""
+    if hacs.hass.state == CoreState.running:
+        async_call_later(hacs.hass, 5, hacs.startup_tasks)
+    else:
+        hacs.hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STARTED, hacs.startup_tasks)
+
+
+async def async_load_hacs(hacs: HacsBase):
+    """Load HACS."""
     if not await async_load_hacs_repository():
-        if hacs.configuration.config_type == "flow":
+        if hacs.configuration.config_type == ConfigurationMode.UI:
             if hacs.configuration.config_entry is not None:
                 await async_remove_entry(hacs.hass, hacs.configuration.config_entry)
         hacs.disable(HacsDisabledReason.LOAD_HACS)
         return False
 
-    # Restore from storefiles
+
+async def async_storrage_restore(hacs: HacsBase):
+    """Setup storrage restore."""
     if not await hacs.data.restore():
         hacs_repo = hacs.get_by_name("hacs/integration")
         hacs_repo.pending_restart = True
-        if hacs.configuration.config_type == "flow":
+        if hacs.configuration.config_type == ConfigurationMode.UI:
             if hacs.configuration.config_entry is not None:
                 await async_remove_entry(hacs.hass, hacs.configuration.config_entry)
         hacs.disable(HacsDisabledReason.RESTORE)
         return False
 
-    # Setup startup tasks
-    if hacs.hass.state == CoreState.running:
-        async_call_later(hacs.hass, 5, hacs.startup_tasks)
-    else:
-        hacs.hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STARTED, hacs.startup_tasks)
-
-    # Set up sensor
-    await async_add_sensor()
 
-    # Mischief managed!
-    await hacs.async_set_stage(HacsStage.WAITING)
-    hacs.log.info(
-        "Setup complete, waiting for Home Assistant before startup tasks starts"
+async def async_setup_sensor(hacs: HacsBase):
+    """Setup HACS sensor."""
+    if hacs.configuration.config_type == ConfigurationMode.YAML:
+        hacs.hass.async_create_task(
+            discovery.async_load_platform(
+                hacs.hass, "sensor", DOMAIN, {}, hacs.configuration.config
+            )
+        )
+        return
+    hacs.hass.async_add_job(
+        hacs.hass.config_entries.async_forward_entry_setup(
+            hacs.configuration.config_entry, "sensor"
+        )
     )
-    return True