python-ember-mug 0.9.0b1__tar.gz → 0.9.0b2__tar.gz

{python_ember_mug-0.9.0b1 → python_ember_mug-0.9.0b2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: python-ember-mug
-Version: 0.9.0b1
+Version: 0.9.0b2
 Summary: Python Library for Ember Mugs.
 Project-URL: Documentation, https://sopelj.github.io/python-ember-mug/
 Project-URL: Source code, https://github.com/sopelj/python-ember-mug/

{python_ember_mug-0.9.0b1 → python_ember_mug-0.9.0b2}/ember_mug/__init__.py

@@ -5,4 +5,4 @@ __all__ = ("EmberMug",)
 
 __author__ = """Jesse Sopel"""
 __email__ = "jesse.sopel@gmail.com"
-__version__ = "0.9.0b1"
+__version__ = "0.9.0b2"

{python_ember_mug-0.9.0b1 → python_ember_mug-0.9.0b2}/ember_mug/consts.py

@@ -35,6 +35,7 @@ class DeviceModel(str, Enum):
     MUG_2_14_OZ = "CM19P"  # or CM21L?
     TRAVEL_MUG_12_OZ = "TM19"
     TUMBLER_16_OZ = "CM21XL"
+    UNKNOWN_DEVICE = "Unknown Mug or Tumbler"
 
 
 class DeviceColour(str, Enum):

{python_ember_mug-0.9.0b1 → python_ember_mug-0.9.0b2}/ember_mug/mug.py

@@ -8,7 +8,7 @@ from datetime import datetime, timezone
 from enum import Enum
 from functools import cached_property
 from time import time
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Concatenate, Literal, ParamSpec, TypeVar
 
 from bleak import BleakClient, BleakError
 from bleak_retry_connector import establish_connection
@@ -18,7 +18,6 @@ from .consts import (
     IS_LINUX,
     MUG_NAME_REGEX,
     PUSH_EVENT_BATTERY_IDS,
-    DeviceType,
     LiquidState,
     MugCharacteristic,
     PushEvent,
@@ -36,7 +35,7 @@ from .utils import (
 )
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncIterator, Callable
+    from collections.abc import AsyncIterator, Awaitable, Callable
 
     from bleak.backends.characteristic import BleakGATTCharacteristic
     from bleak.backends.device import BLEDevice
@@ -46,6 +45,31 @@ logger = logging.getLogger(__name__)
 
 DISCONNECT_DELAY = 120
 
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+def require_attribute(
+    attr_name: str,
+) -> Callable[[Callable[Concatenate[EmberMug, P], Awaitable[T]]], Callable[Concatenate[EmberMug, P], Awaitable[T]]]:
+    """Require an attribute to be available on the device."""
+
+    def decorator(
+        func: Callable[Concatenate[EmberMug, P], Awaitable[T]],
+    ) -> Callable[Concatenate[EmberMug, P], Awaitable[T]]:
+        """Inner decorator."""
+
+        async def wrapper(self: EmberMug, *args: P.args, **kwargs: P.kwargs) -> T:
+            if self.has_attribute(attr_name) is False:
+                raise NotImplementedError(
+                    f"The {self.data.model_info.device_type} " "does not have a name attribute",
+                )
+            return await func(self, *args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
 
 class EmberMug:
     """Handle actual the actual mug connection and update states."""
@@ -72,11 +96,6 @@ class EmberMug:
         self._latest_events: dict[int, float] = {}
         self._client_kwargs: dict[str, str] = {}
 
-        # Just shortcuts, the value doesn't change once initialized
-        self.is_tumbler = self.data.model_info.device_type == DeviceType.TUMBLER
-        self.is_travel_mug = self.data.model_info.device_type == DeviceType.TRAVEL_MUG
-        self.is_cup = self.data.model_info.device_type == DeviceType.CUP
-
         logger.debug("New mug connection initialized.")
         self.set_client_options(**kwargs)
 
@@ -90,6 +109,15 @@ class EmberMug:
         """Shortcut to model name."""
         return self.data.model_info.model
 
+    @property
+    def can_write(self) -> bool:
+        """Check if the mug can support write operations."""
+        return self.data.udsk is not None
+
+    def has_attribute(self, attribute: str) -> bool:
+        """Check whether the device has the given attribute."""
+        return attribute in self.data.model_info.update_attributes
+
     async def _ensure_connection(self) -> None:
         """Connect to mug."""
         if self._connect_lock.locked():
@@ -214,19 +242,15 @@ class EmberMug:
         """Get Battery percent from mug gatt."""
         return BatteryInfo.from_bytes(await self._read(MugCharacteristic.BATTERY))
 
+    @require_attribute("led_colour")
     async def get_led_colour(self) -> Colour:
         """Get RGBA colours from mug gatt."""
-        if self.is_travel_mug is True:
-            msg = "The Travel Mug does not have an LED colour attribute"
-            raise NotImplementedError(msg)
         colour_data = await self._read(MugCharacteristic.LED)
         return Colour(*bytearray(colour_data))
 
+    @require_attribute("led_colour")
     async def set_led_colour(self, colour: Colour) -> None:
         """Set new target temp for mug."""
-        if self.is_travel_mug is True:
-            msg = "The Travel Mug does not have an LED colour attribute"
-            raise NotImplementedError(msg)
         await self._write(MugCharacteristic.LED, colour.as_bytearray())
         self.data.led_colour = colour
 
@@ -251,23 +275,19 @@ class EmberMug:
         liquid_level_bytes = await self._read(MugCharacteristic.LIQUID_LEVEL)
         return bytes_to_little_int(liquid_level_bytes)
 
+    @require_attribute("volume_level")
     async def get_volume_level(self) -> VolumeLevel | None:
         """Get volume level from mug gatt."""
-        if self.is_travel_mug is False:
-            msg = "The Mug and Cup do not have a volume level attribute"
-            raise NotImplementedError(msg)
         volume_bytes = await self._read(MugCharacteristic.VOLUME)
         volume_int = bytes_to_little_int(volume_bytes)
         return VolumeLevel.from_state(volume_int)
 
+    @require_attribute("volume_level")
     async def set_volume_level(self, volume: int | VolumeLevel) -> None:
         """Set volume_level on Travel Mug."""
         if not isinstance(volume, VolumeLevel) and isinstance(volume, int) and volume not in (0, 1, 2):
             msg = "Volume level value should be 0, 1, 2 or a VolumeLevel enum"
             raise ValueError(msg)
-        if self.is_travel_mug is False:
-            msg = "The Mug and Cup do not have a volume level attribute"
-            raise NotImplementedError(msg)
         volume_level = volume if isinstance(volume, VolumeLevel) else VolumeLevel.from_state(volume)
         await self._write(MugCharacteristic.VOLUME, bytearray([volume_level.state]))
         self.data.volume_level = volume_level
@@ -278,22 +298,18 @@ class EmberMug:
         state = bytes_to_little_int(liquid_state_bytes)
         return LiquidState(state)
 
+    @require_attribute("name")
     async def get_name(self) -> str:
         """Get mug name from gatt."""
-        if self.is_cup is True:
-            msg = "The Cup does not have a name attribute"
-            raise NotImplementedError(msg)
         name_bytes: bytearray = await self._read(MugCharacteristic.MUG_NAME)
         return bytes(name_bytes).decode("utf8")
 
+    @require_attribute("name")
     async def set_name(self, name: str) -> None:
         """Assign new name to mug."""
         if MUG_NAME_REGEX.match(name) is None:
             msg = "Name cannot contain any special characters and must be 16 characters or less"
             raise ValueError(msg)
-        if self.is_cup is True:
-            msg = "The Cup does not have a name attribute"
-            raise NotImplementedError(msg)
         await self._write(MugCharacteristic.MUG_NAME, bytearray(name.encode("utf8")))
         self.data.name = name
 

{python_ember_mug-0.9.0b1 → python_ember_mug-0.9.0b2}/ember_mug/utils.py

@@ -111,6 +111,17 @@ def get_model_from_id_and_gen(model_id: int, generation: int) -> DeviceModel | N
     return None
 
 
+def guess_model_from_name(name: str | None) -> DeviceModel | None:
+    """Guess model from BLE name."""
+    if not name:
+        return None
+    if "Travel" in name:
+        return DeviceModel.TRAVEL_MUG_12_OZ
+    if "Cup" in name:
+        return DeviceModel.CUP_6_OZ
+    return DeviceModel.UNKNOWN_DEVICE
+
+
 def get_model_info_from_advertiser_data(advertisement: AdvertisementData) -> ModelInfo:
     """Extract model info from manufacturer data in advertiser data."""
     from ember_mug.data import ModelInfo
@@ -128,7 +139,10 @@ def get_model_info_from_advertiser_data(advertisement: AdvertisementData) -> Mod
             get_model_from_id_and_gen(model_id, generation),
             get_colour_from_int(colour_id),
         )
-    return ModelInfo()
+    logger.debug(
+        "Unable to reliably determine model info from advertiser data." "Falling back to guessing based on name.",
+    )
+    return ModelInfo(guess_model_from_name(advertisement.local_name))
 
 
 async def discover_services(client: BleakClient) -> dict[str, Any]:

{python_ember_mug-0.9.0b1 → python_ember_mug-0.9.0b2}/tests/cli/test_commands.py

@@ -14,7 +14,7 @@ from ember_mug import EmberMug
 from ember_mug.consts import DEFAULT_NAME, DeviceModel, DeviceColour
 from ember_mug.cli.commands import EmberMugCli, discover, fetch_info, find_device, get_mug, get_mug_value, poll_mug
 from ember_mug.data import ModelInfo, MugData
-from tests.conftest import TEST_MAC, mock_connection, TEST_MUG_ADVERTISEMENT
+from ..conftest import TEST_MAC, mock_connection, TEST_MUG_ADVERTISEMENT
 
 
 @pytest.fixture

{python_ember_mug-0.9.0b1 → python_ember_mug-0.9.0b2}/tests/conftest.py

@@ -12,7 +12,7 @@ from bleak.backends.device import BLEDevice
 
 from ember_mug import EmberMug
 from ember_mug.data import ModelInfo, MugData
-from ember_mug.consts import DEFAULT_NAME, EMBER_BLE_SIG, MugCharacteristic, DeviceModel, DeviceColour
+from ember_mug.consts import EMBER_BLE_SIG, MugCharacteristic, DeviceModel, DeviceColour
 
 TEST_MAC = "32:36:a5:be:88:cb"
 TEST_MUG_BLUETOOTH_NAME = "Ember Ceramic Mug"
@@ -41,16 +41,20 @@ class TestManufacturerData(bytes, Enum):
 
 
 def build_advertisement_data(
-    manufacturer_data: TestManufacturerData = TestManufacturerData.UNKNOWN,
+    manufacturer_data: TestManufacturerData | None = None,
     service_uuids: list[MugCharacteristic] | None = None,
     name: str = TEST_MUG_BLUETOOTH_NAME,
 ) -> AdvertisementData:
     if service_uuids is None:
         service_uuids = [MugCharacteristic.STANDARD_SERVICE]
 
+    data_dict: dict[int, bytes] = {}
+    if manufacturer_data is not None:
+        data_dict[EMBER_BLE_SIG] = manufacturer_data
+
     return AdvertisementData(
         local_name=name,
-        manufacturer_data={EMBER_BLE_SIG: manufacturer_data},
+        manufacturer_data=data_dict,
         service_data={},
         service_uuids=[str(service) for service in service_uuids],
         tx_power=1,
@@ -59,6 +63,7 @@ def build_advertisement_data(
     )
 
 
+TEST_UNKNOWN_ADVERTISEMENT = build_advertisement_data(None)
 TEST_MUG_ADVERTISEMENT = build_advertisement_data(TestManufacturerData.MUG_2_BLACK)
 TEST_TUMBLER_ADVERTISEMENT = build_advertisement_data(TestManufacturerData.TUMBLER)
 TEST_TRAVEL_MUG_ADVERTISEMENT = build_advertisement_data(
@@ -83,5 +88,5 @@ async def ember_mug(ble_device: BLEDevice) -> AsyncGenerator[EmberMug | Mock, No
         ble_device,
         ModelInfo(DeviceModel.MUG_2_10_OZ, DeviceColour.BLACK),
     )
-    mug._client = AsyncMock()
+    mug._client = Mock()
     yield mug

{python_ember_mug-0.9.0b1 → python_ember_mug-0.9.0b2}/tests/test_connection.py

@@ -10,16 +10,15 @@ from bleak import BleakError
 from bleak.backends.device import BLEDevice
 
 from ember_mug.consts import (
-    DEFAULT_NAME,
     INITIAL_ATTRS,
     UPDATE_ATTRS,
     MugCharacteristic,
     TemperatureUnit,
     VolumeLevel,
+    DeviceModel,
 )
 from ember_mug.data import Colour, ModelInfo
 from ember_mug.mug import EmberMug
-from .conftest import TEST_MUG_BLUETOOTH_NAME
 
 if TYPE_CHECKING:
 
@@ -159,7 +158,7 @@ async def test_read(
     ember_mug: MockMug,
 ) -> None:
     with patch.object(ember_mug, "_ensure_connection", AsyncMock()):
-        ember_mug._client.read_gatt_char.return_value = b"TEST"
+        ember_mug._client.read_gatt_char = AsyncMock(return_value=b"TEST")
         await ember_mug._read(MugCharacteristic.MUG_NAME)
         ember_mug._client.read_gatt_char.assert_called_with(
             MugCharacteristic.MUG_NAME.uuid,
@@ -175,6 +174,7 @@ async def test_read(
 async def test_write(mock_logger: Mock, ember_mug: MockMug) -> None:
     with patch.object(ember_mug, "_ensure_connection", AsyncMock()):
         test_name = bytearray(b"TEST")
+        ember_mug._client.write_gatt_char = AsyncMock()
         await ember_mug._write(
             MugCharacteristic.MUG_NAME,
             test_name,
@@ -190,7 +190,7 @@ async def test_write(mock_logger: Mock, ember_mug: MockMug) -> None:
         )
 
         ember_mug._client = AsyncMock()
-        ember_mug._client.write_gatt_char.side_effect = BleakError
+        ember_mug._client.write_gatt_char = AsyncMock(side_effect=BleakError)
         with pytest.raises(BleakError):
             await ember_mug._write(
                 MugCharacteristic.MUG_NAME,
@@ -219,9 +219,24 @@ def test_set_device(ember_mug: MockMug) -> None:
     assert ember_mug.device.address == new_device.address
 
 
+def test_can_write(ember_mug: MockMug) -> None:
+    ember_mug.data.udsk = "non-empty"
+    assert ember_mug.can_write is True
+
+    ember_mug.data.udsk = None
+    assert ember_mug.can_write is False
+
+
+def test_has_attribute(ember_mug: MockMug) -> None:
+    ember_mug.data.model_info.model = DeviceModel.CUP_6_OZ
+    assert ember_mug.has_attribute("name") is False
+    ember_mug.data.model_info = ModelInfo(DeviceModel.MUG_2_10_OZ)
+    assert ember_mug.has_attribute("name") is True
+
+
 async def test_get_mug_meta(ember_mug: MockMug) -> None:
     with patch.object(ember_mug, "_ensure_connection", AsyncMock()):
-        ember_mug._client.read_gatt_char.return_value = b"Yw====-ABCDEFGHIJ"
+        ember_mug._client.read_gatt_char = AsyncMock(return_value=b"Yw====-ABCDEFGHIJ")
         meta = await ember_mug.get_meta()
         assert meta.mug_id == "WXc9PT09"
         assert meta.serial_number == "ABCDEFGHIJ"
@@ -230,7 +245,7 @@ async def test_get_mug_meta(ember_mug: MockMug) -> None:
 
 async def test_get_mug_battery(ember_mug: MockMug) -> None:
     with patch.object(ember_mug, "_ensure_connection", AsyncMock()):
-        ember_mug._client.read_gatt_char.return_value = b"5\x01"
+        ember_mug._client.read_gatt_char = AsyncMock(return_value=b"5\x01")
         battery = await ember_mug.get_battery()
         assert battery.percent == 53.00
         assert battery.on_charging_base is True
@@ -239,7 +254,7 @@ async def test_get_mug_battery(ember_mug: MockMug) -> None:
 
 async def test_get_mug_led_colour(ember_mug: MockMug) -> None:
     with patch.object(ember_mug, "_ensure_connection", AsyncMock()):
-        ember_mug._client.read_gatt_char.return_value = b"\xf4\x00\xa1\xff"
+        ember_mug._client.read_gatt_char = AsyncMock(return_value=b"\xf4\x00\xa1\xff")
         colour = await ember_mug.get_led_colour()
         assert colour.as_hex() == "#f400a1"
         ember_mug._client.read_gatt_char.assert_called_once_with(MugCharacteristic.LED.uuid)
@@ -247,6 +262,7 @@ async def test_get_mug_led_colour(ember_mug: MockMug) -> None:
 
 async def test_set_mug_led_colour(ember_mug: MockMug) -> None:
     mock_ensure_connection = AsyncMock()
+    ember_mug._client.write_gatt_char = AsyncMock()
     with patch.object(ember_mug, "_ensure_connection", mock_ensure_connection):
         await ember_mug.set_led_colour(Colour(244, 0, 161))
         mock_ensure_connection.assert_called_once()
@@ -257,8 +273,9 @@ async def test_set_mug_led_colour(ember_mug: MockMug) -> None:
 
 
 async def test_set_volume_level_travel_mug(ember_mug: MockMug) -> None:
-    ember_mug.is_travel_mug = True
+    ember_mug.data.model_info.model = DeviceModel.TRAVEL_MUG_12_OZ
     mock_ensure_connection = AsyncMock()
+    ember_mug._client.write_gatt_char = AsyncMock()
     with patch.object(ember_mug, "_ensure_connection", mock_ensure_connection):
         await ember_mug.set_volume_level(VolumeLevel.HIGH)
         mock_ensure_connection.assert_called_once()
@@ -288,13 +305,14 @@ async def test_set_volume_level_mug(ember_mug: MockMug) -> None:
 
 async def test_get_mug_target_temp(ember_mug: MockMug) -> None:
     with patch.object(ember_mug, "_ensure_connection", AsyncMock()):
-        ember_mug._client.read_gatt_char.return_value = b"\xcd\x15"
+        ember_mug._client.read_gatt_char = AsyncMock(return_value=b"\xcd\x15")
         assert (await ember_mug.get_target_temp()) == 55.81
         ember_mug._client.read_gatt_char.assert_called_once_with(MugCharacteristic.TARGET_TEMPERATURE.uuid)
 
 
 async def test_set_mug_target_temp(ember_mug: MockMug) -> None:
     mock_ensure_connection = AsyncMock()
+    ember_mug._client.write_gatt_char = AsyncMock()
     with patch.object(ember_mug, "_ensure_connection", mock_ensure_connection):
         await ember_mug.set_target_temp(55.81)
         mock_ensure_connection.assert_called_once()
@@ -306,37 +324,42 @@ async def test_set_mug_target_temp(ember_mug: MockMug) -> None:
 
 async def test_get_mug_current_temp(ember_mug: MockMug) -> None:
     with patch.object(ember_mug, "_ensure_connection", AsyncMock()):
-        ember_mug._client.read_gatt_char.return_value = b"\xcd\x15"
+        ember_mug._client.read_gatt_char = AsyncMock(return_value=b"\xcd\x15")
         assert (await ember_mug.get_current_temp()) == 55.81
         ember_mug._client.read_gatt_char.assert_called_once_with(MugCharacteristic.CURRENT_TEMPERATURE.uuid)
 
 
 async def test_get_mug_liquid_level(ember_mug: MockMug) -> None:
     with patch.object(ember_mug, "_ensure_connection", AsyncMock()):
-        ember_mug._client.read_gatt_char.return_value = b"\n"
+        ember_mug._client.read_gatt_char = AsyncMock(return_value=b"\n")
         assert (await ember_mug.get_liquid_level()) == 10
         ember_mug._client.read_gatt_char.assert_called_once_with(MugCharacteristic.LIQUID_LEVEL.uuid)
 
 
 async def test_get_mug_liquid_state(ember_mug: MockMug) -> None:
     with patch.object(ember_mug, "_ensure_connection", AsyncMock()):
-        ember_mug._client.read_gatt_char.return_value = b"\x06"
+        ember_mug._client.read_gatt_char = AsyncMock(return_value=b"\x06")
         assert (await ember_mug.get_liquid_state()) == 6
         ember_mug._client.read_gatt_char.assert_called_once_with(MugCharacteristic.LIQUID_STATE.uuid)
 
 
 async def test_get_mug_name(ember_mug: MockMug) -> None:
     with patch.object(ember_mug, "_ensure_connection", AsyncMock()):
-        ember_mug._client.read_gatt_char.return_value = b"Mug Name"
+        ember_mug._client.read_gatt_char = AsyncMock(return_value=b"Mug Name")
         assert (await ember_mug.get_name()) == "Mug Name"
         ember_mug._client.read_gatt_char.assert_called_once_with(MugCharacteristic.MUG_NAME.uuid)
 
+        ember_mug.data.model_info = ModelInfo(DeviceModel.CUP_6_OZ)
+        with pytest.raises(NotImplementedError):
+            await ember_mug.get_name()
+
 
 async def test_set_mug_name(ember_mug: MockMug) -> None:
     with patch.object(ember_mug, "_ensure_connection", AsyncMock()), pytest.raises(ValueError):
         await ember_mug.set_name("Hé!")
 
     mock_ensure_connection = AsyncMock()
+    ember_mug._client.write_gatt_char = AsyncMock()
     with patch.object(ember_mug, "_ensure_connection", mock_ensure_connection):
         await ember_mug.set_name("Mug name")
         mock_ensure_connection.assert_called()
@@ -345,16 +368,21 @@ async def test_set_mug_name(ember_mug: MockMug) -> None:
             bytearray(b"Mug name"),
         )
 
+        ember_mug.data.model_info = ModelInfo(DeviceModel.CUP_6_OZ)
+        with pytest.raises(NotImplementedError):
+            await ember_mug.set_name("Test")
+
 
 async def test_get_mug_udsk(ember_mug: MockMug) -> None:
     with patch.object(ember_mug, "_ensure_connection", AsyncMock()):
-        ember_mug._client.read_gatt_char.return_value = b"abcd12345"
+        ember_mug._client.read_gatt_char = AsyncMock(return_value=b"abcd12345")
         assert (await ember_mug.get_udsk()) == "YWJjZDEyMzQ1"
         ember_mug._client.read_gatt_char.assert_called_once_with(MugCharacteristic.UDSK.uuid)
 
 
 async def test_set_mug_udsk(ember_mug: MockMug) -> None:
     mock_ensure_connection = AsyncMock()
+    ember_mug._client.write_gatt_char = AsyncMock()
     with patch.object(ember_mug, "_ensure_connection", mock_ensure_connection):
         await ember_mug.set_udsk("abcd12345")
         mock_ensure_connection.assert_called_once()
@@ -366,26 +394,26 @@ async def test_set_mug_udsk(ember_mug: MockMug) -> None:
 
 async def test_get_mug_dsk(ember_mug: MockMug) -> None:
     with patch.object(ember_mug, "_ensure_connection", AsyncMock()):
-        ember_mug._client.read_gatt_char.return_value = b"abcd12345"
+        ember_mug._client.read_gatt_char = AsyncMock(return_value=b"abcd12345")
         assert (await ember_mug.get_dsk()) == "YWJjZDEyMzQ1"
-        ember_mug._client.read_gatt_char.return_value = b"something else"
+        ember_mug._client.read_gatt_char = AsyncMock(return_value=b"something else")
         assert (await ember_mug.get_dsk()) == "c29tZXRoaW5nIGVsc2U="
         ember_mug._client.read_gatt_char.assert_called_with(MugCharacteristic.DSK.uuid)
 
 
 async def test_get_mug_temperature_unit(ember_mug: MockMug) -> None:
     with patch.object(ember_mug, "_ensure_connection", AsyncMock()):
-        ember_mug._client.read_gatt_char.return_value = b"\x01"
+        ember_mug._client.read_gatt_char = AsyncMock(return_value=b"\x01")
         assert (await ember_mug.get_temperature_unit()) == TemperatureUnit.FAHRENHEIT
         ember_mug._client.read_gatt_char.assert_called_once_with(MugCharacteristic.TEMPERATURE_UNIT.uuid)
-        ember_mug._client.read_gatt_char.reset_mock()
-        ember_mug._client.read_gatt_char.return_value = b"\x00"
+        ember_mug._client.read_gatt_char = AsyncMock(return_value=b"\x00")
         assert (await ember_mug.get_temperature_unit()) == TemperatureUnit.CELSIUS
         ember_mug._client.read_gatt_char.assert_called_once_with(MugCharacteristic.TEMPERATURE_UNIT.uuid)
 
 
 async def test_set_mug_temperature_unit(ember_mug: MockMug) -> None:
     mock_ensure_connection = AsyncMock()
+    ember_mug._client.write_gatt_char = AsyncMock()
     with patch.object(ember_mug, "_ensure_connection", mock_ensure_connection):
         await ember_mug.set_temperature_unit(TemperatureUnit.CELSIUS)
         mock_ensure_connection.assert_called_once()
@@ -400,7 +428,7 @@ async def test_mug_ensure_correct_unit(ember_mug: MockMug) -> None:
         ember_mug.data.temperature_unit = TemperatureUnit.CELSIUS
         ember_mug.data.use_metric = True
         mock_set_temp = AsyncMock(return_value=None)
-        with patch.object(ember_mug, 'set_temperature_unit', mock_set_temp):
+        with patch.object(ember_mug, "set_temperature_unit", mock_set_temp):
             await ember_mug.ensure_correct_unit()
             mock_set_temp.assert_not_called()
             ember_mug.data.temperature_unit = TemperatureUnit.FAHRENHEIT
@@ -410,14 +438,14 @@ async def test_mug_ensure_correct_unit(ember_mug: MockMug) -> None:
 
 async def test_get_mug_battery_voltage(ember_mug: MockMug) -> None:
     with patch.object(ember_mug, "_ensure_connection", AsyncMock()):
-        ember_mug._client.read_gatt_char.return_value = b"\x01"
+        ember_mug._client.read_gatt_char = AsyncMock(return_value=b"\x01")
         assert (await ember_mug.get_battery_voltage()) == 1
         ember_mug._client.read_gatt_char.assert_called_once_with(MugCharacteristic.CONTROL_REGISTER_DATA.uuid)
 
 
 async def test_get_mug_date_time_zone(ember_mug: MockMug) -> None:
     with patch.object(ember_mug, "_ensure_connection", AsyncMock()):
-        ember_mug._client.read_gatt_char.return_value = b"c\x0f\xf6\x00"
+        ember_mug._client.read_gatt_char = AsyncMock(return_value=b"c\x0f\xf6\x00")
         date_time = await ember_mug.get_date_time_zone()
         assert isinstance(date_time, datetime)
         assert date_time.timestamp() == 1661990400.0
@@ -426,7 +454,7 @@ async def test_get_mug_date_time_zone(ember_mug: MockMug) -> None:
 
 async def test_read_firmware(ember_mug: MockMug) -> None:
     with patch.object(ember_mug, "_ensure_connection", AsyncMock()):
-        ember_mug._client.read_gatt_char.return_value = b"c\x01\x80\x00\x12\x00"
+        ember_mug._client.read_gatt_char = AsyncMock(return_value=b"c\x01\x80\x00\x12\x00")
         firmware = await ember_mug.get_firmware()
         assert firmware.version == 355
         assert firmware.hardware == 128
@@ -453,8 +481,8 @@ async def test_mug_update_multiple(ember_mug: MockMug) -> None:
     mock_get_name = AsyncMock(return_value="name")
     mock_update_info = AsyncMock()
 
-    with patch.multiple(ember_mug, _ensure_connection=AsyncMock(), get_name=mock_get_name):
-        with patch.object(ember_mug.data, 'update_info', mock_update_info):
+    with patch.multiple(ember_mug, get_name=mock_get_name):
+        with patch.object(ember_mug.data, "update_info", mock_update_info):
             await ember_mug._update_multiple({"name"})
             mock_get_name.assert_called_once()
             mock_update_info.assert_called_once_with(name="name")
@@ -464,10 +492,10 @@ async def test_mug_update_queued_attributes(ember_mug: MockMug) -> None:
     mock_get_name = AsyncMock(return_value="name")
     mock_update_info = AsyncMock()
 
-    with patch.multiple(ember_mug, _ensure_connection=AsyncMock(), get_name=mock_get_name):
+    with patch.multiple(ember_mug, get_name=mock_get_name):
         ember_mug._queued_updates = set()
         assert (await ember_mug.update_queued_attributes()) == []
-        with patch.object(ember_mug.data, 'update_info', mock_update_info):
+        with patch.object(ember_mug.data, "update_info", mock_update_info):
             ember_mug._queued_updates = {"name"}
             await ember_mug.update_queued_attributes()
             mock_update_info.assert_called_once_with(name="name")

{python_ember_mug-0.9.0b1 → python_ember_mug-0.9.0b2}/tests/test_utils.py

@@ -1,8 +1,8 @@
 """Tests for `ember_mug.utils`."""
-from unittest.mock import AsyncMock, MagicMock, Mock, call, patch
+from unittest.mock import AsyncMock, MagicMock, Mock, call, patch, PropertyMock
 
 import pytest
-from bleak import AdvertisementData
+from bleak import AdvertisementData, BleakError
 
 from ember_mug.consts import DeviceModel, DeviceColour, MugCharacteristic
 from ember_mug.utils import (
@@ -16,11 +16,14 @@ from ember_mug.utils import (
     get_colour_from_int,
     get_model_from_single_int_and_services,
     get_model_from_id_and_gen,
+    guess_model_from_name,
+)
+from tests.conftest import (
+    TEST_TUMBLER_ADVERTISEMENT,
+    TEST_MUG_ADVERTISEMENT,
+    TEST_TRAVEL_MUG_ADVERTISEMENT,
+    TEST_UNKNOWN_ADVERTISEMENT,
 )
-from tests.conftest import TEST_TUMBLER_ADVERTISEMENT, TEST_MUG_ADVERTISEMENT, TEST_TRAVEL_MUG_ADVERTISEMENT
-
-
-# 131, 147->152, 150-151
 
 
 def test_bytes_to_little_int() -> None:
@@ -71,6 +74,7 @@ def test_get_colour_from_int(colour_id: int, expected_colour: DeviceColour | Non
         (65, [], DeviceModel.MUG_1_14_OZ),
         (-127, [], DeviceModel.MUG_2_10_OZ),
         (-63, [], DeviceModel.MUG_2_14_OZ),
+        (-60, [], DeviceModel.CUP_6_OZ),
         (0, [], None),
     ],
 )
@@ -82,6 +86,23 @@ def test_get_model_from_single_int_and_services(
     assert get_model_from_single_int_and_services(model_id, service_uuids) == expected_model
 
 
+@pytest.mark.parametrize(
+    "model_name,expected_model",
+    [
+        ("", None),
+        ("Test", DeviceModel.UNKNOWN_DEVICE),
+        ("Ember Ceramic Mug", DeviceModel.UNKNOWN_DEVICE),
+        ("Ember Cup", DeviceModel.CUP_6_OZ),
+        ("Ember Travel Mug", DeviceModel.TRAVEL_MUG_12_OZ),
+    ],
+)
+def test_guess_model_from_name(
+    model_name: str,
+    expected_model: DeviceModel | None,
+) -> None:
+    assert guess_model_from_name(model_name) == expected_model
+
+
 @pytest.mark.parametrize(
     "model_id,generation,expected_model",
     [
@@ -106,6 +127,7 @@ def test_get_model_from_id_and_gen(
 @pytest.mark.parametrize(
     "advertisement,expected_model,expected_colour",
     [
+        (TEST_UNKNOWN_ADVERTISEMENT, DeviceModel.UNKNOWN_DEVICE, None),
         (TEST_MUG_ADVERTISEMENT, DeviceModel.MUG_2_10_OZ, DeviceColour.BLACK),
         (TEST_TUMBLER_ADVERTISEMENT, DeviceModel.TUMBLER_16_OZ, DeviceColour.BLACK),
         (TEST_TRAVEL_MUG_ADVERTISEMENT, DeviceModel.TRAVEL_MUG_12_OZ, DeviceColour.RED),