Fix reconnecting to older devices after they drop wifi for an extende…

…d period (#389)

flux_led/__init__.py

@@ -1,5 +1,5 @@
 """Init file for Flux LED"""
-from .base_device import DeviceType
+from .base_device import DeviceType, DeviceUnavailableException
 from .device import WifiLedBulb
 from .pattern import PresetPattern
 from .scanner import BulbScanner
@@ -13,4 +13,5 @@
     "WifiLedBulb",
     "BulbScanner",
     "utils",
+    "DeviceUnavailableException",
 ]

flux_led/aiodevice.py

@@ -182,9 +182,14 @@ async def async_stop(self) -> None:
     def _async_stop(self) -> None:
         """Shutdown the connection and mark unavailable."""
         self.set_unavailable("Connection closed")
+        self._async_close()
+        self._last_update_time = NEVER_TIME
+
+    def _async_close(self) -> None:
+        """Close the connection."""
         if self._aio_protocol:
             self._aio_protocol.close()
-        self._last_update_time = NEVER_TIME
+            self._aio_protocol = None
 
     async def _async_send_state_query(self) -> None:
         assert self._protocol is not None
@@ -335,9 +340,11 @@ async def async_update(self, force: bool = False) -> None:
                 # to make sure the device is still responding
                 return
         self._last_update_time = now
-        if self._updates_without_response >= MAX_UPDATES_WITHOUT_RESPONSE:
-            if self._aio_protocol:
-                self._aio_protocol.close()
+        if (
+            self._aio_protocol
+            and self._updates_without_response >= MAX_UPDATES_WITHOUT_RESPONSE
+        ):
+            self._async_close()
             self.set_unavailable(
                 f"device stopped responding after {MAX_UPDATES_WITHOUT_RESPONSE} requests to send state"
             )
@@ -851,8 +858,7 @@ async def _async_determine_protocol(self) -> None:
                     async with asyncio_timeout(self.timeout):
                         await self._determine_protocol_future
                 except asyncio.TimeoutError:
-                    if self._aio_protocol:
-                        self._aio_protocol.close()
+                    self._async_close()
                     continue
                 else:
                     return

tests.py

@@ -435,7 +435,6 @@ def read_data(expected):
     @patch("flux_led.WifiLedBulb._read_msg")
     @patch("flux_led.WifiLedBulb.connect")
     def test_rgbww_controller_version_4(self, mock_connect, mock_read, mock_send):
-
         calls = 0
 
         def read_data(expected):
@@ -633,7 +632,6 @@ def read_data(expected):
     def test_rgbww_controller_version_2_after_factory_reset(
         self, mock_connect, mock_read, mock_send
     ):
-
         calls = 0
 
         def read_data(expected):
@@ -667,7 +665,6 @@ def read_data(expected):
     @patch("flux_led.WifiLedBulb._read_msg")
     @patch("flux_led.WifiLedBulb.connect")
     def test_rgbww_controller_version_9(self, mock_connect, mock_read, mock_send):
-
         calls = 0
 
         def read_data(expected):
@@ -829,7 +826,6 @@ def read_data(expected):
     @patch("flux_led.WifiLedBulb._read_msg")
     @patch("flux_led.WifiLedBulb.connect")
     def test_rgbcw_bulb_v4(self, mock_connect, mock_read, mock_send):
-
         calls = 0
 
         def read_data(expected):
@@ -1354,7 +1350,6 @@ def read_data(expected):
     @patch("flux_led.WifiLedBulb._read_msg")
     @patch("flux_led.WifiLedBulb.connect")
     def test_rgbcw_bulb_v9(self, mock_connect, mock_read, mock_send):
-
         calls = 0
 
         def read_data(expected):

tests_aio.py

@@ -4,11 +4,17 @@
 import json
 import logging
 import time
+import sys
 from unittest.mock import MagicMock, call, patch
 
+try:
+    from unittest.mock import AsyncMock
+except ImportError:
+    from unittest.mock import MagicMock as AsyncMock
+
 import pytest
 
