Replace set_device with ble_event_callback, fix tests and add more

CHANGELOG.md

@@ -12,6 +12,7 @@
 
 * Bumped minimum version of bleak to 0.21.0
 * Discover method changed to use advertisement information
+* `set_device` method replaced by `ble_event_callback` and it now also updated model info if not yet set.
 * `model` renamed to `model_info`
 * `include_extra` option removed. `debug` has a similar function now.
 

ember_mug/cli/commands.py

@@ -134,10 +134,10 @@ async def get_mug_value(args: Namespace) -> None:
 async def set_mug_value(args: Namespace) -> None:
     """Set one or more values on the mug."""
     attrs = ("name", "target_temp", "temperature_unit", "led_colour", "volume_level")
-    values = [(attr, value) for attr in attrs if (value := getattr(args, attr))]
+    values = [(attr, value) for attr in attrs if (value := getattr(args, attr, None))]
     if not values:
         print("Please specify at least one attribute and value to set.")
-        options = [f"--{a}" for a in attrs]
+        options = [f"--{a.replace('_', '-')}" for a in attrs]
         print(f'Options: {", ".join(options)}')
         sys.exit(1)
 

ember_mug/mug.py

@@ -10,7 +10,7 @@
 from time import time
 from typing import TYPE_CHECKING, Any, Concatenate, Literal, ParamSpec, TypeVar
 
-from bleak import BleakClient, BleakError
+from bleak import AdvertisementData, BleakClient, BleakError
 from bleak_retry_connector import establish_connection
 
 from .consts import (
@@ -31,6 +31,7 @@
     decode_byte_string,
     discover_services,
     encode_byte_string,
+    get_model_info_from_advertiser_data,
     temp_from_bytes,
 )
 
@@ -99,15 +100,26 @@ def __init__(
         logger.debug("New mug connection initialized.")
         self.set_client_options(**kwargs)
 
-    def set_device(self, ble_device: BLEDevice) -> None:
-        """Set the ble device."""
-        logger.debug("Set new device from %s to %s", self.device, ble_device)
+    def ble_event_callback(self, ble_device: BLEDevice, advertisement_data: AdvertisementData) -> None:
+        """Update BLE Device and, if needed, model information."""
         self.device = ble_device
+        logger.debug("Set new device from %s to %s", self.device, ble_device)
+        if (
+            not self.data.model_info.model
+            and advertisement_data.manufacturer_data
+            and (model_info := get_model_info_from_advertiser_data(advertisement_data))
+        ):
+            logger.debug(
+                "Updated model info from advertisement data (%s) -> %s",
+                advertisement_data,
+                model_info,
+            )
+            self.data.model_info = model_info
 
     @cached_property
     def model_name(self) -> str | None:
         """Shortcut to model name."""
-        return self.data.model_info.model
+        return self.data.model_info.model.value if self.data.model_info.model else None
 
     @property
     def can_write(self) -> bool:

tests/cli/test_commands.py

@@ -11,8 +11,17 @@
 from pytest import CaptureFixture
 
 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.consts import DeviceModel, DeviceColour
+from ember_mug.cli.commands import (
+    EmberMugCli,
+    discover,
+    fetch_info,
+    find_device,
+    get_mug,
+    get_mug_value,
+    poll_mug,
+    set_mug_value,
+)
 from ember_mug.data import ModelInfo, MugData
 from ..conftest import TEST_MAC, mock_connection, TEST_MUG_ADVERTISEMENT
 
@@ -223,6 +232,33 @@ async def test_get_mug_value(
     captured = capsys.readouterr()
     assert captured.out == "55.5\ntest\n"
 
+    mock_mug_with_connection.get_name.side_effect = NotImplementedError
+    with pytest.raises(SystemExit):
+        args = mock_namespace(attributes=["name"], raw=True)
+        await get_mug_value(args)
+
+
+async def test_set_mug_value_no_value(capsys: CaptureFixture) -> None:
+    with pytest.raises(SystemExit):
+        await set_mug_value(Namespace())
+    captured = capsys.readouterr()
+    assert captured.out == (
+        "Please specify at least one attribute and value to set.\n"
+        "Options: --name, --target-temp, --temperature-unit, --led-colour, --volume-level\n"
+    )
+
+
+async def test_set_mug_value(mock_mug_with_connection: AsyncMock, capsys: CaptureFixture) -> None:
+    mock_mug_with_connection.data = MugData(ModelInfo())
+    args = mock_namespace(name="test")
+    await set_mug_value(args)
+    mock_mug_with_connection.set_name.assert_called_once_with("test")
+
+    mock_mug_with_connection.reset_mock()
+    mock_mug_with_connection.set_name.side_effect = NotImplementedError("Unable to set name on Cup")
+    with pytest.raises(SystemExit):
+        await set_mug_value(args)
+
 
 def test_ember_cli():
     cli = EmberMugCli()

tests/conftest.py

@@ -78,7 +78,7 @@ def ble_device_fixture() -> BLEDevice:
 
 
 @pytest.fixture()
-def mug_data(ble_device: BLEDevice) -> MugData:
+def mug_data() -> MugData:
     return MugData(ModelInfo())
 
 

tests/test_connection.py

@@ -19,6 +19,7 @@
 )
 from ember_mug.data import Colour, ModelInfo
 from ember_mug.mug import EmberMug
+from tests.conftest import TEST_MUG_ADVERTISEMENT
 
 if TYPE_CHECKING:
 
@@ -207,15 +208,17 @@ async def test_write(mock_logger: Mock, ember_mug: MockMug) -> None:
         assert isinstance(exception, BleakError)
 
 
-def test_set_device(ember_mug: MockMug) -> None:
+def test_ble_event_callback(ember_mug: MockMug) -> None:
     new_device = BLEDevice(
         address="BA:36:a5:be:88:cb",
         name="Ember Ceramic Mug",
         details={},
         rssi=1,
     )
+    ember_mug.data.model_info.model = None
     assert ember_mug.device.address != new_device.address
-    ember_mug.set_device(new_device)
+    ember_mug.ble_event_callback(new_device, TEST_MUG_ADVERTISEMENT)
+    assert ember_mug.model_name == DeviceModel.MUG_2_10_OZ
     assert ember_mug.device.address == new_device.address
 
 
@@ -479,23 +482,21 @@ async def test_mug_update_all(ember_mug: MockMug) -> None:
 
 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, get_name=mock_get_name):
-        with patch.object(ember_mug.data, "update_info", mock_update_info):
+        with patch.object(ember_mug.data, "update_info") as mock_update_info:
             await ember_mug._update_multiple({"name"})
             mock_get_name.assert_called_once()
             mock_update_info.assert_called_once_with(name="name")
 
 
 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, 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") as mock_update_info:
             ember_mug._queued_updates = {"name"}
             await ember_mug.update_queued_attributes()
             mock_update_info.assert_called_once_with(name="name")