Rework Proxmoxve to use a DataUpdateCoordinator (#45068)

Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>

homeassistant/components/proxmoxve/__init__.py

@@ -1,9 +1,10 @@
 """Support for Proxmox VE."""
-from enum import Enum
+from datetime import timedelta
 import logging
 
 from proxmoxer import ProxmoxAPI
 from proxmoxer.backends.https import AuthenticationError
+from proxmoxer.core import ResourceException
 from requests.exceptions import SSLError
 import voluptuous as vol
 
@@ -14,11 +15,14 @@ from homeassistant.const import (
     CONF_USERNAME,
     CONF_VERIFY_SSL,
 )
+from homeassistant.core import HomeAssistant
 import homeassistant.helpers.config_validation as cv
+from homeassistant.helpers.update_coordinator import (
+    CoordinatorEntity,
+    DataUpdateCoordinator,
+)
 
-_LOGGER = logging.getLogger(__name__)
+PLATFORMS = ["binary_sensor"]
-
-
 DOMAIN = "proxmoxve"
 PROXMOX_CLIENTS = "proxmox_clients"
 CONF_REALM = "realm"
@@ -27,9 +31,17 @@ CONF_NODES = "nodes"
 CONF_VMS = "vms"
 CONF_CONTAINERS = "containers"
 
+COORDINATOR = "coordinator"
+API_DATA = "api_data"
+
 DEFAULT_PORT = 8006
 DEFAULT_REALM = "pam"
 DEFAULT_VERIFY_SSL = True
+TYPE_VM = 0
+TYPE_CONTAINER = 1
+UPDATE_INTERVAL = 60
+
+_LOGGER = logging.getLogger(__name__)
 
 CONFIG_SCHEMA = vol.Schema(
     {
@@ -71,52 +83,191 @@ CONFIG_SCHEMA = vol.Schema(
 )
 
 
-def setup(hass, config):
+async def async_setup(hass: HomeAssistant, config: dict):
-    """Set up the component."""
+    """Set up the platform."""
+    hass.data.setdefault(DOMAIN, {})
 
-    # Create API Clients for later use
+    def build_client() -> ProxmoxAPI:
-    hass.data[PROXMOX_CLIENTS] = {}
+        """Build the Proxmox client connection."""
-    for entry in config[DOMAIN]:
+        hass.data[PROXMOX_CLIENTS] = {}
-        host = entry[CONF_HOST]
+        for entry in config[DOMAIN]:
-        port = entry[CONF_PORT]
+            host = entry[CONF_HOST]
-        user = entry[CONF_USERNAME]
+            port = entry[CONF_PORT]
-        realm = entry[CONF_REALM]
+            user = entry[CONF_USERNAME]
-        password = entry[CONF_PASSWORD]
+            realm = entry[CONF_REALM]
-        verify_ssl = entry[CONF_VERIFY_SSL]
+            password = entry[CONF_PASSWORD]
+            verify_ssl = entry[CONF_VERIFY_SSL]
 
-        try:
+            try:
-            # Construct an API client with the given data for the given host
+                # Construct an API client with the given data for the given host
-            proxmox_client = ProxmoxClient(
+                proxmox_client = ProxmoxClient(
-                host, port, user, realm, password, verify_ssl
+                    host, port, user, realm, password, verify_ssl
+                )
+                proxmox_client.build_client()
+            except AuthenticationError:
+                _LOGGER.warning(
+                    "Invalid credentials for proxmox instance %s:%d", host, port
+                )
+                continue
+            except SSLError:
+                _LOGGER.error(
+                    'Unable to verify proxmox server SSL. Try using "verify_ssl: false"'
+                )
+                continue
+
+            return proxmox_client
+
+    proxmox_client = await hass.async_add_executor_job(build_client)
+
+    async def async_update_data() -> dict:
+        """Fetch data from API endpoint."""
+
+        proxmox = proxmox_client.get_api_client()
+
+        def poll_api() -> dict:
+            data = {}
+
+            for host_config in config[DOMAIN]:
+                host_name = host_config["host"]
+
+                data[host_name] = {}
+
+                for node_config in host_config["nodes"]:
+                    node_name = node_config["node"]
+                    data[host_name][node_name] = {}
+
+                    for vm_id in node_config["vms"]:
+                        data[host_name][node_name][vm_id] = {}
+
+                        vm_status = call_api_container_vm(
+                            proxmox, node_name, vm_id, TYPE_VM
+                        )
+
+                        if vm_status is None:
+                            _LOGGER.warning("Vm/Container %s unable to be found", vm_id)
+                            data[host_name][node_name][vm_id] = None
+                            continue
+
+                        data[host_name][node_name][vm_id] = parse_api_container_vm(
+                            vm_status
+                        )
+
+                    for container_id in node_config["containers"]:
+                        data[host_name][node_name][container_id] = {}
+
+                        container_status = call_api_container_vm(
+                            proxmox, node_name, container_id, TYPE_CONTAINER
+                        )
+
+                        if container_status is None:
+                            _LOGGER.error(
+                                "Vm/Container %s unable to be found", container_id
+                            )
+                            data[host_name][node_name][container_id] = None
+                            continue
+
+                        data[host_name][node_name][
+                            container_id
+                        ] = parse_api_container_vm(container_status)
+
+            return data
+
+        return await hass.async_add_executor_job(poll_api)
+
+    coordinator = DataUpdateCoordinator(
+        hass,
+        _LOGGER,
+        name="proxmox_coordinator",
+        update_method=async_update_data,
+        update_interval=timedelta(seconds=UPDATE_INTERVAL),
+    )
+
+    hass.data[DOMAIN][COORDINATOR] = coordinator
+
+    # Fetch initial data
+    await coordinator.async_refresh()
+
+    for component in PLATFORMS:
+        await hass.async_create_task(
+            hass.helpers.discovery.async_load_platform(
+                component, DOMAIN, {"config": config}, config
             )
-            proxmox_client.build_client()
-        except AuthenticationError:
-            _LOGGER.warning(
-                "Invalid credentials for proxmox instance %s:%d", host, port
-            )
-            continue
-        except SSLError:
-            _LOGGER.error(
-                'Unable to verify proxmox server SSL. Try using "verify_ssl: false"'
-            )
-            continue
-
-        hass.data[PROXMOX_CLIENTS][f"{host}:{port}"] = proxmox_client
-
-    if hass.data[PROXMOX_CLIENTS]:
-        hass.helpers.discovery.load_platform(
-            "binary_sensor", DOMAIN, {"entries": config[DOMAIN]}, config
         )
-        return True
 
-    return False
+    return True
 
 
-class ProxmoxItemType(Enum):
+def parse_api_container_vm(status):
-    """Represents the different types of machines in Proxmox."""
+    """Get the container or vm api data and return it formatted in a dictionary.
 
-    qemu = 0
+    It is implemented in this way to allow for more data to be added for sensors
-    lxc = 1
+    in the future.
+    """
+
+    return {"status": status["status"], "name": status["name"]}
+
+
+def call_api_container_vm(proxmox, node_name, vm_id, machine_type):
+    """Make proper api calls."""
+    status = None
+
+    try:
+        if machine_type == TYPE_VM:
+            status = proxmox.nodes(node_name).qemu(vm_id).status.current.get()
+        elif machine_type == TYPE_CONTAINER:
+            status = proxmox.nodes(node_name).lxc(vm_id).status.current.get()
+    except ResourceException:
+        return None
+
+    return status
+
+
+class ProxmoxEntity(CoordinatorEntity):
+    """Represents any entity created for the Proxmox VE platform."""
+
+    def __init__(
+        self,
+        coordinator: DataUpdateCoordinator,
+        unique_id,
+        name,
+        icon,
+        host_name,
+        node_name,
+        vm_id=None,
+    ):
+        """Initialize the Proxmox entity."""
+        super().__init__(coordinator)
+
+        self.coordinator = coordinator
+        self._unique_id = unique_id
+        self._name = name
+        self._host_name = host_name
+        self._icon = icon
+        self._available = True
+        self._node_name = node_name
+        self._vm_id = vm_id
+
+        self._state = None
+
+    @property
+    def unique_id(self) -> str:
+        """Return the unique ID for this sensor."""
+        return self._unique_id
+
+    @property
+    def name(self) -> str:
+        """Return the name of the entity."""
+        return self._name
+
+    @property
+    def icon(self) -> str:
+        """Return the mdi icon of the entity."""
+        return self._icon
+
+    @property
+    def available(self) -> bool:
+        """Return True if entity is available."""
+        return self.coordinator.last_update_success and self._available
 
 
 class ProxmoxClient:

homeassistant/components/proxmoxve/binary_sensor.py

@@ -1,112 +1,95 @@
 """Binary sensor to read Proxmox VE data."""
 import logging
 
-from homeassistant.components.binary_sensor import BinarySensorEntity
+from homeassistant.const import STATE_OFF, STATE_ON
-from homeassistant.const import ATTR_ATTRIBUTION, CONF_HOST, CONF_PORT
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from . import CONF_CONTAINERS, CONF_NODES, CONF_VMS, PROXMOX_CLIENTS, ProxmoxItemType
+from . import COORDINATOR, DOMAIN, ProxmoxEntity
 
-ATTRIBUTION = "Data provided by Proxmox VE"
 _LOGGER = logging.getLogger(__name__)
 
 
-def setup_platform(hass, config, add_entities, discovery_info=None):
+async def async_setup_platform(hass, config, add_entities, discovery_info=None):
-    """Set up the sensor platform."""
+    """Set up binary sensors."""
+    if discovery_info is None:
+        return
+
+    coordinator = hass.data[DOMAIN][COORDINATOR]
 
     sensors = []
 
-    for entry in discovery_info["entries"]:
+    for host_config in discovery_info["config"][DOMAIN]:
-        port = entry[CONF_PORT]
+        host_name = host_config["host"]
 
-        for node in entry[CONF_NODES]:
+        for node_config in host_config["nodes"]:
-            for virtual_machine in node[CONF_VMS]:
+            node_name = node_config["node"]
-                sensors.append(
+
-                    ProxmoxBinarySensor(
+            for vm_id in node_config["vms"]:
-                        hass.data[PROXMOX_CLIENTS][f"{entry[CONF_HOST]}:{port}"],
+                coordinator_data = coordinator.data[host_name][node_name][vm_id]
-                        node["node"],
+
-                        ProxmoxItemType.qemu,
+                # unfound vm case
-                        virtual_machine,
+                if coordinator_data is None:
-                    )
+                    continue
+
+                vm_name = coordinator_data["name"]
+                vm_status = create_binary_sensor(
+                    coordinator, host_name, node_name, vm_id, vm_name
                 )
+                sensors.append(vm_status)
 
-            for container in node[CONF_CONTAINERS]:
+            for container_id in node_config["containers"]:
-                sensors.append(
+                coordinator_data = coordinator.data[host_name][node_name][container_id]
-                    ProxmoxBinarySensor(
+
-                        hass.data[PROXMOX_CLIENTS][f"{entry[CONF_HOST]}:{port}"],
+                # unfound container case
-                        node["node"],
+                if coordinator_data is None:
-                        ProxmoxItemType.lxc,
+                    continue
-                        container,
+
-                    )
+                container_name = coordinator_data["name"]
+                container_status = create_binary_sensor(
+                    coordinator, host_name, node_name, container_id, container_name
                 )
+                sensors.append(container_status)
 
-    add_entities(sensors, True)
+    add_entities(sensors)
 
 
-class ProxmoxBinarySensor(BinarySensorEntity):
+def create_binary_sensor(coordinator, host_name, node_name, vm_id, name):
+    """Create a binary sensor based on the given data."""
+    return ProxmoxBinarySensor(
+        coordinator=coordinator,
+        unique_id=f"proxmox_{node_name}_{vm_id}_running",
+        name=f"{node_name}_{name}_running",
+        icon="",
+        host_name=host_name,
+        node_name=node_name,
+        vm_id=vm_id,
+    )
+
+
+class ProxmoxBinarySensor(ProxmoxEntity):
     """A binary sensor for reading Proxmox VE data."""
 
-    def __init__(self, proxmox_client, item_node, item_type, item_id):
+    def __init__(
-        """Initialize the binary sensor."""
+        self,
-        self._proxmox_client = proxmox_client
+        coordinator: DataUpdateCoordinator,
-        self._item_node = item_node
+        unique_id,
-        self._item_type = item_type
+        name,
-        self._item_id = item_id
+        icon,
-
+        host_name,
-        self._vmname = None
+        node_name,
-        self._name = None
+        vm_id,
+    ):
+        """Create the binary sensor for vms or containers."""
+        super().__init__(
+            coordinator, unique_id, name, icon, host_name, node_name, vm_id
+        )
 
         self._state = None
 
     @property
-    def name(self):
+    def state(self):
-        """Return the name of the entity."""
+        """Return the state of the binary sensor."""
-        return self._name
+        data = self.coordinator.data[self._host_name][self._node_name][self._vm_id]
-
+        if data["status"] == "running":
-    @property
+            return STATE_ON
-    def is_on(self):
+        return STATE_OFF
-        """Return true if VM/container is running."""
-        return self._state
-
-    @property
-    def device_state_attributes(self):
-        """Return device attributes of the entity."""
-        return {
-            "node": self._item_node,
-            "vmid": self._item_id,
-            "vmname": self._vmname,
-            "type": self._item_type.name,
-            ATTR_ATTRIBUTION: ATTRIBUTION,
-        }
-
-    def update(self):
-        """Check if the VM/Container is running."""
-        item = self.poll_item()
-
-        if item is None:
-            _LOGGER.warning("Failed to poll VM/container %s", self._item_id)
-            return
-
-        self._state = item["status"] == "running"
-
-    def poll_item(self):
-        """Find the VM/Container with the set item_id."""
-        items = (
-            self._proxmox_client.get_api_client()
-            .nodes(self._item_node)
-            .get(self._item_type.name)
-        )
-        item = next(
-            (item for item in items if item["vmid"] == str(self._item_id)), None
-        )
-
-        if item is None:
-            _LOGGER.warning("Couldn't find VM/Container with the ID %s", self._item_id)
-            return None
-
-        if self._vmname is None:
-            self._vmname = item["name"]
-
-        if self._name is None:
-            self._name = f"{self._item_node} {self._vmname} running"
-
-        return item