Merge pull request #18 from RobertD502/add_options_flow

Add options flow

README.md

@@ -46,9 +46,7 @@ Custom Home Assistant component for controlling and monitoring PetKit devices an
 
 - **If using another PetKit integration that uses the petkit domain, you will need to delete it prior to installing this integration.**
 
-- **The current polling interval is set to 2 minutes. If you would like to set a different polling interval, change the `DEFAULT_SCAN_INTERVAL` in the constants.py file (in seconds).**
-
-- If you are running Home Assistant as a Docker container, the `TZ` environment variable must be set.
+- **If you are running Home Assistant as a Docker container, the `TZ` environment variable must be set.**
 
 ## Important - Please Read:
 ### Note About PetKit Account:
@@ -96,6 +94,9 @@ Alternatively, follow the steps below:
 1. Navigate to the Home Assistant Integrations page (Settings --> Devices & Services)
 2. Click the `+ ADD INTEGRATION` button in the lower right-hand corner
 3. Search for `PetKit`
+4. If using a PetKit Asia account, be sure the PetKit Asia account option is checked.
+
+**The current polling interval is set to 2 minutes. If you would like to set a different polling interval, change the polling interval option (via the UI). Keep in mind, setting the polling interval too short may result in your account getting rate limited/blocked.**
 
 # Devices
 

custom_components/petkit/__init__.py

@@ -2,10 +2,10 @@
 from __future__ import annotations
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_EMAIL
+from homeassistant.const import CONF_EMAIL, CONF_PASSWORD
 from homeassistant.core import HomeAssistant
 
-from .const import DOMAIN, LOGGER, PLATFORMS
+from .const import ASIA_ACCOUNT, DOMAIN, LOGGER, PETKIT_COORDINATOR, PLATFORMS, POLLING_INTERVAL, UPDATE_LISTENER
 from .coordinator import PetKitDataUpdateCoordinator
 from .util import async_validate_api, NoDevicesError
 
@@ -15,10 +15,15 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     coordinator = PetKitDataUpdateCoordinator(hass, entry)
     await coordinator.async_config_entry_first_refresh()
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = {
+        PETKIT_COORDINATOR: coordinator
+    }
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
+    update_listener = entry.add_update_listener(async_update_options)
+    hass.data[DOMAIN][entry.entry_id][UPDATE_LISTENER] = update_listener
+
     return True
 
 
@@ -30,3 +35,32 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         if not hass.data[DOMAIN]:
             del hass.data[DOMAIN]
     return unload_ok
+
+async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+    """Migrate old entry."""
+
+    if entry.version == 1:
+        email = entry.data[CONF_EMAIL]
+        password = entry.data[CONF_PASSWORD]
+
+        LOGGER.debug('Migrating PetKit config entry')
+        entry.version = 2
+
+        hass.config_entries.async_update_entry(
+            entry,
+            data={
+                CONF_EMAIL: email,
+                CONF_PASSWORD: password,
+            },
+            options={
+                ASIA_ACCOUNT: False,
+                POLLING_INTERVAL: 120,
+            },
+        )
+
+    return True
+
+async def async_update_options(hass: HomeAssistant, entry: ConfigEntry) -> None:
+    """ Update options. """
+
+    await hass.config_entries.async_reload(entry.entry_id)

custom_components/petkit/binary_sensor.py

@@ -12,7 +12,7 @@
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from .const import DOMAIN, FEEDERS, LITTER_BOXES, WATER_FOUNTAINS
+from .const import DOMAIN, FEEDERS, LITTER_BOXES, PETKIT_COORDINATOR, WATER_FOUNTAINS
 from .coordinator import PetKitDataUpdateCoordinator
 
 
@@ -21,7 +21,7 @@ async def async_setup_entry(
 ) -> None:
     """Set Up PetKit Binary Sensor Entities."""
 
-    coordinator: PetKitDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator: PetKitDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id][PETKIT_COORDINATOR]
 
     binary_sensors = []
 

custom_components/petkit/button.py

@@ -15,7 +15,7 @@
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from .const import DOMAIN, FEEDERS, LITTER_BOXES, WATER_FOUNTAINS
+from .const import DOMAIN, FEEDERS, LITTER_BOXES, PETKIT_COORDINATOR, WATER_FOUNTAINS
 from .coordinator import PetKitDataUpdateCoordinator
 from .exceptions import PetKitBluetoothError
 
@@ -25,7 +25,7 @@ async def async_setup_entry(
 ) -> None:
     """Set Up PetKit Button Entities."""
 
-    coordinator: PetKitDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator: PetKitDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id][PETKIT_COORDINATOR]
 
     buttons = []
 

