Improved VM Status reporting back to HA. Added extra attributes for VMs and Containers.

Author: MrD3y5eL

custom_components/unraid/const.py

@@ -1,6 +1,7 @@
 """Constants for the Unraid integration."""
 from homeassistant.const import Platform
 
+
 DOMAIN = "unraid"
 DEFAULT_PORT = 22
 DEFAULT_CHECK_INTERVAL = 300  # seconds

custom_components/unraid/coordinator.py

@@ -32,16 +32,22 @@ def __init__(self, hass: HomeAssistant, api: UnraidAPI, entry: ConfigEntry) -> N
     async def _async_update_data(self) -> Dict[str, Any]:
         """Fetch data from Unraid."""
         try:
+            # Fetch VM data first for faster switch response
+            vms = await self.api.get_vms()
+            
+            # Then fetch the rest of the data
             data = {
+                "vms": vms,
                 "system_stats": await self.api.get_system_stats(),
                 "docker_containers": await self.api.get_docker_containers(),
-                "vms": await self.api.get_vms(),
                 "user_scripts": await self.api.get_user_scripts(),
             }
+            
             if self.has_ups:
                 ups_info = await self.api.get_ups_info()
-                if ups_info:  # Only add UPS info if it's not empty
+                if ups_info:
                     data["ups_info"] = ups_info
+                    
             return data
         except Exception as err:
             _LOGGER.error("Error communicating with Unraid: %s", err)

custom_components/unraid/manifest.json

@@ -10,6 +10,6 @@
   "issue_tracker": "https://github.com/domalab/ha-unraid/issues",
   "requirements": [],
   "ssdp": [],
-  "version": "0.1.3",
+  "version": "0.1.4",
   "zeroconf": []
 }

custom_components/unraid/switch.py

@@ -1,11 +1,14 @@
 """Switch platform for Unraid."""
 from __future__ import annotations
 
+from typing import Any, Dict
+
 from homeassistant.components.switch import SwitchEntity
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
+from homeassistant.exceptions import HomeAssistantError
 
 from .const import DOMAIN
 from .coordinator import UnraidDataUpdateCoordinator
@@ -79,9 +82,21 @@ def is_on(self) -> bool:
         """Return true if the container is running."""
         for container in self.coordinator.data["docker_containers"]:
             if container["name"] == self._container_name:
-                return container["status"].lower() == "running"
+                return container["state"] == "running"
         return False
 
+    @property
+    def extra_state_attributes(self) -> Dict[str, Any]:
+        """Return the state attributes."""
+        for container in self.coordinator.data["docker_containers"]:
+            if container["name"] == self._container_name:
+                return {
+                    "container_id": container["id"],
+                    "status": container["status"],
+                    "image": container["image"]
+                }
+        return {}
+
     async def async_turn_on(self, **kwargs) -> None:
         """Turn the container on."""
         await self.coordinator.api.start_container(self._container_name)
@@ -100,7 +115,27 @@ def __init__(self, coordinator: UnraidDataUpdateCoordinator, vm_name: str) -> No
         super().__init__(coordinator, f"vm_{vm_name}")
         self._vm_name = vm_name
         self._attr_name = f"Unraid VM {vm_name}"
-        self._attr_icon = "mdi:desktop-classic"
+        self._attr_entity_registry_enabled_default = True
+        self._attr_assumed_state = False
+
+    @property
+    def icon(self) -> str:
+        """Return the icon to use for the VM."""
+        for vm in self.coordinator.data["vms"]:
+            if vm["name"] == self._vm_name:
+                if vm.get("os_type") == "windows":
+                    return "mdi:microsoft-windows"
+                elif vm.get("os_type") == "linux":
+                    return "mdi:linux"
+                return "mdi:desktop-tower"
+        return "mdi:desktop-tower"
+
+    @property
+    def available(self) -> bool:
+        """Return if the switch is available."""
+        if not self.coordinator.last_update_success:
+            return False
+        return any(vm["name"] == self._vm_name for vm in self.coordinator.data["vms"])
 
     @property
     def is_on(self) -> bool:
@@ -109,13 +144,28 @@ def is_on(self) -> bool:
             if vm["name"] == self._vm_name:
                 return vm["status"].lower() == "running"
         return False
+    
+    @property
+    def extra_state_attributes(self) -> Dict[str, Any]:
+        """Return the state attributes."""
+        for vm in self.coordinator.data["vms"]:
+            if vm["name"] == self._vm_name:
+                return {
+                    "os_type": vm.get("os_type", "unknown"),
+                    "status": vm.get("status", "unknown"),
+                }
+        return {}
 
     async def async_turn_on(self, **kwargs) -> None:
         """Turn the VM on."""
-        await self.coordinator.api.start_vm(self._vm_name)
+        success = await self.coordinator.api.start_vm(self._vm_name)
+        if not success:
+            raise HomeAssistantError(f"Failed to start VM {self._vm_name}")
         await self.coordinator.async_request_refresh()
 
     async def async_turn_off(self, **kwargs) -> None:
         """Turn the VM off."""
