Initial Commit

custom_components/truenas/sensor.py

@@ -0,0 +1,176 @@
+"""TrueNAS sensor platform."""
+
+from __future__ import annotations
+
+from logging import getLogger
+from datetime import date, datetime
+from decimal import Decimal
+
+from homeassistant.components.sensor import SensorEntity
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.typing import StateType
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+
+from .coordinator import TrueNASCoordinator
+from .entity import TrueNASEntity, async_add_entities
+from .sensor_types import (
+    SENSOR_SERVICES,
+    SENSOR_TYPES,
+)
+
+_LOGGER = getLogger(__name__)
+
+
+# ---------------------------
+#   async_setup_entry
+# ---------------------------
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    _async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Set up entry for TrueNAS component."""
+    dispatcher = {
+        "TrueNASSensor": TrueNASSensor,
+        "TrueNASUptimeSensor": TrueNASUptimeSensor,
+        "TrueNASClousyncSensor": TrueNASClousyncSensor,
+        "TrueNASDatasetSensor": TrueNASDatasetSensor,
+    }
+    await async_add_entities(hass, config_entry, dispatcher)
+
+
+# ---------------------------
+#   TrueNASSensor
+# ---------------------------
+class TrueNASSensor(TrueNASEntity, SensorEntity):
+    """Define an TrueNAS sensor."""
+
+    def __init__(
+        self,
+        coordinator: TrueNASCoordinator,
+        entity_description,
+        uid: str | None = None,
+    ):
+        super().__init__(coordinator, entity_description, uid)
+        self._attr_suggested_unit_of_measurement = (
+            self.entity_description.suggested_unit_of_measurement
+        )
+
+    @property
+    def native_value(self) -> StateType | date | datetime | Decimal:
+        """Return the value reported by the sensor."""
+        return self._data[self.entity_description.data_attribute]
+
+    @property
+    def native_unit_of_measurement(self) -> str | None:
+        """Return the unit the value is expressed in."""
+        if self.entity_description.native_unit_of_measurement:
+            if self.entity_description.native_unit_of_measurement.startswith("data__"):
+                uom = self.entity_description.native_unit_of_measurement[6:]
+                if uom in self._data:
+                    return self._data[uom]
+
+            return self.entity_description.native_unit_of_measurement
+
+        return None
+
+
+# ---------------------------
+#   TrueNASUptimeSensor
+# ---------------------------
+class TrueNASUptimeSensor(TrueNASSensor):
+    """Define an TrueNAS Uptime sensor."""
+
+    async def restart(self) -> None:
+        """Restart TrueNAS systen."""
+        await self.hass.async_add_executor_job(
+            self.coordinator.api.query,
+            "system/reboot",
+            "post",
+        )
+
+    async def stop(self) -> None:
+        """Shutdown TrueNAS systen."""
+        await self.hass.async_add_executor_job(
+            self.coordinator.api.query,
+            "system/shutdown",
+            "post",
+        )
+
+
+# ---------------------------
+#   TrueNASDatasetSensor
+# ---------------------------
+class TrueNASDatasetSensor(TrueNASSensor):
+    """Define an TrueNAS Dataset sensor."""
+
+    async def snapshot(self) -> None:
+        """Create dataset snapshot."""
+        ts = datetime.now().isoformat(sep="_", timespec="microseconds")
+        await self.hass.async_add_executor_job(
+            self.coordinator.api.query,
+            "zfs/snapshot",
+            "post",
+            {"dataset": f"{self._data['name']}", "name": f"custom-{ts}"},
+        )
+
+
+# ---------------------------
+#   TrueNASClousyncSensor
+# ---------------------------
+class TrueNASClousyncSensor(TrueNASSensor):
+    """Define an TrueNAS Cloudsync sensor."""
+
+    async def start(self) -> None:
+        """Run cloudsync job."""
+        tmp_job = await self.hass.async_add_executor_job(
+            self.coordinator.api.query, f"cloudsync/id/{self._data['id']}"
+        )
+
+        if "job" not in tmp_job:
+            _LOGGER.error(
+                "Clousync job %s (%s) invalid",
+                self._data["description"],
+                self._data["id"],
+            )
+            return
+        if tmp_job["job"]["state"] in ["WAITING", "RUNNING"]:
+            _LOGGER.warning(
+                "Clousync job %s (%s) is already running",
+                self._data["description"],
+                self._data["id"],
+            )
+            return
+
+        await self.hass.async_add_executor_job(
+            self.coordinator.api.query, f"cloudsync/id/{self._data['id']}/sync", "post"
+        )
+
+    async def stop(self) -> None:
+        """Abort cloudsync job."""
+        tmp_job = await self.hass.async_add_executor_job(
+            self.coordinator.api.query, f"cloudsync/id/{self._data['id']}"
+        )
+
+        if "job" not in tmp_job:
+            _LOGGER.error(
+                "Clousync job %s (%s) invalid",
+                self._data["description"],
+                self._data["id"],
+            )
+            return
+        if tmp_job["job"]["state"] not in ["WAITING", "RUNNING"]:
+            _LOGGER.warning(
+                "Clousync job %s (%s) is not running",
+                self._data["description"],
+                self._data["id"],
+            )
+            return
+
+        await self.hass.async_add_executor_job(
+            self.coordinator.api.query,
+            f"cloudsync/id/{self._data['id']}/abort",
+            "post",
+            None,
+        )