Merge pull request #137 from briodan/master

Add Restart Button for containers

custom_components/monitor_docker/__init__.py

@@ -34,12 +34,15 @@
     CONF_SENSORNAME,
     CONF_SWITCHENABLED,
     CONF_SWITCHNAME,
+    CONF_BUTTONENABLED,
+    CONF_BUTTONNAME,
     CONFIG,
     CONTAINER_INFO_ALLINONE,
     DEFAULT_NAME,
     DEFAULT_RETRY,
     DEFAULT_SENSORNAME,
     DEFAULT_SWITCHNAME,
+    DEFAULT_BUTTONNAME,
     DOMAIN,
     MONITORED_CONDITIONS_LIST,
     PRECISION,
@@ -68,7 +71,11 @@
         vol.Optional(CONF_SWITCHENABLED, default=True): vol.Any(
             cv.boolean, cv.ensure_list(cv.string)
         ),
+        vol.Optional(CONF_BUTTONENABLED, default=False): vol.Any(
+            cv.boolean, cv.ensure_list(cv.string)
+        ),
         vol.Optional(CONF_SWITCHNAME, default=DEFAULT_SWITCHNAME): cv.string,
+        vol.Optional(CONF_BUTTONNAME, default=DEFAULT_BUTTONNAME): cv.string,
         vol.Optional(CONF_CERTPATH, default=""): cv.string,
         vol.Optional(CONF_RETRY, default=DEFAULT_RETRY): cv.positive_int,
         vol.Optional(CONF_MEMORYCHANGE, default=100): cv.positive_int,

custom_components/monitor_docker/button.py

