From 0be18f7703c76b032941417f9301d7f5867cd0bb Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sat, 8 Apr 2023 09:41:45 -1000 Subject: [PATCH] Fix reconnecting to older devices after they drop wifi for an extended period (#389) --- flux_led/__init__.py | 3 +- flux_led/aiodevice.py | 18 +++++++---- tests.py | 5 --- tests_aio.py | 74 ++++++++++++++++++++++++++----------------- 4 files changed, 59 insertions(+), 41 deletions(-) diff --git a/flux_led/__init__.py b/flux_led/__init__.py index 9134e610..d00e7abc 100644 --- a/flux_led/__init__.py +++ b/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", ] diff --git a/flux_led/aiodevice.py b/flux_led/aiodevice.py index 14b40e43..b0c3dd6f 100644 --- a/flux_led/aiodevice.py +++ b/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 diff --git a/tests.py b/tests.py index 731ebf2f..f1289d64 100644 --- a/tests.py +++ b/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): diff --git a/tests_aio.py b/tests_aio.py index 7b553d0f..90b59fa8 100644 --- a/tests_aio.py +++ b/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 ):