diff --git a/roborock/api.py b/roborock/api.py index 8e7b097..f7d7591 100644 --- a/roborock/api.py +++ b/roborock/api.py @@ -186,10 +186,18 @@ def __init__(self, endpoint: str, device_info: DeviceData, queue_timeout: int = self._listeners: list[Callable[[str, CacheableAttribute, RoborockBase], None]] = [] self.is_available: bool = True self.queue_timeout = queue_timeout + self._status_type: type[Status] = ModelStatus.get( + self.device_info.model, S7MaxVStatus + ) def __del__(self) -> None: self.release() + @property + def status_type(self) -> type[Status]: + """Gets the status type for this device""" + return self._status_type + def release(self): self.sync_disconnect() [item.stop() for item in self.cache.values()] @@ -253,9 +261,6 @@ def on_message_received(self, messages: list[RoborockMessage]) -> None: data_protocol = RoborockDataProtocol(int(data_point_number)) self._logger.debug(f"Got device update for {data_protocol.name}: {data_point}") if data_protocol in ROBOROCK_DATA_STATUS_PROTOCOL: - _cls: type[Status] = ModelStatus.get( - self.device_info.model, S7MaxVStatus - ) # Default to S7 MAXV if we don't have the data if self.cache[CacheableAttribute.status].value is None: self._logger.debug( f"Got status update({data_protocol.name}) before get_status was called." @@ -263,7 +268,7 @@ def on_message_received(self, messages: list[RoborockMessage]) -> None: self.cache[CacheableAttribute.status]._value = {} value = self.cache[CacheableAttribute.status].value value[data_protocol.name] = data_point - status = _cls.from_dict(value) + status = self._status_type.from_dict(value) for listener in self._listeners: listener(self.device_info.device.duid, CacheableAttribute.status, status) elif data_protocol in ROBOROCK_DATA_CONSUMABLE_PROTOCOL: @@ -406,10 +411,7 @@ async def send_command( return response async def get_status(self) -> Status | None: - _cls: type[Status] = ModelStatus.get( - self.device_info.model, S7MaxVStatus - ) # Default to S7 MAXV if we don't have the data - return _cls.from_dict(await self.cache[CacheableAttribute.status].async_value()) + return self._status_type.from_dict(await self.cache[CacheableAttribute.status].async_value()) async def get_dnd_timer(self) -> DnDTimer | None: return DnDTimer.from_dict(await self.cache[CacheableAttribute.dnd_timer].async_value())