Add sensors to Google Drive

homeassistant/components/google_drive/__init__.py

@@ -6,7 +6,7 @@
 
 from google_drive_api.exceptions import GoogleDriveApiError
 
-from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers import instance_id
@@ -19,13 +19,13 @@
 
 from .api import AsyncConfigEntryAuth, DriveClient
 from .const import DOMAIN
+from .coordinator import GoogleDriveConfigEntry, GoogleDriveDataUpdateCoordinator
 
 DATA_BACKUP_AGENT_LISTENERS: HassKey[list[Callable[[], None]]] = HassKey(
     f"{DOMAIN}.backup_agent_listeners"
 )
 
-
-type GoogleDriveConfigEntry = ConfigEntry[DriveClient]
+PLATFORMS = (Platform.SENSOR,)
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: GoogleDriveConfigEntry) -> bool:
@@ -41,11 +41,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: GoogleDriveConfigEntry)
     await auth.async_get_access_token()
 
     client = DriveClient(await instance_id.async_get(hass), auth)
-    entry.runtime_data = client
 
     # Test we can access Google Drive and raise if not
     try:
-        await client.async_create_ha_root_folder_if_not_exists()
+        folder_id, _ = await client.async_create_ha_root_folder_if_not_exists()
     except GoogleDriveApiError as err:
         raise ConfigEntryNotReady from err
 
@@ -55,11 +54,21 @@ def async_notify_backup_listeners() -> None:
 
     entry.async_on_unload(entry.async_on_state_change(async_notify_backup_listeners))
 
+    entry.runtime_data = GoogleDriveDataUpdateCoordinator(
+        hass, entry=entry, client=client, backup_folder_id=folder_id
+    )
+    await entry.runtime_data.async_config_entry_first_refresh()
+
+    # Set up all platforms for this device/entry.
+    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
+
     return True
 
 
 async def async_unload_entry(
     hass: HomeAssistant, entry: GoogleDriveConfigEntry
 ) -> bool:
     """Unload a config entry."""
+    await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+
     return True

homeassistant/components/google_drive/api.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from collections.abc import AsyncIterator, Callable, Coroutine
+from dataclasses import dataclass
 import json
 import logging
 from typing import Any
@@ -27,6 +28,16 @@
 _LOGGER = logging.getLogger(__name__)
 
 
+@dataclass
+class StorageQuotaData:
+    """Class to represent storage quota data."""
+
+    limit: int
+    usage: int
+    usage_in_drive: int
+    usage_in_trash: int
+
+
 class AsyncConfigEntryAuth(AbstractAuth):
     """Provide Google Drive authentication tied to an OAuth2 based config entry."""
 
@@ -95,6 +106,18 @@ async def async_get_email_address(self) -> str:
         res = await self._api.get_user(params={"fields": "user(emailAddress)"})
         return str(res["user"]["emailAddress"])
 
+    async def async_get_storage_quota(self) -> StorageQuotaData:
+        """Get storage quota of the current user."""
+        res = await self._api.get_user(params={"fields": "storageQuota"})
+
+        storageQuota = res["storageQuota"]
+        return StorageQuotaData(
+            limit=int(storageQuota.get("limit", 0)),
+            usage=int(storageQuota.get("usage", 0)),
+            usage_in_drive=int(storageQuota.get("usageInDrive", 0)),
+            usage_in_trash=int(storageQuota.get("usageInTrash", 0)),
+        )
+
     async def async_create_ha_root_folder_if_not_exists(self) -> tuple[str, str]:
         """Create Home Assistant folder if it doesn't exist."""
         fields = "id,name"

homeassistant/components/google_drive/backup.py

@@ -68,7 +68,8 @@ def __init__(self, config_entry: GoogleDriveConfigEntry) -> None:
         assert config_entry.unique_id
         self.name = config_entry.title
         self.unique_id = slugify(config_entry.unique_id)
