From 06fd32b5016c00c4d4c6362d4031276791d18c56 Mon Sep 17 00:00:00 2001
From: MattHag <16444067+MattHag@users.noreply.github.com>
Date: Sun, 13 Oct 2024 12:06:35 +0200
Subject: [PATCH] Test and refactor process_notification

Related #2273
---
 lib/logitech_receiver/diversion.py        | 85 ++++++++++++-----------
 tests/logitech_receiver/test_diversion.py | 29 ++++++++
 2 files changed, 72 insertions(+), 42 deletions(-)

diff --git a/lib/logitech_receiver/diversion.py b/lib/logitech_receiver/diversion.py
index 8652814a..b9c368ab 100644
--- a/lib/logitech_receiver/diversion.py
+++ b/lib/logitech_receiver/diversion.py
@@ -1495,51 +1495,52 @@ def evaluate_rules(feature, notification: HIDPPNotification, device):
     rules.evaluate(feature, notification, device, True)
 
 
-# process a notification
-def process_notification(device, notification: HIDPPNotification, feature):
+def process_notification(device, notification: HIDPPNotification, feature) -> None:
+    """Processes HID++ notifications."""
     global keys_down, g_keys_down, m_keys_down, mr_key_down, key_down, key_up, thumb_wheel_displacement
     key_down, key_up = None, None
     # need to keep track of keys that are down to find a new key down
-    if feature == FEATURE.REPROG_CONTROLS_V4 and notification.address == 0x00:
-        new_keys_down = struct.unpack("!4H", notification.data[:8])
-        for key in new_keys_down:
-            if key and key not in keys_down:
-                key_down = key
-        for key in keys_down:
-            if key and key not in new_keys_down:
-                key_up = key
-        keys_down = new_keys_down
-    # and also G keys down
-    elif feature == FEATURE.GKEY and notification.address == 0x00:
-        new_g_keys_down = struct.unpack("<I", notification.data[:4])[0]
-        for i in range(32):
-            if new_g_keys_down & (0x01 << i) and not g_keys_down & (0x01 << i):
-                key_down = CONTROL["G" + str(i + 1)]
-            if g_keys_down & (0x01 << i) and not new_g_keys_down & (0x01 << i):
-                key_up = CONTROL["G" + str(i + 1)]
-        g_keys_down = new_g_keys_down
-    # and also M keys down
-    elif feature == FEATURE.MKEYS and notification.address == 0x00:
-        new_m_keys_down = struct.unpack("!1B", notification.data[:1])[0]
-        for i in range(1, 9):
-            if new_m_keys_down & (0x01 << (i - 1)) and not m_keys_down & (0x01 << (i - 1)):
-                key_down = CONTROL["M" + str(i)]
-            if m_keys_down & (0x01 << (i - 1)) and not new_m_keys_down & (0x01 << (i - 1)):
-                key_up = CONTROL["M" + str(i)]
-        m_keys_down = new_m_keys_down
-    # and also MR key
-    elif feature == FEATURE.MR and notification.address == 0x00:
-        new_mr_key_down = struct.unpack("!1B", notification.data[:1])[0]
-        if not mr_key_down and new_mr_key_down:
-            key_down = CONTROL["MR"]
-        if mr_key_down and not new_mr_key_down:
-            key_up = CONTROL["MR"]
-        mr_key_down = new_mr_key_down
-    # keep track of thumb wheel movment
-    elif feature == FEATURE.THUMB_WHEEL and notification.address == 0x00:
-        if notification.data[4] <= 0x01:  # when wheel starts, zero out last movement
-            thumb_wheel_displacement = 0
-        thumb_wheel_displacement += signed(notification.data[0:2])
+    if notification.address == 0x00:
+        if feature == FEATURE.REPROG_CONTROLS_V4:
+            new_keys_down = struct.unpack("!4H", notification.data[:8])
+            for key in new_keys_down:
+                if key and key not in keys_down:
+                    key_down = key
+            for key in keys_down:
+                if key and key not in new_keys_down:
+                    key_up = key
+            keys_down = new_keys_down
+        # and also G keys down
+        elif feature == FEATURE.GKEY:
+            new_g_keys_down = struct.unpack("<I", notification.data[:4])[0]
+            for i in range(32):
+                if new_g_keys_down & (0x01 << i) and not g_keys_down & (0x01 << i):
+                    key_down = CONTROL["G" + str(i + 1)]
+                if g_keys_down & (0x01 << i) and not new_g_keys_down & (0x01 << i):
+                    key_up = CONTROL["G" + str(i + 1)]
+            g_keys_down = new_g_keys_down
+        # and also M keys down
+        elif feature == FEATURE.MKEYS:
+            new_m_keys_down = struct.unpack("!1B", notification.data[:1])[0]
+            for i in range(1, 9):
+                if new_m_keys_down & (0x01 << (i - 1)) and not m_keys_down & (0x01 << (i - 1)):
+                    key_down = CONTROL["M" + str(i)]
+                if m_keys_down & (0x01 << (i - 1)) and not new_m_keys_down & (0x01 << (i - 1)):
+                    key_up = CONTROL["M" + str(i)]
+            m_keys_down = new_m_keys_down
+        # and also MR key
+        elif feature == FEATURE.MR:
+            new_mr_key_down = struct.unpack("!1B", notification.data[:1])[0]
+            if not mr_key_down and new_mr_key_down:
+                key_down = CONTROL["MR"]
+            if mr_key_down and not new_mr_key_down:
+                key_up = CONTROL["MR"]
+            mr_key_down = new_mr_key_down
+        # keep track of thumb wheel movement
+        elif feature == FEATURE.THUMB_WHEEL:
+            if notification.data[4] <= 0x01:  # when wheel starts, zero out last movement
+                thumb_wheel_displacement = 0
+            thumb_wheel_displacement += signed(notification.data[0:2])
 
     GLib.idle_add(evaluate_rules, feature, notification, device)
 
diff --git a/tests/logitech_receiver/test_diversion.py b/tests/logitech_receiver/test_diversion.py
index dcff0fc8..3b574be1 100644
--- a/tests/logitech_receiver/test_diversion.py
+++ b/tests/logitech_receiver/test_diversion.py
@@ -6,6 +6,8 @@ from unittest.mock import mock_open
 import pytest
 
 from logitech_receiver import diversion
+from logitech_receiver.base import HIDPPNotification
+from logitech_receiver.hidpp20_constants import FEATURE
 
 
 @pytest.fixture
@@ -88,3 +90,30 @@ def test_key_is_down():
     result = diversion.key_is_down(key=diversion.CONTROL.G2)
 
     assert result is False
+
+
+@pytest.mark.parametrize(
+    "feature, data",
+    [
+        (
+            FEATURE.REPROG_CONTROLS_V4,
+            [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08],
+        ),
+        (FEATURE.GKEY, [0x01, 0x02, 0x03, 0x04]),
+        (FEATURE.MKEYS, [0x01, 0x02, 0x03, 0x04]),
+        (FEATURE.MR, [0x01, 0x02, 0x03, 0x04]),
+        (FEATURE.THUMB_WHEEL, [0x01, 0x02, 0x03, 0x04, 0x05]),
+        (FEATURE.DEVICE_UNIT_ID, [0x01, 0x02, 0x03, 0x04, 0x05]),
+    ],
+)
+def test_process_notification(feature, data):
+    device_mock = mock.Mock()
+    notification = HIDPPNotification(
+        report_id=0x01,
+        devnumber=1,
+        sub_id=0x13,
+        address=0x00,
+        data=bytes(data),
+    )
+
+    diversion.process_notification(device_mock, notification, feature)