Add button and sensor entities

custom_components/unifi_voucher/__init__.py

@@ -3,6 +3,10 @@
 
 from homeassistant.core import HomeAssistant
 from homeassistant.config_entries import ConfigEntry
+from homeassistant.exceptions import (
+ ConfigEntryAuthFailed,
+ ConfigEntryNotReady,
+)
 
 from .const import (
  DOMAIN,
@@ -12,36 +16,46 @@
 
 
 async def async_setup(hass: HomeAssistant, config: dict) -> bool:
- """Set up DUniFi Hotspot Manager component."""
+ """Set up UniFi Hotspot Manager component."""
  hass.data.setdefault(DOMAIN, {})
 
  return True
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
  """Set up platform from a ConfigEntry."""
  hass.data.setdefault(DOMAIN, {})
- hass.data[DOMAIN][entry.entry_id] = coordinator = UnifiVoucherCoordinator(
- hass=hass,
- config_entry=entry,
- )
- await coordinator.async_config_entry_first_refresh()
 
- await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
- entry.async_on_unload(entry.add_update_listener(async_reload_entry))
+ try:
+ coordinator = UnifiVoucherCoordinator(
+ hass=hass,
+ config_entry=config_entry,
+ )
+ await coordinator.initialize()
+ await coordinator.async_config_entry_first_refresh()
+
+ except AuthenticationRequired as err:
+ raise ConfigEntryAuthFailed from err
+
+ except Exception as err:
+ raise ConfigEntryNotReady from err
+
+ hass.data[DOMAIN][config_entry.entry_id] = coordinator
+ await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS)
+ config_entry.async_on_unload(config_entry.add_update_listener(async_reload_entry))
 
  return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
  """Unload a config entry."""
- if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
+ if unload_ok := await hass.config_entries.async_unload_platforms(config_entry, PLATFORMS):
  # Remove config entry from domain.
- hass.data[DOMAIN].pop(entry.entry_id)
+ hass.data[DOMAIN].pop(config_entry.entry_id)
  return unload_ok
 
 
-async def async_reload_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
+async def async_reload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> None:
  """Reload config entry."""
- await async_unload_entry(hass, entry)
- await async_setup_entry(hass, entry)
+ await async_unload_entry(hass, config_entry)
+ await async_setup_entry(hass, config_entry)

custom_components/unifi_voucher/api.py

@@ -3,11 +3,19 @@
 
 import asyncio
 from dataclasses import dataclass
+from typing import TypedDict
+
+from datetime import (
+ timedelta,
+ datetime,
+)
 
 from aiohttp import CookieJar
 import aiounifi
+from aiounifi.interfaces.api_handlers import APIHandler
 from aiounifi.models.configuration import Configuration
 from aiounifi.models.api import (
+ ApiItem,
  ApiRequest,
  TypedApiResponse,
 )
@@ -17,6 +25,7 @@
  HomeAssistant,
 )
 from homeassistant.helpers import aiohttp_client
