Add new features to air purifiers (#451)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: J. Nick Koston <nick@koston.org>
Co-authored-by: J. Nick Koston <nick@home-assistant.io>

switchbot/__init__.py

@@ -11,6 +11,7 @@ from bleak_retry_connector import (
 from .adv_parser import SwitchbotSupportedType, parse_advertisement_data
 from .const import (
     AirPurifierMode,
+    AirQualityLevel,
     BulbColorMode,
     CeilingLightColorMode,
     ClimateAction,
@@ -67,6 +68,7 @@ from .models import SwitchBotAdvertisement
 
 __all__ = [
     "AirPurifierMode",
+    "AirQualityLevel",
     "BulbColorMode",
     "CeilingLightColorMode",
     "ClimateAction",

switchbot/adv_parser.py

@@ -492,50 +492,50 @@ SUPPORTED_TYPES: dict[str | bytes, SwitchbotSupportedType] = {
         "manufacturer_id": 2409,
     },
     "*": {
-        "modelName": SwitchbotModel.AIR_PURIFIER,
-        "modelFriendlyName": "Air Purifier",
+        "modelName": SwitchbotModel.AIR_PURIFIER_US,
+        "modelFriendlyName": "Air Purifier US",
         "func": process_air_purifier,
         "manufacturer_id": 2409,
     },
     b"\x0a": {
-        "modelName": SwitchbotModel.AIR_PURIFIER,
-        "modelFriendlyName": "Air Purifier",
+        "modelName": SwitchbotModel.AIR_PURIFIER_US,
+        "modelFriendlyName": "Air Purifier US",
         "func": process_air_purifier,
         "manufacturer_id": 2409,
     },
     "+": {
-        "modelName": SwitchbotModel.AIR_PURIFIER,
-        "modelFriendlyName": "Air Purifier",
+        "modelName": SwitchbotModel.AIR_PURIFIER_JP,
+        "modelFriendlyName": "Air Purifier JP",
         "func": process_air_purifier,
         "manufacturer_id": 2409,
     },
     b"\x0b": {
-        "modelName": SwitchbotModel.AIR_PURIFIER,
-        "modelFriendlyName": "Air Purifier",
+        "modelName": SwitchbotModel.AIR_PURIFIER_JP,
+        "modelFriendlyName": "Air Purifier JP",
         "func": process_air_purifier,
         "manufacturer_id": 2409,
     },
     "7": {
-        "modelName": SwitchbotModel.AIR_PURIFIER_TABLE,
-        "modelFriendlyName": "Air Purifier Table",
+        "modelName": SwitchbotModel.AIR_PURIFIER_TABLE_US,
+        "modelFriendlyName": "Air Purifier Table US",
         "func": process_air_purifier,
         "manufacturer_id": 2409,
     },
     b"\x17": {
-        "modelName": SwitchbotModel.AIR_PURIFIER_TABLE,
-        "modelFriendlyName": "Air Purifier Table",
+        "modelName": SwitchbotModel.AIR_PURIFIER_TABLE_US,
+        "modelFriendlyName": "Air Purifier Table US",
         "func": process_air_purifier,
         "manufacturer_id": 2409,
     },
     "8": {
-        "modelName": SwitchbotModel.AIR_PURIFIER_TABLE,
-        "modelFriendlyName": "Air Purifier Table",
+        "modelName": SwitchbotModel.AIR_PURIFIER_TABLE_JP,
+        "modelFriendlyName": "Air Purifier Table JP",
         "func": process_air_purifier,
         "manufacturer_id": 2409,
     },
     b"\x18": {
-        "modelName": SwitchbotModel.AIR_PURIFIER_TABLE,
-        "modelFriendlyName": "Air Purifier Table",
+        "modelName": SwitchbotModel.AIR_PURIFIER_TABLE_JP,
+        "modelFriendlyName": "Air Purifier Table JP",
         "func": process_air_purifier,
         "manufacturer_id": 2409,
     },

switchbot/const/__init__.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from ..enum import StrEnum
-from .air_purifier import AirPurifierMode
+from .air_purifier import AirPurifierMode, AirQualityLevel
 from .climate import ClimateAction, ClimateMode, SmartThermostatRadiatorMode
 from .evaporative_humidifier import (
     HumidifierAction,
@@ -85,8 +85,10 @@ class SwitchbotModel(StrEnum):
     K10_VACUUM = "K10+ Vacuum"
     K10_PRO_VACUUM = "K10+ Pro Vacuum"
     K10_PRO_COMBO_VACUUM = "K10+ Pro Combo Vacuum"
-    AIR_PURIFIER = "Air Purifier"
-    AIR_PURIFIER_TABLE = "Air Purifier Table"
+    AIR_PURIFIER_US = "Air Purifier US"
+    AIR_PURIFIER_JP = "Air Purifier JP"
+    AIR_PURIFIER_TABLE_US = "Air Purifier Table US"
+    AIR_PURIFIER_TABLE_JP = "Air Purifier Table JP"
     HUB3 = "Hub3"
     LOCK_ULTRA = "Lock Ultra"
     LOCK_LITE = "Lock Lite"
@@ -115,6 +117,7 @@ __all__ = [
     "DEFAULT_RETRY_TIMEOUT",
     "DEFAULT_SCAN_TIMEOUT",
     "AirPurifierMode",
+    "AirQualityLevel",
     "BulbColorMode",
     "CeilingLightColorMode",
     "ClimateAction",

switchbot/devices/air_purifier.py

@@ -3,17 +3,19 @@
 from __future__ import annotations
 
 import logging
-import struct
-from typing import Any
+from typing import Any, ClassVar
 
 from bleak.backends.device import BLEDevice
 
 from ..adv_parsers.air_purifier import get_air_purifier_mode
 from ..const import SwitchbotModel
 from ..const.air_purifier import AirPurifierMode, AirQualityLevel
+from ..const.light import ColorMode
+from ..helpers import _UNPACK_UINT16_BE
+from .base_light import SwitchbotSequenceBaseLight
 from .device import (
     SwitchbotEncryptedDevice,
-    SwitchbotSequenceDevice,
+    SwitchbotOperationError,
     update_after_operation,
 )
 
@@ -30,13 +32,51 @@ COMMAND_SET_MODE = {
     AirPurifierMode.PET.name.lower(): f"{COMMAND_HEAD}01010400",
 }
 DEVICE_GET_BASIC_SETTINGS_KEY = "570f4d81"
+COMMAND_SET_PERCENTAGE = f"{COMMAND_HEAD}02{{percentage:02x}}"
+READ_LED_SETTINGS_COMMAND = "570f4d05"
+READ_LED_STATUS_COMMAND = "570f4d07"
 
 
-class SwitchbotAirPurifier(SwitchbotSequenceDevice, SwitchbotEncryptedDevice):
+class SwitchbotAirPurifier(SwitchbotSequenceBaseLight, SwitchbotEncryptedDevice):
     """Representation of a Switchbot Air Purifier."""
 
     _turn_on_command = f"{COMMAND_HEAD}010100"
     _turn_off_command = f"{COMMAND_HEAD}010000"
+    _open_child_lock_command = f"{COMMAND_HEAD}0301"
+    _close_child_lock_command = f"{COMMAND_HEAD}0300"
+    _open_wireless_charging_command = f"{COMMAND_HEAD}0d01"
+    _close_wireless_charging_command = f"{COMMAND_HEAD}0d00"
+    _open_light_sensitive_switch_command = f"{COMMAND_HEAD}0702"
+    _turn_led_on_command = f"{COMMAND_HEAD}0701"
+    _turn_led_off_command = f"{COMMAND_HEAD}0700"
+    _set_rgb_command = _set_brightness_command = f"{COMMAND_HEAD}0501{{}}"
+    _get_basic_info_command = [
+        DEVICE_GET_BASIC_SETTINGS_KEY,
+        READ_LED_SETTINGS_COMMAND,
+        READ_LED_STATUS_COMMAND,
+    ]
+
+    _PM25_MODELS: ClassVar[frozenset[SwitchbotModel]] = frozenset(
+        {
+            SwitchbotModel.AIR_PURIFIER_US,
+            SwitchbotModel.AIR_PURIFIER_TABLE_US,
+        }
+    )
+
+    _LEVEL_MODES: ClassVar[frozenset[str]] = frozenset(
+        {
+            AirPurifierMode.LEVEL_1.name.lower(),
+            AirPurifierMode.LEVEL_2.name.lower(),
+            AirPurifierMode.LEVEL_3.name.lower(),
+        }
+    )
+
+    _WIRELESS_MODELS: ClassVar[frozenset[SwitchbotModel]] = frozenset(
+        {
+            SwitchbotModel.AIR_PURIFIER_TABLE_US,
+            SwitchbotModel.AIR_PURIFIER_TABLE_JP,
+        }
+    )
 
     def __init__(
         self,
@@ -44,7 +84,7 @@ class SwitchbotAirPurifier(SwitchbotSequenceDevice, SwitchbotEncryptedDevice):
         key_id: str,
         encryption_key: str,
         interface: int = 0,
-        model: SwitchbotModel = SwitchbotModel.AIR_PURIFIER,
+        model: SwitchbotModel = SwitchbotModel.AIR_PURIFIER_US,
         **kwargs: Any,
     ) -> None:
         super().__init__(device, key_id, encryption_key, model, interface, **kwargs)
@@ -55,20 +95,48 @@ class SwitchbotAirPurifier(SwitchbotSequenceDevice, SwitchbotEncryptedDevice):
         device: BLEDevice,
         key_id: str,
         encryption_key: str,
-        model: SwitchbotModel = SwitchbotModel.AIR_PURIFIER,
+        model: SwitchbotModel = SwitchbotModel.AIR_PURIFIER_US,
         **kwargs: Any,
     ) -> bool:
         return await super().verify_encryption_key(
             device, key_id, encryption_key, model, **kwargs
         )
 
+    @property
+    def color_modes(self) -> set[ColorMode]:
+        """Return the supported color modes."""
+        return {ColorMode.RGB}
+
+    @property
+    def color_mode(self) -> ColorMode:
+        """Return the current color mode."""
+        return ColorMode.RGB
+
     async def get_basic_info(self) -> dict[str, Any] | None:
         """Get device basic settings."""
-        if not (_data := await self._get_basic_info()):
+        if not (
+            res := await self._get_basic_info_by_multi_commands(
+                self._get_basic_info_command
+            )
+        ):
             return None
 
-        _LOGGER.debug("data: %s", _data)
+        _data, led_settings, led_status = res[0], res[1], res[2]
+
+        _LOGGER.debug(
+            "%s %s basic info %s", self._model, self._device.address, _data.hex()
+        )
+        _LOGGER.debug(
+            "%s %s led settings %s",
+            self._model,
+            self._device.address,
+            led_settings.hex(),
+        )
+        _LOGGER.debug(
+            "%s %s led_status %s", self._model, self._device.address, led_status.hex()
+        )
         isOn = bool(_data[2] & 0b10000000)
+        wireless_charging = bool(_data[2] & 0b01000000)
         version_info = (_data[2] & 0b00110000) >> 4
         _mode = _data[2] & 0b00000111
         isAqiValid = bool(_data[3] & 0b00000100)
@@ -76,11 +144,16 @@ class SwitchbotAirPurifier(SwitchbotSequenceDevice, SwitchbotEncryptedDevice):
         _aqi_level = (_data[4] & 0b00000110) >> 1
         aqi_level = AirQualityLevel(_aqi_level).name.lower()
         speed = _data[6] & 0b01111111
-        pm25 = struct.unpack("<H", _data[12:14])[0] & 0xFFF
+        pm25 = _UNPACK_UINT16_BE(_data, 12)[0] & 0xFFF
         firmware = _data[15] / 10.0
         mode = get_air_purifier_mode(_mode, speed)
+        self._state["r"] = led_settings[2]
+        self._state["g"] = led_settings[3]
+        self._state["b"] = led_settings[4]
+        brightness = led_settings[5]
+        light_sensitive = bool(led_status[1] & 0x02)
 
-        return {
+        data = {
             "isOn": isOn,
             "version_info": version_info,
             "mode": mode,
@@ -88,21 +161,16 @@ class SwitchbotAirPurifier(SwitchbotSequenceDevice, SwitchbotEncryptedDevice):
             "child_lock": child_lock,
             "aqi_level": aqi_level,
             "speed": speed,
-            "pm25": pm25,
             "firmware": firmware,
+            "brightness": brightness,
+            "light_sensitive": light_sensitive,
         }
+        if self._model in self._WIRELESS_MODELS:
+            data["wireless_charging"] = wireless_charging
 
-    async def _get_basic_info(self) -> bytes | None:
-        """Return basic info of device."""
-        _data = await self._send_command(
-            key=DEVICE_GET_BASIC_SETTINGS_KEY, retry=self._retry_count
-        )
-
-        if _data in (b"\x07", b"\x00"):
-            _LOGGER.error("Unsuccessful, please try again")
-            return None
-
-        return _data
+        if self._model in self._PM25_MODELS:
+            return data | {"pm25": pm25}
+        return data
 
     @update_after_operation
     async def set_preset_mode(self, preset_mode: str) -> bool:
@@ -110,9 +178,105 @@ class SwitchbotAirPurifier(SwitchbotSequenceDevice, SwitchbotEncryptedDevice):
         result = await self._send_command(COMMAND_SET_MODE[preset_mode])
         return self._check_command_result(result, 0, {1})
 
-    def get_current_percentage(self) -> Any:
-        """Return cached percentage."""
-        return self._get_adv_value("speed")
+    @update_after_operation
+    async def set_percentage(self, percentage: int) -> bool:
+        """Set percentage."""
+        if not 0 <= percentage <= 100:
+            raise ValueError("Percentage must be between 0 and 100")
+        self._validate_current_mode()
+
+        result = await self._send_command(
+            COMMAND_SET_PERCENTAGE.format(percentage=percentage)
+        )
+        return self._check_command_result(result, 0, {1})
+
+    def _validate_current_mode(self) -> None:
+        """Validate current mode for setting percentage."""
+        current_mode = self.get_current_mode()
+        if current_mode not in self._LEVEL_MODES:
+            raise ValueError("Percentage can only be set in LEVEL modes.")
+
+    @update_after_operation
+    async def set_brightness(self, brightness: int) -> bool:
+        """Set brightness."""
+        self._validate_brightness(brightness)
+        r, g, b = (
+            self._state.get("r", 0),
+            self._state.get("g", 0),
+            self._state.get("b", 0),
+        )
+        hex_data = f"{r:02X}{g:02X}{b:02X}{brightness:02X}"
+        result = await self._send_command(self._set_brightness_command.format(hex_data))
+        return self._check_command_result(result, 0, {1})
+
+    @update_after_operation
+    async def set_rgb(self, brightness: int, r: int, g: int, b: int) -> bool:
+        """
+        Set rgb.
+
+        Note: byte order is reversed from base class (RGB+brightness
+        instead of brightness+RGB).
+        """
+        self._validate_brightness(brightness)
+        self._validate_rgb(r, g, b)
+        hex_data = f"{r:02X}{g:02X}{b:02X}{brightness:02X}"
+        result = await self._send_command(self._set_rgb_command.format(hex_data))
+        return self._check_command_result(result, 0, {1})
+
+    @update_after_operation
+    async def turn_led_on(self) -> bool:
+        """Turn on LED."""
+        result = await self._send_command(self._turn_led_on_command)
+        return self._check_command_result(result, 0, {1})
+
+    @update_after_operation
+    async def turn_led_off(self) -> bool:
+        """Turn off LED."""
+        result = await self._send_command(self._turn_led_off_command)
+        return self._check_command_result(result, 0, {1})
+
+    @update_after_operation
+    async def open_light_sensitive_switch(self) -> bool:
+        """
+        Open the light sensitive switch.
+
+        This will allow the LED to automatically adjust brightness based on ambient light.
+        The LED will turn on in dark environments and turn off in bright environments.
+        """
+        result = await self._send_command(self._open_light_sensitive_switch_command)
+        return self._check_command_result(result, 0, {1})
+
+    @update_after_operation
+    async def close_light_sensitive_switch(self) -> bool:
+        """
+        Close the light sensitive switch.
+
+        Since the current protocol does not support obtaining the LED status,
+        sending an on or off command will turn off the light sensitive switch.
+        """
+        result = await self._send_command(self._turn_led_on_command)
+        return self._check_command_result(result, 0, {1})
+
+    def _check_wireless_charging_supported(self) -> None:
+        if self._model not in self._WIRELESS_MODELS:
+            raise SwitchbotOperationError(
+                "Wireless charging is only available on table versions"
+                f" (current model={self._model})"
+            )
+
+    @update_after_operation
+    async def open_wireless_charging(self) -> bool:
+        """Enable the wireless charging pad (table models only)."""
+        self._check_wireless_charging_supported()
+        result = await self._send_command(self._open_wireless_charging_command)
+        return self._check_command_result(result, 0, {1})
+
+    @update_after_operation
+    async def close_wireless_charging(self) -> bool:
+        """Disable the wireless charging pad (table models only)."""
+        self._check_wireless_charging_supported()
+        result = await self._send_command(self._close_wireless_charging_command)
+        return self._check_command_result(result, 0, {1})
 
     def is_on(self) -> bool | None:
         """Return air purifier state from cache."""
@@ -129,3 +293,19 @@ class SwitchbotAirPurifier(SwitchbotSequenceDevice, SwitchbotEncryptedDevice):
     def get_current_mode(self) -> Any:
         """Return cached mode."""
         return self._get_adv_value("mode")
+
+    def is_child_lock_on(self) -> bool | None:
+        """Return child lock state from cache."""
+        return self._get_adv_value("child_lock")
+
+    def is_wireless_charging_on(self) -> bool | None:
+        """Return wireless charging state from cache."""
+        return self._get_adv_value("wireless_charging")
+
+    def get_current_percentage(self) -> int | None:
+        """Return cached percentage."""
+        return self._get_adv_value("speed")
+
+    def is_light_sensitive_on(self) -> bool | None:
+        """Return light sensitive state from cache."""
+        return self._get_adv_value("light_sensitive")

switchbot/devices/base_light.py

@@ -75,10 +75,29 @@ class SwitchbotBaseLight(SwitchbotDevice):
         """Return the current effect."""
         return self._get_adv_value("effect")
 
+    @staticmethod
+    def _validate_brightness(brightness: int) -> None:
+        if not 0 <= brightness <= 100:
+            raise ValueError("Brightness must be between 0 and 100")
+
+    @staticmethod
+    def _validate_rgb(r: int, g: int, b: int) -> None:
+        if not 0 <= r <= 255:
+            raise ValueError("r must be between 0 and 255")
+        if not 0 <= g <= 255:
+            raise ValueError("g must be between 0 and 255")
+        if not 0 <= b <= 255:
+            raise ValueError("b must be between 0 and 255")
+
+    @staticmethod
+    def _validate_color_temp(color_temp: int) -> None:
+        if not 2700 <= color_temp <= 6500:
+            raise ValueError("Color Temp must be between 2700 and 6500")
+
     @update_after_operation
     async def set_brightness(self, brightness: int) -> bool:
         """Set brightness."""
-        assert 0 <= brightness <= 100, "Brightness must be between 0 and 100"
+        self._validate_brightness(brightness)
         hex_brightness = f"{brightness:02X}"
         self._check_function_support(self._set_brightness_command)
         result = await self._send_command(
@@ -89,8 +108,8 @@ class SwitchbotBaseLight(SwitchbotDevice):
     @update_after_operation
     async def set_color_temp(self, brightness: int, color_temp: int) -> bool:
         """Set color temp."""
-        assert 0 <= brightness <= 100, "Brightness must be between 0 and 100"
-        assert 2700 <= color_temp <= 6500, "Color Temp must be between 2700 and 6500"
+        self._validate_brightness(brightness)
+        self._validate_color_temp(color_temp)
         hex_data = f"{brightness:02X}{color_temp:04X}"
         self._check_function_support(self._set_color_temp_command)
         result = await self._send_command(self._set_color_temp_command.format(hex_data))
@@ -99,10 +118,8 @@ class SwitchbotBaseLight(SwitchbotDevice):
     @update_after_operation
     async def set_rgb(self, brightness: int, r: int, g: int, b: int) -> bool:
         """Set rgb."""
-        assert 0 <= brightness <= 100, "Brightness must be between 0 and 100"
-        assert 0 <= r <= 255, "r must be between 0 and 255"
-        assert 0 <= g <= 255, "g must be between 0 and 255"
-        assert 0 <= b <= 255, "b must be between 0 and 255"
+        self._validate_brightness(brightness)
+        self._validate_rgb(r, g, b)
         self._check_function_support(self._set_rgb_command)
         hex_data = f"{brightness:02X}{r:02X}{g:02X}{b:02X}"
         result = await self._send_command(self._set_rgb_command.format(hex_data))
@@ -141,7 +158,8 @@ class SwitchbotBaseLight(SwitchbotDevice):
         """Get device basic settings by sending multiple commands."""
         results = []
         for command in commands:
-            if not (result := await self._get_basic_info(command)):
+            result = await self._send_command(command)
+            if not self._check_command_result(result, 0, {1}):
                 return None
             results.append(result)
         return results

switchbot/devices/device.py

@@ -78,10 +78,10 @@ API_MODEL_TO_ENUM: dict[str, SwitchbotModel] = {
     "WoFan2": SwitchbotModel.CIRCULATOR_FAN,
     "WoHub2": SwitchbotModel.HUB2,
     "WoRollerShade": SwitchbotModel.ROLLER_SHADE,
-    "WoAirPurifierJP": SwitchbotModel.AIR_PURIFIER,
-    "WoAirPurifierUS": SwitchbotModel.AIR_PURIFIER,
-    "WoAirPurifierJPPro": SwitchbotModel.AIR_PURIFIER_TABLE,
-    "WoAirPurifierUSPro": SwitchbotModel.AIR_PURIFIER_TABLE,
+    "WoAirPurifierJP": SwitchbotModel.AIR_PURIFIER_JP,
+    "WoAirPurifierUS": SwitchbotModel.AIR_PURIFIER_US,
+    "WoAirPurifierJPPro": SwitchbotModel.AIR_PURIFIER_TABLE_JP,
+    "WoAirPurifierUSPro": SwitchbotModel.AIR_PURIFIER_TABLE_US,
     "WoSweeperMini": SwitchbotModel.K10_VACUUM,
     "WoSweeperMiniPro": SwitchbotModel.K10_PRO_VACUUM,
     "91AgWZ1n": SwitchbotModel.K10_PRO_COMBO_VACUUM,
@@ -216,6 +216,8 @@ class SwitchbotBaseDevice:
     _open_command: str | None = None
     _close_command: str | None = None
     _press_command: str | None = None
+    _open_child_lock_command: str | None = None
+    _close_child_lock_command: str | None = None
 
     def __init__(
         self,
@@ -935,6 +937,20 @@ class SwitchbotBaseDevice:
         result = await self._send_command(self._press_command)
         return self._check_command_result(result, 0, {1})
 
+    @update_after_operation
+    async def open_child_lock(self) -> bool:
+        """Open the child lock."""
+        self._check_function_support(self._open_child_lock_command)
+        result = await self._send_command(self._open_child_lock_command)
+        return self._check_command_result(result, 0, {1})
+
+    @update_after_operation
+    async def close_child_lock(self) -> bool:
+        """Close the child lock."""
+        self._check_function_support(self._close_child_lock_command)
+        result = await self._send_command(self._close_child_lock_command)
+        return self._check_command_result(result, 0, {1})
+
 
 class SwitchbotDevice(SwitchbotBaseDevice):
     """

tests/test_adv_parser.py

@@ -2364,8 +2364,8 @@ def test_s10_with_empty_data() -> None:
                 "sequence_number": 161,
             },
             "7",
-            "Air Purifier Table",
-            SwitchbotModel.AIR_PURIFIER_TABLE,
+            "Air Purifier Table US",
+            SwitchbotModel.AIR_PURIFIER_TABLE_US,
         ),
         AdvTestCase(
             b'\xcc\x8d\xa2\xa7\x92>\t"\x80\x000\x00\x0f\x00\x00',
@@ -2382,8 +2382,8 @@ def test_s10_with_empty_data() -> None:
                 "sequence_number": 9,
             },
             "*",
-            "Air Purifier",
-            SwitchbotModel.AIR_PURIFIER,
+            "Air Purifier US",
+            SwitchbotModel.AIR_PURIFIER_US,
         ),
         AdvTestCase(
             b"\xcc\x8d\xa2\xa7\xe4\xa6\x0b\x83\x88d\x00\xea`\x00\x00",
@@ -2400,8 +2400,8 @@ def test_s10_with_empty_data() -> None:
                 "sequence_number": 11,
             },
             "+",
-            "Air Purifier",
-            SwitchbotModel.AIR_PURIFIER,
+            "Air Purifier JP",
+            SwitchbotModel.AIR_PURIFIER_JP,
         ),
         AdvTestCase(
             b"\xcc\x8d\xa2\xa7\xc1\xae\x9b\x81\x8c\xb2\x00\x01\x94\x00\x00",
@@ -2418,8 +2418,8 @@ def test_s10_with_empty_data() -> None:
                 "sequence_number": 155,
             },
             "8",
-            "Air Purifier Table",
-            SwitchbotModel.AIR_PURIFIER_TABLE,
+            "Air Purifier Table JP",
+            SwitchbotModel.AIR_PURIFIER_TABLE_JP,
         ),
         AdvTestCase(
             b"\xcc\x8d\xa2\xa7\xc1\xae\x9e\xa1\x8c\x800\x01\x95\x00\x00",
@@ -2436,8 +2436,8 @@ def test_s10_with_empty_data() -> None:
                 "sequence_number": 158,
             },
             "8",
-            "Air Purifier Table",
-            SwitchbotModel.AIR_PURIFIER_TABLE,
+            "Air Purifier Table JP",
+            SwitchbotModel.AIR_PURIFIER_TABLE_JP,
         ),
         AdvTestCase(
             b"\xcc\x8d\xa2\xa7\xc1\xae\x9e\x05\x8c\x800\x01\x95\x00\x00",
@@ -2454,8 +2454,8 @@ def test_s10_with_empty_data() -> None:
                 "sequence_number": 158,
             },
             "8",
-            "Air Purifier Table",
-            SwitchbotModel.AIR_PURIFIER_TABLE,
+            "Air Purifier Table JP",
+            SwitchbotModel.AIR_PURIFIER_TABLE_JP,
         ),
     ],
 )
@@ -2491,7 +2491,9 @@ def test_air_purifier_passive() -> None:
         },
         rssi=-97,
     )
-    result = parse_advertisement_data(ble_device, adv_data, SwitchbotModel.AIR_PURIFIER)
+    result = parse_advertisement_data(
+        ble_device, adv_data, SwitchbotModel.AIR_PURIFIER_US
+    )
     assert result == SwitchBotAdvertisement(
         address="aa:bb:cc:dd:ee:ff",
         data={
@@ -2509,8 +2511,8 @@ def test_air_purifier_passive() -> None:
             },
             "isEncrypted": False,
             "model": "*",
-            "modelFriendlyName": "Air Purifier",
-            "modelName": SwitchbotModel.AIR_PURIFIER,
+            "modelFriendlyName": "Air Purifier US",
+            "modelName": SwitchbotModel.AIR_PURIFIER_US,
         },
         device=ble_device,
         rssi=-97,

tests/test_air_purifier.py

@@ -3,38 +3,53 @@ from unittest.mock import AsyncMock, MagicMock, patch
 import pytest
 from bleak.backends.device import BLEDevice
 
-from switchbot import SwitchBotAdvertisement, SwitchbotEncryptedDevice, SwitchbotModel
+from switchbot import (
+    SwitchBotAdvertisement,
+    SwitchbotEncryptedDevice,
+    SwitchbotModel,
+    SwitchbotOperationError,
+)
 from switchbot.const.air_purifier import AirPurifierMode
 from switchbot.devices import air_purifier
 
 from .test_adv_parser import generate_ble_device
 
 common_params = [
-    (b"7\x00\x00\x95-\x00", "7"),
-    (b"*\x00\x00\x15\x04\x00", "*"),
-    (b"+\x00\x00\x15\x04\x00", "+"),
-    (b"8\x00\x00\x95-\x00", "8"),
+    (b"7\x00\x00\x95-\x00", "7", SwitchbotModel.AIR_PURIFIER_TABLE_US),
+    (b"*\x00\x00\x15\x04\x00", "*", SwitchbotModel.AIR_PURIFIER_US),
+    (b"+\x00\x00\x15\x04\x00", "+", SwitchbotModel.AIR_PURIFIER_JP),
+    (b"8\x00\x00\x95-\x00", "8", SwitchbotModel.AIR_PURIFIER_TABLE_JP),
 ]
 
 
 def create_device_for_command_testing(
-    rawAdvData: bytes, model: str, init_data: dict | None = None
+    rawAdvData: bytes,
+    model: str,
+    model_name: SwitchbotModel,
+    init_data: dict | None = None,
 ):
     ble_device = generate_ble_device("aa:bb:cc:dd:ee:ff", "any")
     device = air_purifier.SwitchbotAirPurifier(
-        ble_device, "ff", "ffffffffffffffffffffffffffffffff"
+        ble_device,
+        "ff",
+        "ffffffffffffffffffffffffffffffff",
+        model=model_name,
     )
+    device.update = AsyncMock()
     device.update_from_advertisement(
-        make_advertisement_data(ble_device, rawAdvData, model, init_data)
+        make_advertisement_data(ble_device, rawAdvData, model, model_name, init_data)
     )
     device._send_command = AsyncMock()
     device._check_command_result = MagicMock()
-    device.update = AsyncMock()
     return device
 
 
 def make_advertisement_data(
-    ble_device: BLEDevice, rawAdvData: bytes, model: str, init_data: dict | None = None
+    ble_device: BLEDevice,
+    rawAdvData: bytes,
+    model: str,
+    model_name: SwitchbotModel,
+    init_data: dict | None = None,
 ):
     """Set advertisement data with defaults."""
     if init_data is None:
@@ -59,7 +74,7 @@ def make_advertisement_data(
             "isEncrypted": False,
             "model": model,
             "modelFriendlyName": "Air Purifier",
-            "modelName": SwitchbotModel.AIR_PURIFIER,
+            "modelName": model_name,
         },
         device=ble_device,
         rssi=-80,
@@ -69,15 +84,17 @@ def make_advertisement_data(
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
-    ("rawAdvData", "model"),
+    ("rawAdvData", "model", "model_name"),
     common_params,
 )
 @pytest.mark.parametrize(
     "pm25",
     [150],
 )
-async def test_status_from_proceess_adv(rawAdvData, model, pm25):
-    device = create_device_for_command_testing(rawAdvData, model, {"pm25": pm25})
+async def test_status_from_process_adv(rawAdvData, model, model_name, pm25):
+    device = create_device_for_command_testing(
+        rawAdvData, model, model_name, {"pm25": pm25}
+    )
     assert device.get_current_percentage() == 100
     assert device.is_on() is True
     assert device.get_current_aqi_level() == "excellent"
@@ -87,55 +104,61 @@ async def test_status_from_proceess_adv(rawAdvData, model, pm25):
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
-    ("rawAdvData", "model"),
+    ("rawAdvData", "model", "model_name"),
     common_params,
 )
-async def test_get_basic_info_returns_none_when_no_data(rawAdvData, model):
-    device = create_device_for_command_testing(rawAdvData, model)
-    device._get_basic_info = AsyncMock(return_value=None)
+async def test_get_basic_info_returns_none_when_no_data(rawAdvData, model, model_name):
+    device = create_device_for_command_testing(rawAdvData, model, model_name)
+    device._get_basic_info_by_multi_commands = AsyncMock(return_value=None)
 
     assert await device.get_basic_info() is None
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
-    ("rawAdvData", "model"),
+    ("rawAdvData", "model", "model_name"),
     common_params,
 )
 @pytest.mark.parametrize(
     "mode", ["level_1", "level_2", "level_3", "auto", "pet", "sleep"]
 )
-async def test_set_preset_mode(rawAdvData, model, mode):
-    device = create_device_for_command_testing(rawAdvData, model, {"mode": mode})
+async def test_set_preset_mode(rawAdvData, model, model_name, mode):
+    device = create_device_for_command_testing(
+        rawAdvData, model, model_name, {"mode": mode}
+    )
     await device.set_preset_mode(mode)
     assert device.get_current_mode() == mode
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
-    ("rawAdvData", "model"),
+    ("rawAdvData", "model", "model_name"),
     common_params,
 )
-async def test_turn_on(rawAdvData, model):
-    device = create_device_for_command_testing(rawAdvData, model, {"isOn": True})
+async def test_turn_on(rawAdvData, model, model_name):
+    device = create_device_for_command_testing(
+        rawAdvData, model, model_name, {"isOn": True}
+    )
     await device.turn_on()
     assert device.is_on() is True
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
-    ("rawAdvData", "model"),
+    ("rawAdvData", "model", "model_name"),
     common_params,
 )
-async def test_turn_off(rawAdvData, model):
-    device = create_device_for_command_testing(rawAdvData, model, {"isOn": False})
+async def test_turn_off(rawAdvData, model, model_name):
+    device = create_device_for_command_testing(
+        rawAdvData, model, model_name, {"isOn": False}
+    )
     await device.turn_off()
     assert device.is_on() is False
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
-    ("rawAdvData", "model"),
+    ("rawAdvData", "model", "model_name"),
     common_params,
 )
 @pytest.mark.parametrize(
@@ -146,8 +169,8 @@ async def test_turn_off(rawAdvData, model):
         (b"\x01\x02\x03", b"\x01\x02\x03"),
     ],
 )
-async def test__get_basic_info(rawAdvData, model, response, expected):
-    device = create_device_for_command_testing(rawAdvData, model)
+async def test__get_basic_info(rawAdvData, model, model_name, response, expected):
+    device = create_device_for_command_testing(rawAdvData, model, model_name)
     device._send_command = AsyncMock(return_value=response)
     result = await device._get_basic_info()
     assert result == expected
@@ -155,33 +178,61 @@ async def test__get_basic_info(rawAdvData, model, response, expected):
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
-    ("rawAdvData", "model"),
+    "device_case",
     common_params,
 )
 @pytest.mark.parametrize(
-    ("basic_info", "result"),
+    "info_case",
     [
         (
             bytearray(
-                b"\x01\xa7\xe9\x8c\x08\x00\xb2\x01\x96\x00\x00\x00\xf0\x00\x00\x17"
+                b"\x01\xa7\xe9\x8c\x08\x00\xb2\x01\x96\x00\x00\x00\x00\x01\x00\x17"
             ),
-            [True, 2, "level_2", True, False, "excellent", 50, 240, 2.3],
+            bytearray(b"\x01\x01\x11\x22\x33\x44"),
+            bytearray(b"\x01\x03"),
+            [
+                True,
+                2,
+                "level_2",
+                True,
+                False,
+                "excellent",
+                50,
+                1,
+                2.3,
+                0x44,
+                True,
+            ],
         ),
         (
             bytearray(
-                b"\x01\xa8\xec\x8c\x08\x00\xb2\x01\x96\x00\x00\x00\xf0\x00\x00\x17"
+                b"\x01\xa8\xec\x8c\x08\x00\xb2\x01\x96\x00\x00\x00\x01\x00\x00\x17"
             ),
-            [True, 2, "pet", True, False, "excellent", 50, 240, 2.3],
+            bytearray(b"\x01\x01\xaa\xbb\xcc\x1e"),
+            bytearray(b"\x01\x00"),
+            [
+                True,
+                2,
+                "pet",
+                True,
+                False,
+                "excellent",
+                50,
+                256,
+                2.3,
+                0x1E,
+                False,
+            ],
         ),
     ],
 )
-async def test_get_basic_info(rawAdvData, model, basic_info, result):
-    device = create_device_for_command_testing(rawAdvData, model)
-
-    async def mock_get_basic_info():
-        return basic_info
-
-    device._get_basic_info = AsyncMock(side_effect=mock_get_basic_info)
+async def test_get_basic_info(device_case, info_case):
+    rawAdvData, model, model_name = device_case
+    basic_info, led_settings, led_status, result = info_case
+    device = create_device_for_command_testing(rawAdvData, model, model_name)
+    device._get_basic_info_by_multi_commands = AsyncMock(
+        return_value=[basic_info, led_settings, led_status]
+    )
 
     info = await device.get_basic_info()
     assert info["isOn"] == result[0]
@@ -191,8 +242,16 @@ async def test_get_basic_info(rawAdvData, model, basic_info, result):
     assert info["child_lock"] == result[4]
     assert info["aqi_level"] == result[5]
     assert info["speed"] == result[6]
-    assert info["pm25"] == result[7]
+    if model_name not in (
+        SwitchbotModel.AIR_PURIFIER_JP,
+        SwitchbotModel.AIR_PURIFIER_TABLE_JP,
+    ):
+        assert info["pm25"] == result[7]
+    else:
+        assert "pm25" not in info
     assert info["firmware"] == result[8]
+    assert info["brightness"] == result[9]
+    assert info["light_sensitive"] == result[10]
 
 
 @pytest.mark.asyncio
@@ -214,7 +273,7 @@ async def test_verify_encryption_key(mock_parent_verify):
         ble_device,
         key_id,
         encryption_key,
-        SwitchbotModel.AIR_PURIFIER,
+        SwitchbotModel.AIR_PURIFIER_US,
     )
 
     assert result is True
@@ -229,3 +288,197 @@ def test_get_modes():
         "sleep",
         "pet",
     ]
+
+
+@pytest.mark.asyncio
+async def test_air_purifier_color_and_led_properties():
+    raw_adv, model, model_name = common_params[0]
+    device = create_device_for_command_testing(
+        raw_adv,
+        model,
+        model_name,
+    )
+
+    assert device.color_modes == {air_purifier.ColorMode.RGB}
+    assert device.color_mode == air_purifier.ColorMode.RGB
+
+
+@pytest.mark.asyncio
+async def test_set_percentage_validation_and_command():
+    raw_adv, model, model_name = common_params[0]
+    device = create_device_for_command_testing(
+        raw_adv,
+        model,
+        model_name,
+        {"mode": "level_2"},
+    )
+    device._check_command_result = MagicMock(return_value=True)
+    device._send_command = AsyncMock(return_value=b"\x01")
+
+    assert await device.set_percentage(25) is True
+    device._send_command.assert_called_with(
+        air_purifier.COMMAND_SET_PERCENTAGE.format(percentage=25)
+    )
+
+    with pytest.raises(ValueError, match="Percentage must be between 0 and 100"):
+        await device.set_percentage(-1)
+    with pytest.raises(ValueError, match="Percentage must be between 0 and 100"):
+        await device.set_percentage(101)
+
+    invalid_mode_device = create_device_for_command_testing(
+        raw_adv,
+        model,
+        model_name,
+        {"mode": "auto"},
+    )
+    with pytest.raises(ValueError, match="Percentage can only be set in LEVEL modes"):
+        await invalid_mode_device.set_percentage(10)
+
+
+@pytest.mark.asyncio
+async def test_set_brightness_validation_and_command():
+    raw_adv, model, model_name = common_params[0]
+    device = create_device_for_command_testing(raw_adv, model, model_name)
+    device._state = {"r": 1, "g": 2, "b": 3}
+    device._check_command_result = MagicMock(return_value=True)
+    device._send_command = AsyncMock(return_value=b"\x01")
+
+    assert await device.set_brightness(10) is True
+    device._send_command.assert_called_with(
+        device._set_brightness_command.format("0102030A")
+    )
+
+    with pytest.raises(ValueError, match="Brightness must be between 0 and 100"):
+        await device.set_brightness(101)
+
+
+@pytest.mark.asyncio
+async def test_set_rgb_validation_and_command():
+    raw_adv, model, model_name = common_params[0]
+    device = create_device_for_command_testing(raw_adv, model, model_name)
+    device._check_command_result = MagicMock(return_value=True)
+    device._send_command = AsyncMock(return_value=b"\x01")
+
+    assert await device.set_rgb(20, 1, 2, 3) is True
+    device._send_command.assert_called_with(device._set_rgb_command.format("01020314"))
+
+    with pytest.raises(ValueError, match="Brightness must be between 0 and 100"):
+        await device.set_rgb(101, 1, 2, 3)
+    with pytest.raises(ValueError, match="r must be between 0 and 255"):
+        await device.set_rgb(10, 256, 2, 3)
+    with pytest.raises(ValueError, match="g must be between 0 and 255"):
+        await device.set_rgb(10, 1, 256, 3)
+    with pytest.raises(ValueError, match="b must be between 0 and 255"):
+        await device.set_rgb(10, 1, 2, 256)
+
+
+@pytest.mark.asyncio
+async def test_led_and_light_sensitive_commands():
+    raw_adv, model, model_name = common_params[0]
+    device = create_device_for_command_testing(
+        raw_adv, model, model_name, {"led_status": True}
+    )
+    device._check_command_result = MagicMock(return_value=True)
+    device._send_command = AsyncMock(return_value=b"\x01")
+
+    assert await device.turn_led_on() is True
+    device._send_command.assert_called_with(device._turn_led_on_command)
+
+    assert await device.turn_led_off() is True
+    device._send_command.assert_called_with(device._turn_led_off_command)
+
+    assert await device.open_light_sensitive_switch() is True
+    device._send_command.assert_called_with(device._open_light_sensitive_switch_command)
+
+    assert await device.close_light_sensitive_switch() is True
+    device._send_command.assert_called_with(device._turn_led_on_command)
+
+    device_off = create_device_for_command_testing(
+        raw_adv,
+        model,
+        model_name,
+    )
+    device_off._check_command_result = MagicMock(return_value=True)
+    device_off._send_command = AsyncMock(return_value=b"\x01")
+    assert await device_off.close_light_sensitive_switch() is True
+    device_off._send_command.assert_called_with(device_off._turn_led_on_command)
+
+
+@pytest.mark.asyncio
+async def test_air_purifier_cache_getters():
+    raw_adv, model, model_name = common_params[0]
+    device = create_device_for_command_testing(
+        raw_adv,
+        model,
+        model_name,
+        {
+            "child_lock": True,
+            "wireless_charging": True,
+            "light_sensitive": True,
+            "speed": 88,
+        },
+    )
+
+    assert device.get_current_percentage() == 88
+    assert device.is_child_lock_on() is True
+    assert device.is_wireless_charging_on() is True
+    assert device.is_light_sensitive_on() is True
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "operation_case",
+    [
+        ("open_child_lock", "_open_child_lock_command"),
+        ("close_child_lock", "_close_child_lock_command"),
+    ],
+)
+async def test_child_lock_operations(operation_case):
+    """Child lock commands should always be forwarded correctly."""
+    raw_adv, model, model_name = common_params[0]
+    device = create_device_for_command_testing(raw_adv, model, model_name)
+    operation_name, command_attr = operation_case
+    command = getattr(device, command_attr)
+
+    device._check_function_support = MagicMock()
+    device._check_command_result = MagicMock(return_value=True)
+    device._send_command = AsyncMock(return_value=b"\x01")
+
+    operation = getattr(device, operation_name)
+    assert await operation() is True
+
+    device._check_function_support.assert_called_with(command)
+    device._send_command.assert_called_with(command)
+    device._check_command_result.assert_called_with(b"\x01", 0, {1})
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("raw_adv", "model", "model_name", "supported"),
+    [
+        (*common_params[0], True),
+        (*common_params[3], True),
+        (*common_params[1], False),
+        (*common_params[2], False),
+    ],
+)
+async def test_wireless_charging_model_support(raw_adv, model, model_name, supported):
+    """Wireless charging operations should only succeed for table variants."""
+    device = create_device_for_command_testing(raw_adv, model, model_name)
+    if supported:
+        device._check_command_result = MagicMock(return_value=True)
+        device._send_command = AsyncMock(return_value=b"\x01")
+
+        assert await device.open_wireless_charging() is True
+        assert await device.close_wireless_charging() is True
+
+        assert device._send_command.call_args_list == [
+            ((device._open_wireless_charging_command,),),
+            ((device._close_wireless_charging_command,),),
+        ]
+
+    else:
+        with pytest.raises(SwitchbotOperationError):
+            await device.open_wireless_charging()
+        with pytest.raises(SwitchbotOperationError):
+            await device.close_wireless_charging()

tests/test_bulb.py

@@ -81,15 +81,10 @@ async def test_default_info():
 )
 async def test_get_basic_info_returns_none(basic_info, version_info):
     device = create_device_for_command_testing()
-
-    async def mock_get_basic_info(arg):
-        if arg == device._get_basic_info_command[1]:
-            return basic_info
-        if arg == device._get_basic_info_command[0]:
-            return version_info
-        return None
-
-    device._get_basic_info = AsyncMock(side_effect=mock_get_basic_info)
+    device._send_command = AsyncMock(side_effect=[version_info, basic_info])
+    device._check_command_result = MagicMock(
+        side_effect=[bool(version_info), bool(basic_info)]
+    )
 
     assert await device.get_basic_info() is None
 
@@ -124,15 +119,10 @@ async def test_get_basic_info_returns_none(basic_info, version_info):
 async def test_get_basic_info(info_data, result):
     """Test getting basic info from the color bulb."""
     device = create_device_for_command_testing()
-
-    async def mock_get_basic_info(args: str) -> list[int] | None:
-        if args == device._get_basic_info_command[1]:
-            return info_data["basic_info"]
-        if args == device._get_basic_info_command[0]:
-            return info_data["version_info"]
-        return None
-
-    device._get_basic_info = AsyncMock(side_effect=mock_get_basic_info)
+    device._send_command = AsyncMock(
+        side_effect=[info_data["version_info"], info_data["basic_info"]]
+    )
+    device._check_command_result = MagicMock(side_effect=[True, True])
     info = await device.get_basic_info()
 
     assert info["isOn"] is result[0]

tests/test_ceiling_light.py

@@ -76,15 +76,10 @@ async def test_default_info():
 )
 async def test_get_basic_info_returns_none(basic_info, version_info):
     device = create_device_for_command_testing()
-
-    async def mock_get_basic_info(arg):
-        if arg == device._get_basic_info_command[1]:
-            return basic_info
-        if arg == device._get_basic_info_command[0]:
-            return version_info
-        return None
-
-    device._get_basic_info = AsyncMock(side_effect=mock_get_basic_info)
+    device._send_command = AsyncMock(side_effect=[version_info, basic_info])
+    device._check_command_result = MagicMock(
+        side_effect=[bool(version_info), bool(basic_info)]
+    )
 
     assert await device.get_basic_info() is None
 
@@ -119,15 +114,10 @@ async def test_get_basic_info_returns_none(basic_info, version_info):
 async def test_get_basic_info(info_data, result):
     """Test getting basic info from the ceiling light."""
     device = create_device_for_command_testing()
-
-    async def mock_get_basic_info(args: str) -> list[int] | None:
-        if args == device._get_basic_info_command[1]:
-            return info_data["basic_info"]
-        if args == device._get_basic_info_command[0]:
-            return info_data["version_info"]
-        return None
-
-    device._get_basic_info = AsyncMock(side_effect=mock_get_basic_info)
+    device._send_command = AsyncMock(
+        side_effect=[info_data["version_info"], info_data["basic_info"]]
+    )
+    device._check_command_result = MagicMock(side_effect=[True, True])
     info = await device.get_basic_info()
 
     assert info["isOn"] is result[0]

tests/test_lock.py

@@ -38,7 +38,7 @@ def test_lock_init(model: str):
 @pytest.mark.parametrize(
     "model",
     [
-        SwitchbotModel.AIR_PURIFIER,
+        SwitchbotModel.AIR_PURIFIER_JP,
     ],
 )
 def test_lock_init_with_invalid_model(model: str):

tests/test_strip_light.py

@@ -123,14 +123,10 @@ async def test_get_basic_info_returns_none(basic_info, version_info, device_case
     adv_info, dev_cls = device_case
     device = create_device_for_command_testing(adv_info, dev_cls)
 
-    async def mock_get_basic_info(arg):
-        if arg == device._get_basic_info_command[1]:
-            return basic_info
-        if arg == device._get_basic_info_command[0]:
-            return version_info
-        return None
-
-    device._get_basic_info = AsyncMock(side_effect=mock_get_basic_info)
+    device._send_command = AsyncMock(side_effect=[version_info, basic_info])
+    device._check_command_result = MagicMock(
+        side_effect=[bool(version_info), bool(basic_info)]
+    )
 
     assert await device.get_basic_info() is None
 
@@ -167,14 +163,10 @@ async def test_strip_light_get_basic_info(info_data, result, device_case):
     adv_info, dev_cls = device_case
     device = create_device_for_command_testing(adv_info, dev_cls)
 
-    async def mock_get_basic_info(args: str) -> list[int] | None:
-        if args == device._get_basic_info_command[1]:
-            return info_data["basic_info"]
-        if args == device._get_basic_info_command[0]:
-            return info_data["version_info"]
-        return None
-
-    device._get_basic_info = AsyncMock(side_effect=mock_get_basic_info)
+    device._send_command = AsyncMock(
+        side_effect=[info_data["version_info"], info_data["basic_info"]]
+    )
+    device._check_command_result = MagicMock(side_effect=[True, True])
     info = await device.get_basic_info()
 
     assert info["isOn"] is result[0]