-from flux_led import aiodevice, aioscanner
+from flux_led import aiodevice, aioscanner, DeviceUnavailableException
 from flux_led.aio import AIOWifiLedBulb
 from flux_led.aioprotocol import AIOLEDENETProtocol
 from flux_led.aioscanner import AIOBulbScanner, LEDENETDiscovery
@@ -682,7 +688,7 @@ def _updated_callback(*args, **kwargs):
         pass
 
     task = asyncio.create_task(light.async_setup(_updated_callback))
-    transport, _ = await mock_aio_protocol()
+    transport, original_aio_protocol = await mock_aio_protocol()
     light._aio_protocol.data_received(
         b"\x81\xA3#\x25\x01\x10\x64\x00\x00\x00\x04\x00\xf0\xd5"
     )
@@ -713,8 +719,10 @@ def _updated_callback(*args, **kwargs):
         await light.async_update()
     assert light.available is False
 
-    with pytest.raises(RuntimeError):
-        await light.async_update()
+    # simulate reconnect
+    await light.async_update()
+    assert light._aio_protocol != original_aio_protocol
+    light._aio_protocol = original_aio_protocol
 
     transport.reset_mock()
     light._aio_protocol.data_received(
@@ -725,7 +733,7 @@ def _updated_callback(*args, **kwargs):
     assert transport.mock_calls[0][0] == "write"
     assert (
         transport.mock_calls[0][1][0]
-        == b"\xb0\xb1\xb2\xb3\x00\x01\x01\x05\x00\x04\x81\x8a\x8b\x96\xfd"
+        == b"\xb0\xb1\xb2\xb3\x00\x01\x01\x06\x00\x04\x81\x8a\x8b\x96\xfe"
     )
     assert light.available is True
     light._aio_protocol.data_received(
@@ -1835,8 +1843,8 @@ def _updated_callback(*args, **kwargs):
 
 
 @pytest.mark.asyncio
-async def test_async_set_brightness_rgbww(mock_aio_protocol):
-    """Test we can set brightness rgbww."""
+async def test_async_stop(mock_aio_protocol):
+    """Test we can stop without throwing."""
     light = AIOWifiLedBulb("192.168.1.166")
 
     def _updated_callback(*args, **kwargs):
@@ -1852,6 +1860,22 @@ def _updated_callback(*args, **kwargs):
     await light.async_stop()
     await asyncio.sleep(0)  # make sure nothing throws
 
+
+@pytest.mark.asyncio
+async def test_async_set_brightness_rgbww(mock_aio_protocol):
+    """Test we can set brightness rgbww."""
+    light = AIOWifiLedBulb("192.168.1.166")
+
+    def _updated_callback(*args, **kwargs):
+        pass
+
+    task = asyncio.create_task(light.async_setup(_updated_callback))
+    transport, protocol = await mock_aio_protocol()
+    light._aio_protocol.data_received(
+        b"\x81\x25\x23\x61\x05\x10\xb6\x00\x98\x19\x04\x25\x0f\xde"
+    )
+    await task
+
     transport.reset_mock()
     await light.async_set_brightness(255)
     assert transport.mock_calls[0][0] == "write"
@@ -1878,9 +1902,6 @@ def _updated_callback(*args, **kwargs):
     )
     await task
 
-    await light.async_stop()
-    await asyncio.sleep(0)  # make sure nothing throws
-
     transport.reset_mock()
     await light.async_set_brightness(255)
     assert transport.mock_calls[0][0] == "write"
@@ -1909,9 +1930,6 @@ def _updated_callback(*args, **kwargs):
     )
     await task
 
-    await light.async_stop()
-    await asyncio.sleep(0)  # make sure nothing throws
-
     transport.reset_mock()
     await light.async_set_brightness(255)
     assert transport.mock_calls[0][0] == "write"
@@ -1940,9 +1958,6 @@ def _updated_callback(*args, **kwargs):
     )
     await task
 
-    await light.async_stop()
-    await asyncio.sleep(0)  # make sure nothing throws
-
     transport.reset_mock()
     await light.async_set_brightness(255)
     assert transport.mock_calls[0][0] == "write"
@@ -1971,9 +1986,6 @@ def _updated_callback(*args, **kwargs):
     )
     await task
 