-        await self.coordinator.api.stop_vm(self._vm_name)
+        success = await self.coordinator.api.stop_vm(self._vm_name)
+        if not success:
+            raise HomeAssistantError(f"Failed to stop VM {self._vm_name}")
         await self.coordinator.async_request_refresh()

custom_components/unraid/unraid.py

@@ -5,9 +5,53 @@
 from typing import Dict, List, Any, Optional
 import re
 from async_timeout import timeout
+from enum import Enum
+import json
 
 _LOGGER = logging.getLogger(__name__)
 
+class VMState(Enum):
+    """VM states matching Unraid/libvirt states."""
+    RUNNING = 'running'
+    STOPPED = 'shut off'
+    PAUSED = 'paused'
+    IDLE = 'idle'
+    IN_SHUTDOWN = 'in shutdown'
+    CRASHED = 'crashed'
+    SUSPENDED = 'pmsuspended'
+
+    @classmethod
+    def is_running(cls, state: str) -> bool:
+        """Check if the state represents a running VM."""
+        return state.lower() == cls.RUNNING.value
+
+    @classmethod
+    def parse(cls, state: str) -> str:
+        """Parse the VM state string."""
+        state = state.lower().strip()
+        try:
+            return next(s.value for s in cls if s.value == state)
+        except StopIteration:
+            return state
+
+class ContainerStates(Enum):
+    """Docker container states."""
+    RUNNING = 'running'
+    EXITED = 'exited'
+    PAUSED = 'paused'
+    RESTARTING = 'restarting'
+    DEAD = 'dead'
+    CREATED = 'created'
+
+    @classmethod
+    def parse(cls, state: str) -> str:
+        """Parse the container state string."""
+        state = state.lower().strip()
+        try:
+            return next(s.value for s in cls if s.value == state)
+        except StopIteration:
+            return state
+
 class UnraidAPI:
     """API client for interacting with Unraid servers."""
 
@@ -397,18 +441,35 @@ async def get_docker_containers(self) -> List[Dict[str, Any]]:
         """Fetch information about Docker containers."""
         try:
             _LOGGER.debug("Fetching Docker container information")
-            result = await self.execute_command("docker ps -a --format '{{.Names}}|{{.State}}'")
+            # Get basic container info with proven format
+            result = await self.execute_command("docker ps -a --format '{{.Names}}|{{.State}}|{{.ID}}|{{.Image}}'")
             if result.exit_status != 0:
                 _LOGGER.error("Docker container list command failed with exit status %d", result.exit_status)
                 return []
-            
+
             containers = []
             for line in result.stdout.splitlines():
                 parts = line.split('|')
-                if len(parts) == 2:
-                    containers.append({"name": parts[0], "status": parts[1]})
+                if len(parts) == 4:  # Now expecting 4 parts
+                    container_name = parts[0].strip()
+                    # Get container icon if available
+                    icon_path = f"/var/lib/docker/unraid/images/{container_name}-icon.png"
+                    icon_result = await self.execute_command(
+                        f"[ -f {icon_path} ] && (base64 {icon_path}) || echo ''"
+                    )
+                    icon_data = icon_result.stdout[0] if icon_result.exit_status == 0 else ""
+
+                    containers.append({
+                        "name": container_name,
+                        "state": ContainerStates.parse(parts[1].strip()),
+                        "status": parts[1].strip(),
+                        "id": parts[2].strip(),
+                        "image": parts[3].strip(),
+                        "icon": icon_data
+                    })
                 else:
                     _LOGGER.warning("Unexpected format in docker container output: %s", line)
+
             return containers
         except Exception as e:
             _LOGGER.error("Error getting docker containers: %s", str(e))
@@ -417,8 +478,8 @@ async def get_docker_containers(self) -> List[Dict[str, Any]]:
     async def start_container(self, container_name: str) -> bool:
         """Start a Docker container."""
         try:
-            _LOGGER.debug("Starting Docker container: %s", container_name)
-            result = await self.execute_command(f"docker start {container_name}")
+            _LOGGER.debug("Starting container: %s", container_name)
+            result = await self.execute_command(f'docker start "{container_name}"')
             if result.exit_status != 0:
                 _LOGGER.error("Failed to start container %s: %s", container_name, result.stderr)
                 return False
@@ -427,12 +488,12 @@ async def start_container(self, container_name: str) -> bool:
         except Exception as e:
             _LOGGER.error("Error starting container %s: %s", container_name, str(e))
             return False
-        
+
     async def stop_container(self, container_name: str) -> bool:
         """Stop a Docker container."""
         try:
-            _LOGGER.debug("Stopping Docker container: %s", container_name)
-            result = await self.execute_command(f"docker stop {container_name}")
+            _LOGGER.debug("Stopping container: %s", container_name)
+            result = await self.execute_command(f'docker stop "{container_name}"')
             if result.exit_status != 0:
                 _LOGGER.error("Failed to stop container %s: %s", container_name, result.stderr)
                 return False
@@ -455,43 +516,100 @@ async def get_vms(self) -> List[Dict[str, Any]]:
             for line in result.stdout.splitlines():
                 if line.strip():
                     name = line.strip()