-        self._client = config_entry.runtime_data
+        self._coordinator = config_entry.runtime_data
+        self._client = config_entry.runtime_data.client
 
     async def async_upload_backup(
         self,
@@ -84,6 +85,7 @@ async def async_upload_backup(
         """
         try:
             await self._client.async_upload_backup(open_stream, backup)
+            await self._coordinator.async_request_refresh()
         except (GoogleDriveApiError, HomeAssistantError, TimeoutError) as err:
             raise BackupAgentError(f"Failed to upload backup: {err}") from err
 

homeassistant/components/google_drive/const.py

@@ -2,4 +2,8 @@
 
 from __future__ import annotations
 
+from datetime import timedelta
+
 DOMAIN = "google_drive"
+
+SCAN_INTERVAL = timedelta(hours=6)

homeassistant/components/google_drive/coordinator.py

@@ -0,0 +1,77 @@
+"""DataUpdateCoordinator for WLED."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import logging
+
+from google_drive_api.exceptions import GoogleDriveApiError
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .api import DriveClient, StorageQuotaData
+from .const import DOMAIN, SCAN_INTERVAL
+
+type GoogleDriveConfigEntry = ConfigEntry[GoogleDriveDataUpdateCoordinator]
+
+_LOGGER = logging.getLogger(__name__)
+
+
+@dataclass
+class GoogleDriveCoordinatorData:
+    """Class to hold coordinator data."""
+
+    storage_quota: StorageQuotaData
+    email_address: str
+
+
+class GoogleDriveDataUpdateCoordinator(
+    DataUpdateCoordinator[GoogleDriveCoordinatorData]
+):
+    """Class to manage fetching Google Drive data from single endpoint."""
+
+    client: DriveClient
+    config_entry: GoogleDriveConfigEntry
+    _email_address: str
+    backup_folder_id: str
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        *,
+        client: DriveClient,
+        backup_folder_id: str,
+        entry: GoogleDriveConfigEntry,
+    ) -> None:
+        """Initialize Google Drive data updater."""
+        self.client = client
+        self.backup_folder_id = backup_folder_id
+
+        super().__init__(
+            hass,
+            _LOGGER,
+            config_entry=entry,
+            name=DOMAIN,
+            update_interval=SCAN_INTERVAL,
+        )
+
+    async def _async_setup(self) -> None:
+        """Do initialization logic."""
+        self._email_address = await self.client.async_get_email_address()
+
+    async def _async_update_data(self) -> GoogleDriveCoordinatorData:
+        """Fetch data from Google Drive."""
+        try:
+            storage_quota = await self.client.async_get_storage_quota()
+            return GoogleDriveCoordinatorData(
+                storage_quota=storage_quota,
+                email_address=self._email_address,
+            )
+        except GoogleDriveApiError as error:
+            raise UpdateFailed(
+                translation_domain=DOMAIN,
+                translation_key="invalid_response_google_drive_error",
+                translation_placeholders={"error": str(error)},
+            ) from error

homeassistant/components/google_drive/entity.py

@@ -0,0 +1,25 @@
+"""Define the Google Drive entity."""
+
+from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
+
+from .const import DOMAIN
+from .coordinator import GoogleDriveDataUpdateCoordinator
+
+
+class GoogleDriveEntity(CoordinatorEntity[GoogleDriveDataUpdateCoordinator]):
+    """Defines a base Google Drive entity."""
+
+    _attr_has_entity_name = True
+
+    @property
+    def device_info(self) -> DeviceInfo:
+        """Return device information about this Google Drive device."""
+        return DeviceInfo(
+            identifiers={(DOMAIN, str(self.coordinator.config_entry.unique_id))},
+            name=self.coordinator.data.email_address,
+            manufacturer="Google",
+            model="Google Drive",
+            configuration_url=f"https://drive.google.com/drive/folders/{self.coordinator.backup_folder_id}",
+            entry_type=DeviceEntryType.SERVICE,
+        )

homeassistant/components/google_drive/sensor.py

@@ -0,0 +1,116 @@
+"""Support for GoogleDrive sensors."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from datetime import datetime
+
+from homeassistant.components.sensor import (
+    SensorDeviceClass,
+    SensorEntity,
+    SensorEntityDescription,
+)
+from homeassistant.const import EntityCategory, UnitOfInformation
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
+from homeassistant.helpers.typing import StateType
+from homeassistant.util import slugify
+
+from .coordinator import (
+    GoogleDriveConfigEntry,
+    GoogleDriveCoordinatorData,
+    GoogleDriveDataUpdateCoordinator,
+)
+from .entity import GoogleDriveEntity
+
+# Coordinator is used to centralize the data updates
+PARALLEL_UPDATES = 0
+
+
+@dataclass(frozen=True, kw_only=True)
+class GoogleDriveSensorEntityDescription(SensorEntityDescription):
+    """Describes GoogleDrive sensor entity."""
+
+    value_fn: Callable[[GoogleDriveCoordinatorData], datetime | StateType]
+
+
+SENSORS: tuple[GoogleDriveSensorEntityDescription, ...] = (
+    GoogleDriveSensorEntityDescription(
+        key="storage_total",
+        translation_key="storage_total",
+        native_unit_of_measurement=UnitOfInformation.BYTES,
+        suggested_unit_of_measurement=UnitOfInformation.GIBIBYTES,
+        suggested_display_precision=0,
+        device_class=SensorDeviceClass.DATA_SIZE,
+        entity_category=EntityCategory.DIAGNOSTIC,
+        value_fn=lambda data: data.storage_quota.limit,
+    ),
+    GoogleDriveSensorEntityDescription(
+        key="storage_used",
+        translation_key="storage_used",
+        native_unit_of_measurement=UnitOfInformation.BYTES,
+        suggested_unit_of_measurement=UnitOfInformation.GIBIBYTES,
+        suggested_display_precision=0,
+        device_class=SensorDeviceClass.DATA_SIZE,
+        entity_category=EntityCategory.DIAGNOSTIC,
+        value_fn=lambda data: data.storage_quota.usage,
+    ),
+    GoogleDriveSensorEntityDescription(
+        key="storage_used_in_drive",
+        translation_key="storage_used_in_drive",
+        native_unit_of_measurement=UnitOfInformation.BYTES,
+        suggested_unit_of_measurement=UnitOfInformation.GIBIBYTES,
+        suggested_display_precision=0,
+        device_class=SensorDeviceClass.DATA_SIZE,
+        entity_category=EntityCategory.DIAGNOSTIC,
+        value_fn=lambda data: data.storage_quota.usage_in_drive,
+        entity_registry_enabled_default=False,
+    ),
+    GoogleDriveSensorEntityDescription(
+        key="storage_used_in_drive_trash",
+        translation_key="storage_used_in_drive_trash",
+        native_unit_of_measurement=UnitOfInformation.BYTES,
+        suggested_unit_of_measurement=UnitOfInformation.GIBIBYTES,
+        suggested_display_precision=0,
+        device_class=SensorDeviceClass.DATA_SIZE,
+        entity_category=EntityCategory.DIAGNOSTIC,
+        value_fn=lambda data: data.storage_quota.usage_in_trash,
+        entity_registry_enabled_default=False,
+    ),
+)
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    entry: GoogleDriveConfigEntry,
+    async_add_entities: AddConfigEntryEntitiesCallback,
+) -> None:
+    """Set up GoogleDrive sensor based on a config entry."""
+    coordinator = entry.runtime_data
+    async_add_entities(
+        GoogleDriveSensorEntity(coordinator, description) for description in SENSORS
+    )
+
+
+class GoogleDriveSensorEntity(GoogleDriveEntity, SensorEntity):
+    """Defines a Google Drive sensor entity."""
+
+    entity_description: GoogleDriveSensorEntityDescription
+
+    def __init__(
+        self,
+        coordinator: GoogleDriveDataUpdateCoordinator,
+        description: GoogleDriveSensorEntityDescription,
+    ) -> None:
+        """Initialize a Google Drive sensor entity."""
+        super().__init__(coordinator=coordinator)
+        self.entity_description = description
+        self._attr_unique_id = (
+            f"{slugify(coordinator.config_entry.unique_id)}_{description.key}"
+        )
+
+    @property
+    def native_value(self) -> datetime | StateType:
+        """Return the state of the sensor."""
+        return self.entity_description.value_fn(self.coordinator.data)

homeassistant/components/google_drive/strings.json

@@ -42,5 +42,21 @@
         "title": "[%key:common::config_flow::title::reauth%]"
       }
     }
+  },
+  "entity": {
+    "sensor": {
+      "storage_total": {
+        "name": "Total available storage"
+      },
+      "storage_used": {
+        "name": "Used storage"
+      },
+      "storage_used_in_drive": {
+        "name": "Used storage in Drive"
+      },
+      "storage_used_in_drive_trash": {
+        "name": "Used storage in Drive Trash"
+      }
+    }
   }
 }

tests/components/google_drive/conftest.py

@@ -42,6 +42,17 @@ def mock_api() -> Generator[MagicMock]:
         "homeassistant.components.google_drive.api.GoogleDriveApi"
     ) as mock_api_cl:
         mock_api = mock_api_cl.return_value
+        mock_api.get_user = AsyncMock(
+            return_value={
+                "user": {"emailAddress": TEST_USER_EMAIL},
+                "storageQuota": {
+                    "limit": "1000",
+                    "usage": "100",
+                    "usageInDrive": "50",
+                    "usageInTrash": "10",
+                },
+            }
+        )
         yield mock_api