custom_components/petkit/config_flow.py

@@ -9,28 +9,38 @@
 
 from homeassistant import config_entries
 from homeassistant.const import CONF_EMAIL, CONF_PASSWORD
+from homeassistant.core import callback
 from homeassistant.data_entry_flow import FlowResult
 import homeassistant.helpers.config_validation as cv
 
-from .const import DEFAULT_NAME, DOMAIN
+from .const import ASIA_ACCOUNT, DEFAULT_NAME, DOMAIN, POLLING_INTERVAL
 from .util import NoDevicesError, async_validate_api
 
 
 DATA_SCHEMA = vol.Schema(
     {
         vol.Required(CONF_EMAIL): cv.string,
         vol.Required(CONF_PASSWORD): cv.string,
+        vol.Optional(ASIA_ACCOUNT): cv.boolean,
     }
 )
 
 
 class PetKitConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     """Handle a config flow for PetKit integration."""
 
-    VERSION = 1
+    VERSION = 2
 
     entry: config_entries.ConfigEntry | None
 
+    @staticmethod
+    @callback
+    def async_get_options_flow(
+        config_entry: config_entries.ConfigEntry,
+    ) -> PetKitOptionsFlowHandler:
+        """Get the options flow for this handler."""
+        return PetKitOptionsFlowHandler(config_entry)
+
     async def async_step_reauth(self, entry_data: Mapping[str, Any]) -> FlowResult:
         """Handle re-authentication with PetKit."""
 
@@ -47,8 +57,9 @@ async def async_step_reauth_confirm(
         if user_input:
             email = user_input[CONF_EMAIL]
             password = user_input[CONF_PASSWORD]
+            asia_account = user_input[ASIA_ACCOUNT] if ASIA_ACCOUNT in user_input else False
             try:
-                await async_validate_api(self.hass, email, password)
+                await async_validate_api(self.hass, email, password, asia_account)
             except AuthError:
                 errors["base"] = "invalid_auth"
             except ConnectionError:
@@ -67,6 +78,10 @@ async def async_step_reauth_confirm(
                         CONF_EMAIL: email,
                         CONF_PASSWORD: password,
                     },
+                    options={
+                        ASIA_ACCOUNT: asia_account,
+                        POLLING_INTERVAL: self.entry.options[POLLING_INTERVAL],
+                    }
                 )
 
                 await self.hass.config_entries.async_reload(self.entry.entry_id)