-                    status = await self._get_vm_status(name)
-                    vms.append({"name": name, "status": status})
+                    status = await self.get_vm_status(name)
+                    os_type = await self.get_vm_os_info(name)
+                    vms.append({
+                        "name": name, 
+                        "status": status,
+                        "os_type": os_type
+                    })
             return vms
         except Exception as e:
             _LOGGER.error("Error getting VMs: %s", str(e))
             return []
+        
+    async def get_vm_os_info(self, vm_name: str) -> str:
+        """Get the OS type of a VM."""
+        try:
+            # First try to get OS info from VM XML
+            result = await self.execute_command(f'virsh dumpxml "{vm_name}" | grep "<os>"')
+            xml_output = result.stdout
+
+            # Check for Windows-specific indicators
+            if any(indicator in '\n'.join(xml_output).lower() for indicator in ['windows', 'win', 'microsoft']):
+                return 'windows'
+            
+            # Try to get detailed OS info if available
+            result = await self.execute_command(f'virsh domosinfo "{vm_name}" 2>/dev/null')
+            if result.exit_status == 0:
+                os_info = '\n'.join(result.stdout).lower()
+                if any(indicator in os_info for indicator in ['windows', 'win', 'microsoft']):
+                    return 'windows'
+                elif any(indicator in os_info for indicator in ['linux', 'unix', 'ubuntu', 'debian', 'centos', 'fedora', 'rhel']):
+                    return 'linux'
+            
+            # Default to checking common paths in VM name
+            vm_name_lower = vm_name.lower()
+            if any(win_term in vm_name_lower for win_term in ['windows', 'win']):
+                return 'windows'
+            elif any(linux_term in vm_name_lower for linux_term in ['linux', 'ubuntu', 'debian', 'centos', 'fedora', 'rhel']):
+                return 'linux'
+            
+            return 'unknown'
+        except Exception as e:
+            _LOGGER.debug("Error getting OS info for VM %s: %s", vm_name, str(e))
+            return 'unknown'
 
-    async def _get_vm_status(self, vm_name: str) -> str:
-        """Get the status of a specific virtual machine."""
+    async def get_vm_status(self, vm_name: str) -> str:
+        """Get detailed status of a specific virtual machine."""
         try:
             result = await self.execute_command(f"virsh domstate {vm_name}")
             if result.exit_status != 0:
-                _LOGGER.error("VM status command for %s failed with exit status %d", vm_name, result.exit_status)
-                return "unknown"
-            return result.stdout.strip()
+                _LOGGER.error("Failed to get VM status for %s: %s", vm_name, result.stderr)
+                return VMState.CRASHED.value
+            return VMState.parse(result.stdout.strip())
         except Exception as e:
             _LOGGER.error("Error getting VM status for %s: %s", vm_name, str(e))
-            return "unknown"
+            return VMState.CRASHED.value
 
-    async def start_vm(self, vm_name: str) -> bool:
-        """Start a virtual machine."""
+    async def stop_vm(self, vm_name: str) -> bool:
+        """Stop a virtual machine using ACPI shutdown."""
         try:
-            _LOGGER.debug("Starting VM: %s", vm_name)
-            result = await self.execute_command(f"virsh start {vm_name}")
-            return result.exit_status == 0 and "started" in result.stdout.lower()
+            _LOGGER.debug("Stopping VM: %s", vm_name)
+            result = await self.execute_command(f'virsh shutdown "{vm_name}" --mode acpi')
+            success = result.exit_status == 0
+            
+            if success:
+                # Wait for the VM to actually shut down
+                for _ in range(30):  # Wait up to 60 seconds
+                    await asyncio.sleep(2)
+                    status = await self.get_vm_status(vm_name)
+                    if status == VMState.STOPPED.value:
+                        return True
+                return False
+            return False
         except Exception as e:
-            _LOGGER.error("Error starting VM %s: %s", vm_name, str(e))
+            _LOGGER.error("Error stopping VM %s: %s", vm_name, str(e))
             return False
 
-    async def stop_vm(self, vm_name: str) -> bool:
-        """Stop a virtual machine."""
+    async def start_vm(self, vm_name: str) -> bool:
+        """Start a virtual machine and wait for it to be running."""
         try:
-            _LOGGER.debug("Stopping VM: %s", vm_name)
-            result = await self.execute_command(f"virsh shutdown {vm_name}")
-            return result.exit_status == 0 and "shutting down" in result.stdout.lower()
+            _LOGGER.debug("Starting VM: %s", vm_name)
+            result = await self.execute_command(f'virsh start "{vm_name}"')
+            success = result.exit_status == 0
+            
+            if success:
+                # Wait for the VM to actually start
+                for _ in range(15):  # Wait up to 30 seconds
+                    await asyncio.sleep(2)
+                    status = await self.get_vm_status(vm_name)
+                    if status == VMState.RUNNING.value:
+                        return True
+                return False
+            return False
         except Exception as e:
-            _LOGGER.error("Error stopping VM %s: %s", vm_name, str(e))
+            _LOGGER.error("Error starting VM %s: %s", vm_name, str(e))
             return False
 
     async def get_user_scripts(self) -> List[Dict[str, Any]]: