fix COMMAND_DISABLE_NOTIFICATIONS command

switchbot/devices/lock.py

@@ -37,7 +37,12 @@ COMMAND_LOCK = {
     SwitchbotModel.LOCK_PRO: f"{COMMAND_HEADER}0f4e0101000000",
     SwitchbotModel.LOCK_ULTRA: f"{COMMAND_HEADER}0f4e0101000000",
 }
-COMMAND_ENABLE_NOTIFICATIONS = f"{COMMAND_HEADER}0e01001e00008101"
+COMMAND_ENABLE_NOTIFICATIONS = {
+    SwitchbotModel.LOCK: f"{COMMAND_HEADER}0e01001e00008101",
+    SwitchbotModel.LOCK_LITE: f"{COMMAND_HEADER}0e01001e00008101",
+    SwitchbotModel.LOCK_PRO: f"{COMMAND_HEADER}0e01001e00008104",
+    SwitchbotModel.LOCK_ULTRA: f"{COMMAND_HEADER}0e01001e00008107",
+}
 COMMAND_DISABLE_NOTIFICATIONS = f"{COMMAND_HEADER}0e00"
 
 MOVING_STATUSES = {LockStatus.LOCKING, LockStatus.UNLOCKING}
@@ -123,6 +128,7 @@ class SwitchbotLock(SwitchbotSequenceDevice, SwitchbotEncryptedDevice):
         if status in ignore_statuses:
             return True
 
+        await self._enable_notifications()
         result = await self._send_command(command)
         status = self._check_command_result(result, 0, COMMAND_RESULT_EXPECTED_VALUES)
 
@@ -196,12 +202,8 @@ class SwitchbotLock(SwitchbotSequenceDevice, SwitchbotEncryptedDevice):
         return _data
 
     async def _enable_notifications(self) -> bool:
-        if self._notifications_enabled:
-            return True
-        result = await self._send_command(COMMAND_ENABLE_NOTIFICATIONS)
-        if self._check_command_result(result, 0, COMMAND_RESULT_EXPECTED_VALUES):
-            self._notifications_enabled = True
-        return self._notifications_enabled
+        result = await self._send_command(COMMAND_ENABLE_NOTIFICATIONS[self._model])
+        return self._check_command_result(result, 0, COMMAND_RESULT_EXPECTED_VALUES)
 
     async def _disable_notifications(self) -> bool:
         if not self._notifications_enabled:
@@ -212,7 +214,7 @@ class SwitchbotLock(SwitchbotSequenceDevice, SwitchbotEncryptedDevice):
         return not self._notifications_enabled
 
     def _notification_handler(self, _sender: int, data: bytearray) -> None:
-        if self._notifications_enabled and self._check_command_result(data, 0, {0xF}):
+        if self._check_command_result(data, 0, {0xF}):
             if self._expected_disconnect:
                 _LOGGER.debug(
                     "%s: Ignoring lock notification during expected disconnect",

tests/test_lock.py

@@ -372,27 +372,6 @@ async def test_enable_notifications(model: str):
     with patch.object(device, "_send_command", return_value=b"\x01\x00"):
         result = await device._enable_notifications()
         assert result is True
-        assert device._notifications_enabled is True
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize(
-    "model",
-    [
-        SwitchbotModel.LOCK,
-        SwitchbotModel.LOCK_LITE,
-        SwitchbotModel.LOCK_PRO,
-        SwitchbotModel.LOCK_ULTRA,
-    ],
-)
-async def test_enable_notifications_already_enabled(model: str):
-    """Test _enable_notifications when already enabled."""
-    device = create_device_for_command_testing(model)
-    device._notifications_enabled = True
-    with patch.object(device, "_send_command") as mock_send:
-        result = await device._enable_notifications()
-        assert result is True
-        mock_send.assert_not_called()
 
 
 @pytest.mark.asyncio
@@ -467,7 +446,7 @@ def test_notification_handler_not_enabled(model: str):
     """Test _notification_handler when notifications not enabled."""
     device = create_device_for_command_testing(model)
     device._notifications_enabled = False
-    data = bytearray(b"\x0f\x00\x00\x00\x80\x00\x00\x00\x00\x00")
+    data = bytearray(b"\x01\x00\x00\x00\x80\x00\x00\x00\x00\x00")
     with (
         patch.object(device, "_update_lock_status") as mock_update,
         patch.object(