@@ -0,0 +1,259 @@
+"""Monitor Docker button component."""
+
+import asyncio
+import logging
+import re
+from typing import Any
+
+import voluptuous as vol
+from custom_components.monitor_docker.helpers import DockerAPI, DockerContainerAPI
+from homeassistant.components.button import ENTITY_ID_FORMAT, ButtonEntity
+from homeassistant.const import CONF_NAME
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import config_validation as cv
+from homeassistant.helpers import entity_platform
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
+from homeassistant.util import slugify
+
+from .const import (
+    API,
+    ATTR_NAME,
+    ATTR_SERVER,
+    CONF_CONTAINERS,
+    CONF_CONTAINERS_EXCLUDE,
+    CONF_PREFIX,
+    CONF_RENAME,
+    CONF_RENAME_ENITITY,
+    CONF_BUTTONENABLED,
+    CONF_BUTTONNAME,
+    CONFIG,
+    CONTAINER,
+    CONTAINER_INFO_STATE,
+    DOMAIN,
+    SERVICE_RESTART,
+)
+
+SERVICE_RESTART_SCHEMA = vol.Schema({ATTR_NAME: cv.string, ATTR_SERVER: cv.string})
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def async_setup_platform(
+    hass: HomeAssistant,
+    config: ConfigType,
+    async_add_entities: AddEntitiesCallback,
+    discovery_info: DiscoveryInfoType | None = None,
+):
+    """Set up the Monitor Docker Button."""
+
+    async def async_restart(parm) -> None:
+        cname = parm.data[ATTR_NAME]
+        cserver = parm.data.get(ATTR_SERVER, None)
+
+        server_name = name
+        if cserver is not None:
+            if cserver not in hass.data[DOMAIN]:
+                _LOGGER.error("Server '%s' is not configured", cserver)
+                return
+            else:
+                server_name = cserver
+
+        server_config = hass.data[DOMAIN][server_name][CONFIG]
+        server_api = hass.data[DOMAIN][server_name][API]
+
+        if len(server_config[CONF_CONTAINERS]) == 0:
+            if server_api.get_container(cname):
+                await server_api.get_container(cname).restart()
+            else:
+                _LOGGER.error(
+                    "Service restart failed, container '%s'does not exist", cname
+                )
+        elif cname in server_config[CONF_CONTAINERS]:
+            _LOGGER.debug("Trying to restart container '%s'", cname)
+            if server_api.get_container(cname):
+                await server_api.get_container(cname).restart()
+            else:
+                _LOGGER.error(
+                    "Service restart failed, container '%s'does not exist", cname
+                )
+        else:
+            _LOGGER.error(
+                "Service restart failed, container '%s' is not configured", cname
+            )
+
+    def find_rename(d: dict[str, str], item: str) -> str:
+        for k in d:
+            if re.match(k, item):
+                return d[k]
+
+        return item
+
+    if discovery_info is None:
+        return
+
+    instance = discovery_info[CONF_NAME]
+    name = discovery_info[CONF_NAME]
+    api = hass.data[DOMAIN][name][API]
+    config = hass.data[DOMAIN][name][CONFIG]
+
+    # Set or overrule prefix
+    prefix = name
+    if config[CONF_PREFIX]:
+        prefix = config[CONF_PREFIX]
+
+    # Don't create any butoon if disabled
+    if config[CONF_BUTTONENABLED] == False:
+        _LOGGER.debug("[%s]: Button(s) are disabled", instance)
+        return True
+
+    _LOGGER.debug("[%s]: Setting up button(s)", instance)
+
+    buttons = []
+
+    # We support add/re-add of a container
+    if CONTAINER in discovery_info:
+        clist = [discovery_info[CONTAINER]]
+    else:
+        clist = api.list_containers()
+
+    for cname in clist:
+        includeContainer = False
+        if cname in config[CONF_CONTAINERS] or not config[CONF_CONTAINERS]:
+            includeContainer = True
+
+        if config[CONF_CONTAINERS_EXCLUDE] and cname in config[CONF_CONTAINERS_EXCLUDE]:
+            includeContainer = False
+
+        if includeContainer:
+            if (
+                config[CONF_BUTTONENABLED] == True
+                or cname in config[CONF_BUTTONENABLED]
+            ):
+                _LOGGER.debug("[%s] %s: Adding component Button", instance, cname)
+
+                # Only force rename of entityid is requested, to not break backwards compatibility
+                alias_entityid = cname
+                if config[CONF_RENAME_ENITITY]:
+                    alias_entityid = find_rename(config[CONF_RENAME], cname)
+
+                buttons.append(
+                    DockerContainerButton(
+                        api.get_container(cname),
+                        instance=instance,
+                        prefix=prefix,
+                        cname=cname,
+                        alias_entityid=alias_entityid,
+                        alias_name=find_rename(config[CONF_RENAME], cname),
+                        name_format=config[CONF_BUTTONNAME],
+                    )
+                )
+            else:
+                _LOGGER.debug("[%s] %s: NOT Adding component Button", instance, cname)
+
+    if not buttons:
+        _LOGGER.info("[%s]: No containers set-up", instance)
+        return False
+
+    async_add_entities(buttons, True)
+
+    # platform = entity_platform.current_platform.get()
+    # platform.async_register_entity_service(SERVICE_RESTART, {}, "async_restart")
+    hass.services.async_register(
+        DOMAIN, SERVICE_RESTART, async_restart, schema=SERVICE_RESTART_SCHEMA
+    )
+
+    return True
+
+#################################################################
+class DockerContainerButton(ButtonEntity):
+    def __init__(
+        self, 
+        container: DockerContainerAPI,
+        instance: str,
+        prefix: str,
+        cname: str,
+        alias_entityid: str,
+        alias_name: str,
+        name_format: str,
+    ):
+        self._container = container
+        self._instance = instance
+        self._prefix = prefix
+        self._cname = cname
+        self._state = False
+        self._entity_id = ENTITY_ID_FORMAT.format(
+            slugify(self._prefix + "_" + self._cname + "_restart")
+        )
+        self._name = name_format.format(name=alias_name)
+        self._removed = False
+
+    @property
+    def entity_id(self) -> str:
+        """Return the entity id of the button."""
+        return self._entity_id
+
+    @property
+    def name(self) -> str:
+        """Return the name of the sensor."""
+        return self._name
+
+    @property
+    def should_poll(self) -> bool:
+        return False
+
+    @property
+    def icon(self) -> str:
+        return "mdi:docker"
+
+    @property
+    def extra_state_attributes(self) -> dict:
+        return {}
+
+    @property
+    def is_on(self) -> bool:
+        return self._state
+
+    async def async_press(self, **kwargs: Any) -> None:
+        await self._container.restart()
+        self._state = False
+        self.async_schedule_update_ha_state()
+
+    async def async_added_to_hass(self) -> None:
+        """Register callbacks."""
+        self._container.register_callback(self.event_callback, "button")
+
+        # Call event callback for possible information available
+        self.event_callback()
+
+    def event_callback(self, name="", remove=False) -> None:
+        """Callback for update of container information."""
+
+        if remove:
+            # If already called before, do not remove it again
+            if self._removed:
+                return
+
+            _LOGGER.info("[%s] %s: Removing button entity", self._instance, self._cname)
+            self._loop.create_task(self.async_remove())
+            self._removed = True
+            return
+
+        state = None
+
+        try:
+            info = self._container.get_info()
+        except Exception as err:
+            _LOGGER.error(
+                "[%s] %s: Cannot request container info (%s)",
+                self._instance,
+                name,
+                str(err),
+            )
+        else:
+            if info is not None:
+                state = info.get(CONTAINER_INFO_STATE) == "running"
+
+        if state is not self._state:
+            self._state = state
+            self.async_schedule_update_ha_state()

