python-kasa · rytilahti · Sep 14, 2023 · Aug 25, 2023 · Aug 27, 2023 · Aug 30, 2023
diff --git a/kasa/modules/module.py b/kasa/modules/module.py
@@ -43,6 +43,15 @@ def query(self):
         queries to the query that gets executed when Device.update() gets called.
         """
 
+    @property
+    def estimated_query_response_size(self):
+        """Estimated maximum size of query response.
+
+        The inheriting modules implement this to estimate how large a query response
+        will be so that queries can be split should an estimated response be too large
+        """
+        return 256  # Estimate for modules that don't specify
+
     @property
     def data(self):
         """Return the module specific raw data from the last update."""

diff --git a/kasa/modules/usage.py b/kasa/modules/usage.py
@@ -21,6 +21,11 @@ def query(self):
 
         return req
 
+    @property
+    def estimated_query_response_size(self):
+        """Estimated maximum query response size."""
+        return 2048
+
     @property
     def daily_data(self):
         """Return statistics on daily basis."""

diff --git a/kasa/smartbulb.py b/kasa/smartbulb.py
@@ -543,3 +543,8 @@ async def save_preset(self, preset: SmartBulbPreset):
         return await self._query_helper(
             self.LIGHT_SERVICE, "set_preferred_state", preset.dict(exclude_none=True)
         )
+
+    @property
+    def max_device_response_size(self) -> int:
+        """Returns the maximum response size the device can safely construct."""
+        return 4096
diff --git a/kasa/smartdevice.py b/kasa/smartdevice.py
@@ -28,9 +28,6 @@
 
 _LOGGER = logging.getLogger(__name__)
 
-# Certain module queries will crash devices; this list skips those queries
-MODEL_MODULE_SKIPLIST = {"KL125(US)": ["cloud"]}  # Issue #345
-
 
 class DeviceType(Enum):
     """Device type enum."""
@@ -337,20 +334,32 @@ async def _modular_update(self, req: dict) -> None:
             )
             self.add_module("emeter", Emeter(self, self.emeter_type))
 
+        request_list = []
+        est_response_size = 1024 if "system" in req else 0
         for module_name, module in self.modules.items():
             if not module.is_supported:
                 _LOGGER.debug("Module %s not supported, skipping" % module)
                 continue
-            modules_to_skip = MODEL_MODULE_SKIPLIST.get(self.model, [])
-            if module_name in modules_to_skip:
-                _LOGGER.debug(f"Module {module} is excluded for {self.model}, skipping")
-                continue
+
+            est_response_size += module.estimated_query_response_size
+            if est_response_size > self.max_device_response_size:
+                request_list.append(req)
+                req = {}
+                est_response_size = module.estimated_query_response_size
 
             q = module.query()
             _LOGGER.debug("Adding query for %s: %s", module, q)
             req = merge(req, q)
+        request_list.append(req)
 
-        self._last_update = await self.protocol.query(req)
+        responses = [
+            await self.protocol.query(request) for request in request_list if request
+        ]
+
+        update: Dict = {}
+        for response in responses:
+            update = {**update, **response}
+        self._last_update = update
 
     def update_from_discover_info(self, info):
         """Update state from info from the discover call."""
@@ -658,6 +667,11 @@ def get_plug_by_index(self, index: int) -> "SmartDevice":
             )
         return self.children[index]
 
+    @property
+    def max_device_response_size(self) -> int:
+        """Returns the maximum response size the device can safely construct."""
+        return 16 * 1024
+
     @property
     def device_type(self) -> DeviceType:
         """Return the device type."""

diff --git a/kasa/tests/test_smartdevice.py b/kasa/tests/test_smartdevice.py
@@ -39,7 +39,9 @@ async def test_initial_update_emeter(dev, mocker):
     dev._last_update = None
     spy = mocker.spy(dev.protocol, "query")
     await dev.update()
-    assert spy.call_count == 2 + len(dev.children)
+    # Devices with small buffers may require 3 queries
+    expected_queries = 2 if dev.max_device_response_size > 4096 else 3
+    assert spy.call_count == expected_queries + len(dev.children)
 
 
 @no_emeter
@@ -164,6 +166,17 @@ async def test_features(dev):
         assert dev.features == set()
 
 
+async def test_max_device_response_size(dev):
+    """Make sure every device return has a set max response size."""
+    assert dev.max_device_response_size > 0
+
+
+async def test_estimated_response_sizes(dev):
+    """Make sure every module has an estimated response size set."""
+    for mod in dev.modules.values():
+        assert mod.estimated_query_response_size > 0
+
+
 @pytest.mark.parametrize("device_class", smart_device_classes)
 def test_device_class_ctors(device_class):
     """Make sure constructor api not broken for new and existing SmartDevices."""