-    await light.async_stop()
-    await asyncio.sleep(0)  # make sure nothing throws
-
     transport.reset_mock()
     await light.async_set_brightness(255)
     assert transport.mock_calls[0][0] == "write"
@@ -2002,9 +2014,6 @@ def _updated_callback(*args, **kwargs):
     )
     await task
 
-    await light.async_stop()
-    await asyncio.sleep(0)  # make sure nothing throws
-
     transport.reset_mock()
     await light.async_set_brightness(255)
     assert transport.mock_calls[0][0] == "write"
@@ -2033,9 +2042,6 @@ def _updated_callback(*args, **kwargs):
     )
     await task
 
-    await light.async_stop()
-    await asyncio.sleep(0)  # make sure nothing throws
-
     transport.reset_mock()
     await light.async_set_brightness(255)
     assert transport.mock_calls[0][0] == "write"
@@ -2162,6 +2168,7 @@ def _updated_callback(*args, **kwargs):
 
 
 @pytest.mark.asyncio
+@pytest.mark.skipif(sys.version_info[:3][1] in (7,), reason="no AsyncMock in 3.7")
 async def test_wrapped_cct_protocol_device(mock_aio_protocol):
     """Test a wrapped cct protocol device."""
     light = AIOWifiLedBulb("192.168.1.166")
@@ -2170,7 +2177,7 @@ def _updated_callback(*args, **kwargs):
         pass
 
     task = asyncio.create_task(light.async_setup(_updated_callback))
-    transport, protocol = await mock_aio_protocol()
+    transport, original_aio_protocol = await mock_aio_protocol()
     light._aio_protocol.data_received(
         b"\x81\x1C\x23\x61\x00\x05\x00\x64\x64\x64\x03\x64\x0F\xC8"
     )
@@ -2273,9 +2280,12 @@ def _updated_callback(*args, **kwargs):
         # First failure should keep the device in
         # a failure state until we get to an update
         # time
-        with pytest.raises(RuntimeError):
+        with patch.object(
+            light, "_async_connect", AsyncMock(side_effect=asyncio.TimeoutError)
+        ), pytest.raises(DeviceUnavailableException):
             await light.async_update()
 
+    light._aio_protocol = original_aio_protocol
     # Should not raise now that bulb has recovered
     light._last_update_time = aiodevice.NEVER_TIME
     light._aio_protocol.data_received(
@@ -2285,6 +2295,7 @@ def _updated_callback(*args, **kwargs):
 
 
 @pytest.mark.asyncio
+@pytest.mark.skipif(sys.version_info[:3][1] in (7,), reason="no AsyncMock in 3.7")
 async def test_cct_protocol_device(mock_aio_protocol):
     """Test a original cct protocol device."""
     light = AIOWifiLedBulb("192.168.1.166")
@@ -2293,7 +2304,7 @@ def _updated_callback(*args, **kwargs):
         pass
 
     task = asyncio.create_task(light.async_setup(_updated_callback))
-    transport, protocol = await mock_aio_protocol()
+    transport, original_aio_protocol = await mock_aio_protocol()
     light._aio_protocol.data_received(
         b"\x81\x09\x23\x61\x00\x05\x00\x64\x64\x64\x03\x64\x0F\xB5"
     )
@@ -2384,9 +2395,13 @@ def _updated_callback(*args, **kwargs):
         # First failure should keep the device in
         # a failure state until we get to an update
         # time
-        with pytest.raises(RuntimeError):
+        with patch.object(
+            light, "_async_connect", AsyncMock(side_effect=asyncio.TimeoutError)
+        ), pytest.raises(DeviceUnavailableException):
             await light.async_update()
 
+    light._aio_protocol = original_aio_protocol
+
     # Should not raise now that bulb has recovered
     light._last_update_time = aiodevice.NEVER_TIME
     light._aio_protocol.data_received(
@@ -3254,6 +3269,7 @@ def _updated_callback(*args, **kwargs):
 
 
 @pytest.mark.asyncio
+@pytest.mark.skipif(sys.version_info[:3][1] in (7,), reason="no AsyncMock in 3.7")
 async def test_async_config_remotes_no_response(
     mock_aio_protocol, caplog: pytest.LogCaptureFixture
 ):