+import homeassistant.util.dt as dt_util
 
 from .const import (
  LOGGER,
@@ -40,6 +49,28 @@ class UnifiVoucherApiAuthenticationError(UnifiVoucherApiError):
  """Exception to indicate an authentication error."""
 
 
+class UnifiTypedVoucher(TypedDict):
+ """Voucher description."""
+
+ _id: str
+ site_id: str
+ note: str
+ code: str
+ quota: int
+ duration: int
+ qos_overwrite: bool
+ qos_usage_quota: str
+ qos_rate_max_up: int
+ qos_rate_max_down: int
+ used: int
+ create_time: datetime
+ start_time: int
+ end_time: int
+ for_hotspot: bool
+ admin_name: str
+ status: str
+ status_expires: int
+
 @dataclass
 class UnifiVoucherListRequest(ApiRequest):
  """Request object for device list."""
@@ -54,6 +85,7 @@ def create(
  path="/stat/voucher",
  )
 
+
 @dataclass
 class UnifiVoucherCreateRequest(ApiRequest):
  """Request object for voucher create."""
@@ -80,31 +112,153 @@ def create(
  :param byte_quota: quantity of bytes allowed in MB
  :param note: description
  """
- params = {
+ data = {
+ "cmd": "create-voucher",
  "n": number,
  "quota": quota,
  "expire": "custom",
  "expire_number": expire,
  "expire_unit": 1,
+ "down": None,
+ "up": None,
  }
  if up_bandwidth:
- params["up"] = up_bandwidth
+ data["up"] = up_bandwidth
  if down_bandwidth:
- params["down"] = down_bandwidth
+ data["down"] = down_bandwidth
  if byte_quota:
- params["bytes"] = byte_quota
+ data["bytes"] = byte_quota
  if note:
- params["note"] = note
+ data["note"] = note
 
  return cls(
  method="post",
  path="/cmd/hotspot",
- data={
- "cmd": "create-voucher",
- "params": params,
- },
+ data=data,
+ )
+
+
+class UnifiVoucher(ApiItem):
+ """Represents a voucher."""
+
+ raw: UnifiTypedVoucher
+
+ @property
+ def id(self) -> str:
+ """ID of voucher."""
+ return self.raw["_id"]
+
+ @property
+ def site_id(self) -> str:
+ """Site ID."""
+ return self.raw["_id"]
+
+ @property
+ def note(self) -> str:
+ """Note."""
+ return self.raw.get("note") or ""
+
+ @property
+ def code(self) -> str:
+ """Code."""
+ if len(c := self.raw.get("code", "")) > 5:
+ return '%s-%s' % (c[:5], c[5:])
+ return c
+
+ @property
+ def quota(self) -> int:
+ """Nmber of vouchers."""
+ return self.raw.get("quota", 0)
+
+ @property
+ def duration(self) -> timedelta:
+ """Expiration of voucher."""
+ return timedelta(
+ minutes=self.raw.get("duration", 0)
+ )
+
+ @property
+ def qos_overwrite(self) -> bool:
+ """Used count."""
+ return self.raw.get("qos_overwrite", False)
+
+ @property
+ def qos_usage_quota(self) -> int:
+ """Quantity of bytes allowed in MB."""
+ return int(self.raw.get("qos_usage_quota", 0))
+
+ @property
+ def qos_rate_max_up(self) -> int:
+ """Up speed allowed in kbps."""
+ return self.raw.get("qos_rate_max_up", 0)
+
+ @property
+ def qos_rate_max_down(self) -> int:
+ """Down speed allowed in kbps."""
+ return self.raw.get("qos_rate_max_down", 0)
+
+ @property
+ def used(self) -> int:
+ """Number of using; 0 = unlimited."""
+ return self.raw.get("used", 0)
+
+ @property
+ def create_time(self) -> datetime:
+ """Create datetime."""
+ return dt_util.as_local(
+ datetime.fromtimestamp(self.raw["create_time"])
  )
 
+ @property
+ def start_time(self) -> datetime | None:
+ """Start datetime."""
+ if "start_time" in self.raw:
+ return dt_util.as_local(
+ datetime.fromtimestamp(self.raw["start_time"])
+ )
+ return None
+
+ @property
+ def end_time(self) -> datetime | None:
+ """End datetime."""
+ if "end_time" in self.raw:
+ return dt_util.as_local(
+ datetime.fromtimestamp(self.raw["end_time"])
+ )
+ return None
+
+ @property
+ def for_hotspot(self) -> bool:
+ """For hotspot."""
+ return self.raw.get("for_hotspot", False)
+
+ @property
+ def admin_name(self) -> str:
+ """Admin name."""
+ return self.raw.get("admin_name", "")
+
+ @property
+ def status(self) -> str:
+ """Status."""
+ return self.raw.get("status", "")
+
+ @property
+ def status_expires(self) -> int | None:
+ """Status expires."""
+ if self.raw.get("status_expires", 0) > 0:
+ return timedelta(
+ seconds=self.raw.get("status_expires")
+ )
+ return None
+
+
+class UnifiVouchers(APIHandler[UnifiVoucher]):
+ """Represent UniFi vouchers."""
+
+ obj_id_key = "_id"
+ item_cls = UnifiVoucher
+ api_request = UnifiVoucherListRequest.create()
+
 class UnifiVoucherApiClient:
  """API Client."""
 
@@ -132,7 +286,7 @@ def __init__(
  verify_ssl=False,
  cookie_jar=CookieJar(unsafe=True),
  )
- self.api = aiounifi.Controller(
+ self.controller = aiounifi.Controller(
  Configuration(
  session,
  host=host,
@@ -155,7 +309,7 @@ async def async_reconnect(self) -> None:
  """Try to reconnect UniFi Network session."""
  try:
  async with asyncio.timeout(5):
- await self.api.login()
+ await self.controller.login()
  except (
  asyncio.TimeoutError,
  aiounifi.BadGateway,
@@ -171,9 +325,9 @@ async def check_api_user(
  _sites = {}
  try:
  async with asyncio.timeout(10):
- await self.api.login()
- await self.api.sites.update()
- for _unique_id, _site in self.api.sites.items():
+ await self.controller.login()
+ await self.controller.sites.update()
+ for _unique_id, _site in self.controller.sites.items():
  # User must have admin or hotspot permissions 
  if _site.role in ("admin", "hotspot"):
  _sites[_unique_id] = _site
@@ -226,4 +380,4 @@ async def request(
  api_request: ApiRequest,
  ) -> TypedApiResponse:
  """Make a request to the API, retry login on failure."""
- return await self.api.request(api_request)
+ return await self.controller.request(api_request)

custom_components/unifi_voucher/button.py

@@ -46,13 +46,20 @@ async def async_setup_entry(
  """Do setup buttons from a config entry created in the integrations UI."""
  coordinator = hass.data[DOMAIN][config_entry.entry_id]
  entity_descriptions = [
- # TODO
- #UnifiVoucherButtonDescription(
- # key=ATTR_REBOOT,
- # translation_key=ATTR_REBOOT,
- # device_class=ButtonDeviceClass.RESTART,
- # press_action=lambda coordinator: coordinator.async_reboot(),
- #),
+ UnifiVoucherButtonDescription(
+ key="update",
+ icon="mdi:update",
+ translation_key="update",
+ device_class=ButtonDeviceClass.UPDATE,
+ press_action=lambda coordinator: coordinator.async_update_vouchers(),
+ ),
+ UnifiVoucherButtonDescription(
+ key="create",
+ icon="mdi:numeric-positive-1",
+ translation_key="create",
+ device_class=ButtonDeviceClass.RESTART,
+ press_action=lambda coordinator: coordinator.async_create_voucher(),
+ ),
  ]
 
  async_add_entities(

custom_components/unifi_voucher/config_flow.py

@@ -184,7 +184,7 @@ async def async_step_site(
  # Abort if site is already configured
  self._async_abort_entries_match(
  {
- CONF_HOST: user_input[CONF_HOST],
+ CONF_HOST: self.data[CONF_HOST],
  CONF_SITE_ID: self.sites[unique_id].name,
  }
  )

custom_components/unifi_voucher/const.py

@@ -24,8 +24,9 @@
 
 UPDATE_INTERVAL = 120
 
-CONF_SITE_ID = "site"
+CONF_SITE_ID = "site_id"
 
 ATTR_EXTRA_STATE_ATTRIBUTES = "extra_state_attributes"
+ATTR_LAST_PULL = "last_pull"
 ATTR_AVAILABLE = "available"
-ATTR_LAST_PULL = "last_pull"
+ATTR_VOUCHER = "voucher"