From 15aaba2802f76401b40e22e14a15d6dad388cada Mon Sep 17 00:00:00 2001
From: MattHag <16444067+MattHag@users.noreply.github.com>
Date: Mon, 30 Dec 2024 23:43:24 +0100
Subject: [PATCH] notification: Refactor process_receiver_notification

Remove repeated code pattern with generalized implementation. Aim
towards easy extension and code readability.

Related #2273
---
 lib/logitech_receiver/notifications.py        | 28 +++++++------
 tests/logitech_receiver/test_notifications.py | 42 +++++++++++++++++++
 2 files changed, 58 insertions(+), 12 deletions(-)

diff --git a/lib/logitech_receiver/notifications.py b/lib/logitech_receiver/notifications.py
index 59ea1d1d..6bde5b0b 100644
--- a/lib/logitech_receiver/notifications.py
+++ b/lib/logitech_receiver/notifications.py
@@ -48,6 +48,8 @@ if typing.TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
+NotificationHandler = typing.Callable[["Receiver", "HIDPPNotification"], bool]
+
 _hidpp10 = hidpp10.Hidpp10()
 _hidpp20 = hidpp20.Hidpp20()
 _F = hidpp20_constants.SupportedFeature
@@ -68,18 +70,20 @@ def process(device: Device | Receiver, notification: HIDPPNotification):
 
 def process_receiver_notification(receiver: Receiver, hidpp_notification: HIDPPNotification) -> bool | None:
     """Process event messages from receivers."""
-    if hidpp_notification.sub_id == Notification.PAIRING_LOCK:
-        return handle_pairing_lock(receiver, hidpp_notification)
-    elif hidpp_notification.sub_id == Registers.DISCOVERY_STATUS_NOTIFICATION:  # Bolt pairing
-        return handle_discovery_status(receiver, hidpp_notification)
-    elif hidpp_notification.sub_id == Registers.DEVICE_DISCOVERY_NOTIFICATION:  # Bolt pairing
-        return handle_device_discovery(receiver, hidpp_notification)
-    elif hidpp_notification.sub_id == Registers.PAIRING_STATUS_NOTIFICATION:  # Bolt pairing
-        return handle_pairing_status(receiver, hidpp_notification)
-    elif hidpp_notification.sub_id == Registers.PASSKEY_REQUEST_NOTIFICATION:  # Bolt pairing
-        return handle_passkey_request(receiver, hidpp_notification)
-    elif hidpp_notification.sub_id == Registers.PASSKEY_PRESSED_NOTIFICATION:  # Bolt pairing
-        return handle_passkey_pressed(receiver, hidpp_notification)
+    event_handler_mapping: dict[int, NotificationHandler] = {
+        Notification.PAIRING_LOCK: handle_pairing_lock,
+        Registers.DEVICE_DISCOVERY_NOTIFICATION: handle_device_discovery,
+        Registers.DISCOVERY_STATUS_NOTIFICATION: handle_discovery_status,
+        Registers.PAIRING_STATUS_NOTIFICATION: handle_pairing_status,
+        Registers.PASSKEY_PRESSED_NOTIFICATION: handle_passkey_pressed,
+        Registers.PASSKEY_REQUEST_NOTIFICATION: handle_passkey_request,
+    }
+
+    try:
+        handler_func = event_handler_mapping[hidpp_notification.sub_id]
+        return handler_func(receiver, hidpp_notification)
+    except KeyError:
+        pass
 
     assert hidpp_notification.sub_id in [
         Notification.CONNECT_DISCONNECT,
diff --git a/tests/logitech_receiver/test_notifications.py b/tests/logitech_receiver/test_notifications.py
index fd9f578e..6a680255 100644
--- a/tests/logitech_receiver/test_notifications.py
+++ b/tests/logitech_receiver/test_notifications.py
@@ -277,3 +277,45 @@ def test_process_feature_notification(mocker, hidpp_notification, feature):
     result = notifications._process_feature_notification(fake_device, hidpp_notification, feature)
 
     assert result is True
+
+
+def test_process_receiver_notification_invalid(mocker):
+    invalid_sub_id = 0x30
+    notification_data = b"\x02"
+    notification = HIDPPNotification(0, 0, invalid_sub_id, 0, notification_data)
+    mock_receiver = mocker.Mock()
+
+    with pytest.raises(AssertionError):
+        notifications.process_receiver_notification(mock_receiver, notification)
+
+
+def test_handle_device_discovery():
+    receiver: Receiver = Receiver(MockLowLevelInterface(), None, {}, True, None, None)
+    sub_id = Registers.DISCOVERY_STATUS_NOTIFICATION
+    data = b"\x01\x02\x03\x04\x05\x06"
+    notification = HIDPPNotification(0, 0, sub_id, 0, data)
+
+    result = notifications.handle_device_discovery(receiver, notification)
+
+    assert result
+
+
+def test_handle_passkey_request(mocker):
+    receiver_mock = mocker.Mock()
+    data = b"\x01"
+    notification = HIDPPNotification(0, 0, 0, 0, data)
+
+    result = notifications.handle_passkey_request(receiver_mock, notification)
+
+    assert result is True
+
+
+def test_handle_passkey_pressed(mocker):
+    receiver = mocker.Mock()
+    sub_id = Registers.DISCOVERY_STATUS_NOTIFICATION
+    data = b"\x01\x02\x03\x04\x05\x06"
+    notification = HIDPPNotification(0, 0, sub_id, 0, data)
+
+    result = notifications.handle_passkey_pressed(receiver, notification)
+
+    assert result is True