diff --git a/kasa/modules/module.py b/kasa/modules/module.py index 7340d7e1..ff806a99 100644 --- a/kasa/modules/module.py +++ b/kasa/modules/module.py @@ -43,6 +43,15 @@ class Module(ABC): 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 index d1f96e7e..9877795d 100644 --- a/kasa/modules/usage.py +++ b/kasa/modules/usage.py @@ -21,6 +21,11 @@ class Usage(Module): 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 index 2d2f28ca..a09487d2 100644 --- a/kasa/smartbulb.py +++ b/kasa/smartbulb.py @@ -543,3 +543,8 @@ class SmartBulb(SmartDevice): 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 index 5c24c943..81ea135b 100755 --- a/kasa/smartdevice.py +++ b/kasa/smartdevice.py @@ -28,9 +28,6 @@ from .protocol import TPLinkSmartHomeProtocol _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 @@ class SmartDevice: ) 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 @@ class 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 index 0839bc06..c70f544b 100644 --- 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."""