From 4fc7694d492f08b9df5d7fa1924ed51035cff11c Mon Sep 17 00:00:00 2001
From: Yevhenii Vaskivskyi <yevhenii@vaskivskyi.com>
Date: Mon, 20 Nov 2023 18:58:36 +0100
Subject: [PATCH] Add `force_clients` possibility (#386)

---
 asusrouter/asusrouter.py | 45 +++++++++++++++++++++++++++++++++++++++-
 asusrouter/const.py      |  1 +
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/asusrouter/asusrouter.py b/asusrouter/asusrouter.py
index bfad983..84fbfe7 100644
--- a/asusrouter/asusrouter.py
+++ b/asusrouter/asusrouter.py
@@ -13,7 +13,11 @@
 import aiohttp
 
 from asusrouter.connection import Connection
-from asusrouter.const import DEFAULT_CACHE_TIME, DEFAULT_TIMEOUT
+from asusrouter.const import (
+    DEFAULT_CACHE_TIME,
+    DEFAULT_CLIENTS_WAITTIME,
+    DEFAULT_TIMEOUT,
+)
 from asusrouter.error import (
     AsusRouter404Error,
     AsusRouterAccessError,
@@ -54,6 +58,7 @@
     save_state,
     set_state,
 )
+from asusrouter.modules.system import AsusSystem
 from asusrouter.tools import legacy
 from asusrouter.tools.converters import safe_list
 from asusrouter.tools.readers import merge_dicts
@@ -77,6 +82,10 @@ class AsusRouter:
     # ID from the last called service
     _last_id: Optional[int] = None
 
+    # Force clients
+    _force_clients: bool = False
+    _force_clients_waittime: float = DEFAULT_CLIENTS_WAITTIME
+
     def __init__(
         self,
         hostname: str,
@@ -449,6 +458,20 @@ def _drop_data(self, datatype: AsusData, endpoint: Endpoint) -> bool:
 
         return False
 
+    async def _check_prerequisites(self, datatype: AsusData) -> None:
+        """Check prerequisites before fetching data."""
+
+        # Allow forcing clients update if set by the user
+        if datatype == AsusData.CLIENTS:
+            if self._force_clients:
+                _LOGGER.debug("Forcing clients update")
+                await self.async_set_state(AsusSystem.UPDATE_CLIENTS)
+                _LOGGER.debug(
+                    "Waiting for clients to be updated: %s s",
+                    self._force_clients_waittime,
+                )
+                await asyncio.sleep(self._force_clients_waittime)
+
     def _check_state(self, datatype: Optional[AsusData]) -> None:
         """Make sure the state object is available."""
 
@@ -499,6 +522,9 @@ async def async_get_data(self, datatype: AsusData, force: bool = False) -> Any:
         # Mark the data as active
         self._state[datatype].start()
 
+        # Check prerequisites
+        await self._check_prerequisites(datatype)
+
         # Get the data finder
         data_finder = self._where_to_get_data(datatype)
 
@@ -848,6 +874,23 @@ def connected(self) -> bool:
     # <-- Properties
     # ---------------------------
 
+    # ---------------------------
+    # Additional settings -->
+    # ---------------------------
+
+    def set_force_clients(
+        self, value: bool = True, waittime: Optional[float] = None
+    ) -> None:
+        """Set force clients."""
+
+        self._force_clients = value
+        if waittime:
+            self._force_clients_waittime = waittime
+
+    # ---------------------------
+    # <-- Additional settings
+    # ---------------------------
+
     # ---------------------------
     # General management -->
     # ---------------------------
diff --git a/asusrouter/const.py b/asusrouter/const.py
index 8777f56..61662e0 100644
--- a/asusrouter/const.py
+++ b/asusrouter/const.py
@@ -21,6 +21,7 @@ class ContentType(str, Enum):
 
 # Library defaults
 DEFAULT_CACHE_TIME = 5.0
+DEFAULT_CLIENTS_WAITTIME = 5.0
 DEFAULT_SLEEP_TIME = 0.1
 DEFAULT_TIMEOUT = 10.0
 DEFAULT_TIMEOUT_CONNECTION = 180.0