From ac984dd3530c6fbdf3231777855401dd698e2688 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sun, 2 Jan 2022 22:38:34 -1000
Subject: [PATCH] Fix strip types on A2 models (#297)

---
 flux_led/models_db.py | 21 ++++++++++++++++++++-
 flux_led/protocol.py  | 30 ++++++++++++++++++++++++++++++
 tests_aio.py          |  2 +-
 3 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/flux_led/models_db.py b/flux_led/models_db.py
index 808cabf5..5547a886 100755
--- a/flux_led/models_db.py
+++ b/flux_led/models_db.py
@@ -23,6 +23,10 @@
     A1_NUM_TO_PROTOCOL,
     A1_OPERATING_MODE_TO_NUM,
     A1_PROTOCOL_TO_NUM,
+    A2_NUM_TO_OPERATING_MODE,
+    A2_NUM_TO_PROTOCOL,
+    A2_OPERATING_MODE_TO_NUM,
+    A2_PROTOCOL_TO_NUM,
     ADDRESSABLE_RGB_NUM_TO_WIRING,
     ADDRESSABLE_RGB_WIRING_TO_NUM,
     NEW_ADDRESSABLE_NUM_TO_OPERATING_MODE,
@@ -279,6 +283,21 @@ class LEDENETDeviceConfigurationOptions:
     num_to_ic_type=A1_NUM_TO_PROTOCOL,
     ic_type_to_num=A1_PROTOCOL_TO_NUM,
 )
+A2_DEVICE_CONFIG = LEDENETDeviceConfigurationOptions(
+    wiring=True,
+    num_to_wiring=ADDRESSABLE_RGB_NUM_TO_WIRING,
+    wiring_to_num=ADDRESSABLE_RGB_WIRING_TO_NUM,
+    operating_modes=False,
+    num_to_operating_mode=A2_NUM_TO_OPERATING_MODE,
+    operating_mode_to_num=A2_OPERATING_MODE_TO_NUM,
+    pixels=True,
+    segments=True,
+    music_pixels=True,
+    music_segments=True,
+    ic_type=True,
+    num_to_ic_type=A2_NUM_TO_PROTOCOL,
+    ic_type_to_num=A2_PROTOCOL_TO_NUM,
+)
 NEW_ADDRESABLE_DEVICE_CONFIG = LEDENETDeviceConfigurationOptions(
     wiring=True,
     num_to_wiring=ADDRESSABLE_RGB_NUM_TO_WIRING,
@@ -1083,7 +1102,7 @@ def protocol_for_version_num(self, version_num: int) -> str:
         color_modes=COLOR_MODES_ADDRESSABLE,
         channel_map={},
         microphone=True,
-        device_config=NEW_ADDRESABLE_DEVICE_CONFIG,
+        device_config=A2_DEVICE_CONFIG,
     ),
     LEDENETModel(
         model_num=0xA3,
diff --git a/flux_led/protocol.py b/flux_led/protocol.py
index 791ba42d..b0ca7609 100755
--- a/flux_led/protocol.py
+++ b/flux_led/protocol.py
@@ -325,6 +325,36 @@ class LEDENETRawState(NamedTuple):
 }
 A1_OPERATING_MODE_TO_NUM = {v: k for k, v in A1_NUM_TO_OPERATING_MODE.items()}
 
+A2_NUM_TO_PROTOCOL = {
+    1: "UCS1903",
+    2: "SM16703",
+    3: "WS2811",
+    4: "WS2811B",
+    5: "SK6812",
+    6: "INK1003",
+    7: "WS2801",
+    8: "WS2815",
+    9: "APA102",
+    10: "TM1914",
+    11: "UCS2904B",
+}
+A2_PROTOCOL_TO_NUM = {v: k for k, v in A2_NUM_TO_PROTOCOL.items()}
+
+A2_NUM_TO_OPERATING_MODE = {
+    1: COLOR_MODE_RGB,
+    2: COLOR_MODE_RGB,
+    3: COLOR_MODE_RGB,
+    4: COLOR_MODE_RGB,
+    5: COLOR_MODE_RGB,
+    6: COLOR_MODE_RGB,
+    7: COLOR_MODE_RGB,
+    8: COLOR_MODE_RGB,
+    9: COLOR_MODE_RGB,
+    10: COLOR_MODE_RGB,
+    11: COLOR_MODE_RGB,
+}
+A2_OPERATING_MODE_TO_NUM = {v: k for k, v in A2_NUM_TO_OPERATING_MODE.items()}
+
 NEW_ADDRESSABLE_NUM_TO_PROTOCOL = {
     1: "WS2812B",
     2: "SM16703",
diff --git a/tests_aio.py b/tests_aio.py
index 518c2e76..ce4c8bb7 100644
--- a/tests_aio.py
+++ b/tests_aio.py
@@ -929,7 +929,7 @@ def _updated_callback(*args, **kwargs):
         )
     assert len(transport.mock_calls) == 1
     assert transport.mock_calls[0][0] == "write"
-    assert transport.mock_calls[0][1][0] == b"b\x01,\x00\x02\x06\x02\x19\x02\xf0\xa4"
+    assert transport.mock_calls[0][1][0] == b"b\x01,\x00\x02\x05\x02\x19\x02\xf0\xa3"
 
     transport.reset_mock()
     with patch.object(light, "_async_device_config_resync", mock_coro):