@@ -88,8 +103,9 @@ async def async_step_user(
         if user_input:
             email = user_input[CONF_EMAIL]
             password = user_input[CONF_PASSWORD]
+            asia_account = user_input[ASIA_ACCOUNT] if ASIA_ACCOUNT in user_input else False
             try:
-                await async_validate_api(self.hass, email, password)
+                await async_validate_api(self.hass, email, password, asia_account)
             except AuthError:
                 errors["base"] = "invalid_auth"
             except ConnectionError:
@@ -104,11 +120,52 @@ async def async_step_user(
 
                 return self.async_create_entry(
                     title=DEFAULT_NAME,
-                    data={CONF_EMAIL: email, CONF_PASSWORD: password},
+                    data={
+                        CONF_EMAIL: email,
+                        CONF_PASSWORD: password
+                    },
+                    options={
+                        ASIA_ACCOUNT: asia_account,
+                        POLLING_INTERVAL: 120,
+                    }
                 )
 
         return self.async_show_form(
             step_id="user",
             data_schema=DATA_SCHEMA,
             errors=errors,
         )
+
+
+class PetKitOptionsFlowHandler(config_entries.OptionsFlow):
+    """ Handle PetKit integration options. """
+
+    def __init__(self, config_entry: config_entries.ConfigEntry) -> None:
+        """Initialize options flow."""
+        self.config_entry = config_entry
+
+    async def async_step_init(self, user_input=None):
+        """ Manage options. """
+        return await self.async_step_petkit_options()
+
+    async def async_step_petkit_options(self, user_input=None):
+        """Manage the PetKit options."""
+        if user_input is not None:
+            return self.async_create_entry(title="", data=user_input)
+
+        options = {
+            vol.Required(
+                ASIA_ACCOUNT,
+                default=self.config_entry.options.get(
+                    ASIA_ACCOUNT, False
+                ),
+            ): bool,
+            vol.Required(
+                POLLING_INTERVAL,
+                default=self.config_entry.options.get(
+                    POLLING_INTERVAL, 120
+                ),
+            ): int,
+        }
+
+        return self.async_show_form(step_id="petkit_options", data_schema=vol.Schema(options))

custom_components/petkit/const.py

@@ -12,7 +12,6 @@
 LOGGER = logging.getLogger(__package__)
 
 
-DEFAULT_SCAN_INTERVAL = 120
 DOMAIN = "petkit"
 PLATFORMS = [
     Platform.BINARY_SENSOR,
@@ -26,7 +25,11 @@
 ]
 
 DEFAULT_NAME = "PetKit"
+ASIA_ACCOUNT = "asia_account"
+PETKIT_COORDINATOR = "petkit_coordinator"
+POLLING_INTERVAL = "polling_interval"
 TIMEOUT = 20
+UPDATE_LISTENER = "update_listener"
 
 PETKIT_ERRORS = (
     asyncio.TimeoutError,

custom_components/petkit/coordinator.py

@@ -15,7 +15,7 @@
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
-from .const import DEFAULT_SCAN_INTERVAL, DOMAIN, LOGGER, TIMEOUT
+from .const import ASIA_ACCOUNT, DOMAIN, LOGGER, POLLING_INTERVAL, TIMEOUT
 
 
 class PetKitDataUpdateCoordinator(DataUpdateCoordinator):
@@ -30,13 +30,14 @@ def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
             entry.data[CONF_EMAIL],
             entry.data[CONF_PASSWORD],
             session=async_get_clientsession(hass),
+            asia_account=entry.options[ASIA_ACCOUNT],
             timeout=TIMEOUT,
         )
         super().__init__(
             hass,
             LOGGER,
             name=DOMAIN,
-            update_interval=timedelta(seconds=DEFAULT_SCAN_INTERVAL),
+            update_interval=timedelta(seconds=entry.options[POLLING_INTERVAL]),
         )
 
     async def _async_update_data(self) -> PetKitData:

custom_components/petkit/fan.py

@@ -15,6 +15,7 @@
 
 from .const import (
     DOMAIN,
+    PETKIT_COORDINATOR,
     PURIFIERS,
     PURIFIER_MODES,
     PURIFIER_MODE_NAMED,
@@ -29,7 +30,7 @@ async def async_setup_entry(
 ) -> None:
     """Set Up PetKit Fan Entities."""
 
-    coordinator: PetKitDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator: PetKitDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id][PETKIT_COORDINATOR]
 
     fans = []
 

custom_components/petkit/manifest.json

@@ -8,6 +8,6 @@
   "integration_type": "hub",
   "iot_class": "cloud_polling",
   "issue_tracker": "https://github.com/RobertD502/home-assistant-petkit/issues",
-  "requirements": ["petkitaio==0.1.4", "tzlocal>=4.2"],
-  "version": "0.1.4.1"
+  "requirements": ["petkitaio==0.1.5", "tzlocal>=4.2"],
+  "version": "0.1.5"
 }

custom_components/petkit/number.py

@@ -20,7 +20,7 @@
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 from homeassistant.util.unit_system import METRIC_SYSTEM
 
-from .const import DOMAIN, FEEDERS, LITTER_BOXES
+from .const import DOMAIN, FEEDERS, LITTER_BOXES, PETKIT_COORDINATOR
 from .coordinator import PetKitDataUpdateCoordinator
 
 
@@ -29,7 +29,7 @@ async def async_setup_entry(
 ) -> None:
     """Set Up PetKit Number Entities."""
 
-    coordinator: PetKitDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator: PetKitDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id][PETKIT_COORDINATOR]
 
     numbers = []
 

custom_components/petkit/select.py

@@ -27,6 +27,7 @@
     LITTER_TYPE_NAMED,
     MANUAL_FEED_NAMED,
     MINI_FEEDER_MANUAL_FEED_OPTIONS,
+    PETKIT_COORDINATOR,
     WATER_FOUNTAINS,
     WF_MODE_COMMAND,
     WF_MODE_NAMED,
@@ -49,7 +50,7 @@ async def async_setup_entry(
 ) -> None:
     """Set Up PetKit Select Entities."""
 
-    coordinator: PetKitDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator: PetKitDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id][PETKIT_COORDINATOR]
 
     selects = []
 

custom_components/petkit/sensor.py

@@ -32,6 +32,7 @@
     DOMAIN,
     FEEDERS,
     LITTER_BOXES,
+    PETKIT_COORDINATOR,
     PURIFIERS,
     WATER_FOUNTAINS
 )
@@ -51,7 +52,7 @@ async def async_setup_entry(
 ) -> None:
     """Set Up PetKit Sensor Entities."""
 
-    coordinator: PetKitDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator: PetKitDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id][PETKIT_COORDINATOR]
 
     sensors = []