diff --git a/src/blinkstick/clients/blinkstick.py b/src/blinkstick/clients/blinkstick.py index aa00e7c..a3545fb 100644 --- a/src/blinkstick/clients/blinkstick.py +++ b/src/blinkstick/clients/blinkstick.py @@ -14,7 +14,7 @@ ) from blinkstick.decorators import no_backend_required from blinkstick.devices import BlinkStickDevice -from blinkstick.enums import BlinkStickVariant +from blinkstick.enums import BlinkStickVariant, Mode from blinkstick.exceptions import NotConnected from blinkstick.utilities import string_to_info_block_data @@ -373,7 +373,7 @@ def get_led_data(self, count: int) -> list[int]: return device_bytes[2 : 2 + count * 3] - def set_mode(self, mode: int) -> None: + def set_mode(self, mode: Mode | int) -> None: """ Set backend mode for BlinkStick Pro. Device currently supports the following modes: @@ -388,6 +388,10 @@ def set_mode(self, mode: int) -> None: @type mode: int @param mode: Device mode to set """ + # If mode is an enum, get the value + # this will allow the user to pass in the enum directly, and also gate the value to the enum values + if not isinstance(mode, int): + mode = Mode(mode).value control_string = bytes(bytearray([4, mode])) self.backend.control_transfer(0x20, 0x9, 0x0004, 0, control_string) diff --git a/src/blinkstick/enums.py b/src/blinkstick/enums.py index fcc3e2e..28cfd7a 100644 --- a/src/blinkstick/enums.py +++ b/src/blinkstick/enums.py @@ -1,6 +1,6 @@ from __future__ import annotations -from enum import Enum +from enum import Enum, IntEnum class BlinkStickVariant(Enum): @@ -38,3 +38,9 @@ def from_version_attrs( elif version_attribute == 0x203: return BlinkStickVariant.BLINKSTICK_FLEX return BlinkStickVariant.UNKNOWN + + +class Mode(IntEnum): + RGB = 1 + RGB_INVERSE = 2 + ADDRESSABLE = 3 diff --git a/tests/clients/test_blinkstick.py b/tests/clients/test_blinkstick.py index 067ac7a..b5c1244 100644 --- a/tests/clients/test_blinkstick.py +++ b/tests/clients/test_blinkstick.py @@ -3,7 +3,7 @@ import pytest from blinkstick.colors import ColorFormat -from blinkstick.enums import BlinkStickVariant +from blinkstick.enums import BlinkStickVariant, Mode from blinkstick.clients.blinkstick import BlinkStick from pytest_mock import MockFixture @@ -310,3 +310,26 @@ def test_inverse_does_not_affect_max_rgb_value(make_blinkstick): bs.set_max_rgb_value(100) bs.set_inverse(True) assert bs.get_max_rgb_value() == 100 + + +@pytest.mark.parametrize( + "mode, is_valid", + [ + (1, True), + (2, True), + (3, True), + (4, False), + (-1, False), + (Mode.RGB, True), + (Mode.RGB_INVERSE, True), + (Mode.ADDRESSABLE, True), + ], +) +def test_set_mode_raises_on_invalid_mode(make_blinkstick, mode, is_valid): + """Test that set_mode raises an exception when an invalid mode is passed.""" + bs = make_blinkstick() + if is_valid: + bs.set_mode(mode) + else: + with pytest.raises(ValueError): + bs.set_mode("invalid_mode") # noqa