custom_components/monitor_docker/const.py

@@ -28,13 +28,16 @@
 CONF_SENSORNAME = "sensorname"
 CONF_SWITCHENABLED = "switchenabled"
 CONF_SWITCHNAME = "switchname"
+CONF_BUTTONENABLED = "buttonenabled"
+CONF_BUTTONNAME = "buttonname"
 
 DEFAULT_NAME = "Docker"
 DEFAULT_RETRY = 60
 DEFAULT_SENSORNAME = "{name} {sensor}"
 DEFAULT_SWITCHNAME = "{name}"
+DEFAULT_BUTTONNAME = "{name} Restart"
 
-COMPONENTS = ["sensor", "switch"]
+COMPONENTS = ["sensor", "switch", "button"]
 
 SERVICE_RESTART = "restart"
 

custom_components/monitor_docker/helpers.py

@@ -276,7 +276,7 @@ async def _run_docker_events(self) -> None:
                     # Remove the docker info sensors
                     self.remove_entities()
 
-                    # Remove all the sensors/switches, they will be auto created if connection is working again
+                    # Remove all the sensors/switches/buttons, they will be auto created if connection is working again
                     for cname in list(self._containers.keys()):
                         try:
                             await self._container_remove(cname)
@@ -353,7 +353,7 @@ async def _run_docker_events(self) -> None:
                             await self._container_create_destroy()
                     elif event["Action"] == "rename":
                         # during a docker-compose up -d <container> the old container can be renamed
-                        # sensors/switch should be removed before the new container is monitored
+                        # sensors/switch/button should be removed before the new container is monitored
 
                         # New name
                         cname = event["Actor"]["Attributes"]["name"]
@@ -478,7 +478,7 @@ async def _container_add(self, cname: str) -> None:
         result = await self._containers[cname]._initGetContainer()
 
         if result:
-            # Lets wait 1 second before we try to create sensors/switches
+            # Lets wait 1 second before we try to create sensors/switches/buttons
             await asyncio.sleep(1)
 
             for component in COMPONENTS:
@@ -1284,6 +1284,14 @@ async def _restart(self) -> None:
         finally:
             self._busy = False
 
+    #############################################################
+    async def _restart_button(self) -> None:
+        """Called from HA button."""
+        _LOGGER.info("[%s] %s: Restart container", self._instance, self._name)
+
+        self._busy = True
+        await self._restart()
+
     #############################################################
     async def restart(self) -> None:
         """Called from service call."""
@@ -1314,7 +1322,7 @@ def get_stats(self) -> dict:
 
     #############################################################
     def register_callback(self, callback: Callable, variable: str):
-        """Register callback from sensor/switch."""
+        """Register callback from sensor/switch/button."""
         if callback not in self._subscribers:
             _LOGGER.debug(
                 "[%s] %s: Added callback to container, entity: %s",