From 0eb5295208a0fb4cf46d90c94415ef66b3d8b08f Mon Sep 17 00:00:00 2001 From: ilan <31193909+iloveicedgreentea@users.noreply.github.com> Date: Sun, 1 Sep 2024 22:14:53 -0400 Subject: [PATCH] Bugfixes (#13) * fix: some tests * fix: potential bug with command task * feat: centralize writing * fix: tests * feat: get signal info * fix: clear command flag on connect * fix: update version --- madvr/consts.py | 2 +- madvr/madvr.py | 128 +++++++++++++++------------- requirements-test.txt | 1 + setup.py | 2 +- tests/conftest.py | 53 ++++++++++++ tests/test_MadVR.py | 192 +++++++++++++++++++++++++----------------- 6 files changed, 240 insertions(+), 138 deletions(-) create mode 100644 tests/conftest.py diff --git a/madvr/consts.py b/madvr/consts.py index 7a039dc..65781fc 100644 --- a/madvr/consts.py +++ b/madvr/consts.py @@ -1,6 +1,6 @@ """Constants for madvr module.""" -REFRESH_TIME = 60 +REFRESH_TIME = 20 PING_DELAY = 30 COMMAND_TIMEOUT = 3 PING_INTERVAL = 5 diff --git a/madvr/madvr.py b/madvr/madvr.py index 065ba1d..01a35cf 100644 --- a/madvr/madvr.py +++ b/madvr/madvr.py @@ -141,27 +141,21 @@ async def task_handle_queue(self) -> None: """Handle command queue.""" while True: await self.connection_event.wait() - while ( - not self.command_queue.empty() and not self.stop_commands_flag.is_set() - ): + while not self.command_queue.empty() and not self.stop_commands_flag.is_set(): command = await self.command_queue.get() self.logger.debug("sending queue command %s", command) try: await self.send_command(command) except NotImplementedError as err: self.logger.warning("Command not implemented: %s", err) - except (ConnectionError, ConnectionResetError): - self.logger.warning("Envy was turned off manually") - # update state that its off - await self._handle_power_off() + except (ConnectionError, ConnectionResetError, BrokenPipeError): + self.logger.warning("Task Queue: Envy seems to be disconnected") except AttributeError: self.logger.warning("Issue sending command from queue") except RetryExceededError: self.logger.warning("Retry exceeded for command %s", command) except OSError as err: self.logger.error("Unexpected error when sending command: %s", err) - finally: - self.command_queue.task_done() if self.stop_commands_flag.is_set(): self.clear_queue() @@ -197,9 +191,7 @@ async def task_read_notifications(self) -> None: # try to connect otherwise it will mark the device as offline await self._reconnect() except ConnectionError as e: - self.logger.error( - "Connection error when reading notifications: %s", e - ) + self.logger.error("Connection error when reading notifications: %s", e) continue await asyncio.sleep(TASK_CPU_DELAY) @@ -213,15 +205,7 @@ async def send_heartbeat(self, once: bool = False) -> None: """ async def perform_heartbeat() -> None: - if not self.connected: - self.logger.warning("Connection not established") - raise HeartBeatError("Connection not established") - - async with self.lock: - if self.writer: - self.writer.write(self.HEARTBEAT) - await self.writer.drain() - self.logger.debug("Heartbeat complete") + await self._write_with_timeout(self.HEARTBEAT) async def handle_heartbeat_error( err: TimeoutError | OSError | HeartBeatError, @@ -270,9 +254,7 @@ async def task_ping_until_alive(self) -> None: try: await self.open_connection() except ConnectionError as err: - self.logger.error( - "Error opening connection after connectivity check: %s", err - ) + self.logger.error("Error opening connection after connectivity check: %s", err) else: self.logger.debug( "Device is not connectable, retrying in %s seconds", @@ -285,13 +267,18 @@ async def task_ping_until_alive(self) -> None: await asyncio.sleep(self.ping_interval) async def task_refresh_info(self) -> None: - """Task to refresh some device info every minute""" + """Task to refresh some device info every 20s""" while True: # wait until the connection is established await self.connection_event.wait() cmds = [ ["GetMacAddress"], ["GetTemperatures"], + # get signal info in case a change was missed and its sitting in limbo + ["GetIncomingSignalInfo"], + ["GetOutgoingSignalInfo"], + ["GetAspectRatio"], + ["GetMaskingRatio"], ] for cmd in cmds: await self.add_command_to_queue(cmd) @@ -317,6 +304,32 @@ def stop(self) -> None: self.stop_heartbeat.set() self.stop_commands_flag.set() + async def _write_with_timeout(self, data: bytes) -> None: + """Write data to the socket with a timeout.""" + if not self.connected: + self.logger.error("Connection not established. Reconnecting") + await self._reconnect() + + if not self.writer: + self.logger.error("Writer is not initialized. Reconnecting") + await self._reconnect() + + async def write_and_drain() -> None: + if not self.writer: + raise ConnectionError("Writer is not initialized") + self.writer.write(data) + await self.writer.drain() + + try: + async with self.lock: + await asyncio.wait_for(write_and_drain(), timeout=self.connect_timeout) + except TimeoutError: + self.logger.error("Write operation timed out after %s seconds", self.connect_timeout) + await self._reconnect() + except (ConnectionResetError, OSError) as err: + self.logger.error("Error writing to socket: %s", err) + await self._reconnect() + async def _reconnect(self) -> None: """ Initiate a persistent connection to the device. @@ -325,10 +338,10 @@ async def _reconnect(self) -> None: """ # it will not try to connect until ping is successful if await self.is_device_connectable(): - self.logger.info("Device is online") + self.logger.debug("Device is online") try: - self.logger.info("Connecting to Envy: %s:%s", self.host, self.port) + self.logger.debug("Connecting to Envy: %s:%s", self.host, self.port) # Command client self.reader, self.writer = await asyncio.wait_for( @@ -343,14 +356,14 @@ async def _reconnect(self) -> None: await self._set_connected(True) self.stop_heartbeat.clear() # send a heartbeat now + self.logger.debug("Sending heartbeat") await self.send_heartbeat(once=True) self.logger.info("Connection established") + self.stop_commands_flag.clear() except (TimeoutError, HeartBeatError, OSError) as err: - self.logger.error( - "Heartbeat failed. Connection not established %s", err - ) + self.logger.error("Heartbeat failed. Connection not established %s", err) await self._set_connected(False) raise ConnectionError("Heartbeat failed") from err else: @@ -359,15 +372,22 @@ async def _reconnect(self) -> None: await self._handle_power_off() async def is_device_connectable(self) -> bool: - """Check if the device is connectable without ping.""" - try: - async with asyncio.timeout(SMALL_DELAY): - _, writer = await asyncio.open_connection(self.host, self.port) - writer.close() - await writer.wait_closed() - return True - except (asyncio.TimeoutError, ConnectionRefusedError, OSError): - return False + """Check if the device is connectable without ping. The device is only connectable when on.""" + retry_count = 0 + # loop because upgrading firmware can take a few seconds and will kill the connection + while retry_count < 10: + try: + async with asyncio.timeout(SMALL_DELAY): + _, writer = await asyncio.open_connection(self.host, self.port) + writer.close() + await writer.wait_closed() + return True + except (TimeoutError, ConnectionRefusedError, OSError): + await asyncio.sleep(SMALL_DELAY) + retry_count += 1 + continue + self.logger.debug("Device is not connectable") + return False async def _clear_attr(self) -> None: """ @@ -400,7 +420,7 @@ async def open_connection(self) -> None: self.logger.debug("Connection opened") except (AckError, ConnectionError) as err: self.logger.error("Error opening connection: %s", err) - raise ConnectionError("Error opening connection") from err + raise # once connected, try to refresh data once in the case the device was turned connected to while on already cmds = [ @@ -441,9 +461,7 @@ async def _construct_command(self, raw_command: list[str]) -> tuple[bytes, str]: bytes: the value to send in bytes str: the 'msg' field in the Enum used to filter notifications """ - self.logger.debug( - "raw_command: %s -- raw_command length: %s", raw_command, len(raw_command) - ) + self.logger.debug("raw_command: %s -- raw_command length: %s", raw_command, len(raw_command)) skip_val = False # HA seems to always send commands as a list even if you set them as a str @@ -500,9 +518,7 @@ async def _construct_command(self, raw_command: list[str]) -> tuple[bytes, str]: cmd = command_base + Footer.footer.value except KeyError as exc: - raise NotImplementedError( - "Incorrect parameter given for command" - ) from exc + raise NotImplementedError("Incorrect parameter given for command") from exc else: cmd = command_name + Footer.footer.value @@ -529,15 +545,8 @@ async def send_command(self, command: list) -> None: self.logger.debug("Using values: %s %s", cmd, enum_type) - if not self.connected: - self.logger.error("Connection not established") - raise ConnectionError("Device not connected") - try: - async with self.lock: - if self.writer: - self.writer.write(cmd) - await self.writer.drain() + await self._write_with_timeout(cmd) except (ConnectionResetError, TimeoutError, OSError) as err: self.logger.error("Error writing command to socket: %s", err) raise ConnectionError("Failed to send command") from err @@ -574,8 +583,6 @@ async def power_on(self, mac: str = "") -> None: """ Power on the device """ - # start processing commands - self.stop_commands_flag.clear() # use the detected mac or one that is supplied at init or function call mac_to_use = self.mac_address or self.mac or mac @@ -585,9 +592,7 @@ async def power_on(self, mac: str = "") -> None: send_magic_packet(mac_to_use, logger=self.logger) else: # without wol, you cant power on the device - self.logger.warning( - "No mac provided, no action to take. Implement your own WOL automation" - ) + self.logger.warning("No mac provided, no action to take. Implement your own WOL automation") async def power_off(self, standby: bool = False) -> None: """ @@ -599,6 +604,9 @@ async def power_off(self, standby: bool = False) -> None: # set the flag to delay the ping task to avoid race conditions self.powered_off_recently = True if self.connected: - await self.send_command(["Standby"] if standby else ["PowerOff"]) + try: + await self.send_command(["Standby"] if standby else ["PowerOff"]) + except ConnectionError as err: + self.logger.error("Error sending power off command: %s", err) await self.close_connection() # diff --git a/requirements-test.txt b/requirements-test.txt index 3e560d9..ce2e1df 100644 --- a/requirements-test.txt +++ b/requirements-test.txt @@ -9,3 +9,4 @@ pre-commit mypy ruff pydantic +pre-commit diff --git a/setup.py b/setup.py index 5446a5f..18d17ac 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setuptools.setup( name="py_madvr2", - version="1.6.29", + version="1.6.32", author="iloveicedgreentea2", description="A package to control MadVR Envy over IP", long_description=long_description, diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..15fb82c --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,53 @@ +# type: ignore +from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch + +import pytest + +from madvr.madvr import Madvr + + +@pytest.fixture +def mock_madvr(): + with patch("madvr.madvr.asyncio.open_connection", new_callable=AsyncMock), patch( + "madvr.madvr.Madvr.connected", new_callable=PropertyMock, return_value=True + ): + madvr = Madvr("192.168.1.100") + # ignore mypy + # + madvr.writer = AsyncMock() + madvr.reader = AsyncMock() + madvr._set_connected = AsyncMock() + madvr._clear_attr = AsyncMock() + madvr.is_device_connectable = AsyncMock() + madvr.close_connection = AsyncMock() + madvr._construct_command = AsyncMock() + madvr._write_with_timeout = AsyncMock() + madvr.stop = MagicMock() + madvr.stop_commands_flag = MagicMock() + madvr.stop_heartbeat = MagicMock() + madvr.add_command_to_queue = AsyncMock() + madvr._reconnect = AsyncMock() + madvr._write_with_timeout = AsyncMock() + + # Mock the background tasks to prevent warnings + madvr.task_handle_queue = AsyncMock() + madvr.task_read_notifications = AsyncMock() + # madvr.send_heartbeat = AsyncMock() + madvr.task_ping_until_alive = AsyncMock() + madvr.task_refresh_info = AsyncMock() + yield madvr + + +@pytest.fixture +def mock_send_magic_packet(): + with patch("madvr.madvr.send_magic_packet") as mock: + yield mock + + +@pytest.fixture +def mock_wait_for(): + async def mock_wait_for_func(coro, timeout): + return await coro + + with patch("asyncio.wait_for", mock_wait_for_func): + yield diff --git a/tests/test_MadVR.py b/tests/test_MadVR.py index ee55417..cc0ae50 100644 --- a/tests/test_MadVR.py +++ b/tests/test_MadVR.py @@ -1,78 +1,118 @@ -# test_madvr.py +# type: ignore +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch import pytest -import asyncio -import pytest_asyncio -from madvr.madvr import Madvr - - -@pytest_asyncio.fixture(scope="session") -async def madvr_instance(): - print("Creating MadVR instance") - madvr = Madvr(host="192.168.88.38", port=44077) - await madvr.open_connection() - print("Madvr instance created") - yield madvr - print("Closing MadVR instance") - madvr.stop() - await madvr.close_connection() - madvr.ping_task.cancel() - - -@pytest.mark.asyncio(scope="session") -async def test_open_connection(madvr_instance): - assert madvr_instance.connected() is True - print("Connection opened") - - -# @pytest.mark.asyncio -# async def test_close_connection(madvr_instance): -# await madvr_instance.close_connection() -# assert madvr_instance.connected() is False -# await madvr_instance.open_connection() # Reopen for the next tests - - -@pytest.mark.asyncio(scope="session") -async def test_process_info(madvr_instance): - print("Testing process info") - """Verify the process info func works to assign attrs, with modifications""" - await madvr_instance._process_notifications( - 'Welcome\r\nOk\r\nIncomingSignalInfo 3840x2160 23.976p 2D 422 10bit HDR10 2020 TV 16:9\r\nAspectRatio 129x123 1.78 178 "TV"\r\nOutgoingSignalInfo 3840x2160 23.976p 2D 444 12bit HDR10 2020 TV\r\n' - ) - await madvr_instance._process_notifications( - "IncomingSignalInfo 4096x2160 60p 2D 444 12bit HDR10 2020 TV 16:9\r\n" - ) - assert madvr_instance.msg_dict == { - "incoming_res": "4096x2160", - "incoming_frame_rate": "60p", - "incoming_color_space": "444", - "incoming_bit_depth": "12bit", - "is_signal": True, - "hdr_flag": True, - "incoming_colorimetry": "2020", - "incoming_black_levels": "TV", - "incoming_aspect_ratio": "16:9", - "aspect_res": "129x123", - "aspect_dec": 1.78, - "aspect_int": "178", - "aspect_name": '"TV"', - "outgoing_res": "3840x2160", - "outgoing_frame_rate": "23.976p", - "outgoing_color_space": "444", - "outgoing_bit_depth": "12bit", - "outgoing_hdr_flag": True, - "outgoing_colorimetry": "2020", - "outgoing_black_levels": "TV", - } - - -@pytest.mark.asyncio(scope="session") -async def test_send_command(madvr_instance): - print("Testing send command") - """Verify the send command func works""" - try: - await asyncio.wait_for(madvr_instance.read_notifications(), timeout=5) - except asyncio.TimeoutError: - pass - result = await madvr_instance.send_command(["GetIncomingSignalInfo"]) - assert result == "ok" + +from madvr.errors import HeartBeatError + + +@pytest.mark.asyncio +async def test_init(mock_madvr): + assert mock_madvr.host == "192.168.1.100" + assert mock_madvr.port == 44077 # Assuming DEFAULT_PORT is 44077 + assert isinstance(mock_madvr.command_queue, asyncio.Queue) + + +@pytest.mark.asyncio +async def test_is_on_property(mock_madvr): + mock_madvr.msg_dict = {"is_on": True} + assert mock_madvr.is_on is True + + mock_madvr.msg_dict = {"is_on": False} + assert mock_madvr.is_on is False + + +@pytest.mark.asyncio +async def test_mac_address_property(mock_madvr): + mock_madvr.msg_dict = {"mac_address": "00:11:22:33:44:55"} + assert mock_madvr.mac_address == "00:11:22:33:44:55" + + mock_madvr.msg_dict = {} + assert mock_madvr.mac_address == "" + + +@pytest.mark.asyncio +async def test_set_update_callback(mock_madvr): + callback = MagicMock() + mock_madvr.set_update_callback(callback) + assert mock_madvr.update_callback == callback + + +@pytest.mark.asyncio +async def test_async_add_tasks(mock_madvr): + with patch("asyncio.get_event_loop") as mock_loop: + mock_loop.return_value.create_task = AsyncMock() + await mock_madvr.async_add_tasks() + assert len(mock_madvr.tasks) == 5 # Assuming 5 tasks are created + + +@pytest.mark.asyncio +async def test_send_heartbeat(mock_madvr): + await mock_madvr.send_heartbeat(once=True) + mock_madvr._write_with_timeout.assert_called_once_with(mock_madvr.HEARTBEAT) + + +@pytest.mark.asyncio +async def test_send_heartbeat_error(mock_madvr): + mock_madvr._write_with_timeout = AsyncMock(side_effect=TimeoutError) + with pytest.raises(HeartBeatError): + await mock_madvr.send_heartbeat(once=True) + + +@pytest.mark.asyncio +async def test_open_connection(mock_madvr): + await mock_madvr.open_connection() + + mock_madvr._reconnect.assert_called_once() + assert mock_madvr.add_command_to_queue.call_count == 5 + + +@pytest.mark.asyncio +async def test_open_connection_error(mock_madvr): + mock_madvr._reconnect.side_effect = ConnectionError + + with pytest.raises(ConnectionError): + await mock_madvr.open_connection() + + mock_madvr.add_command_to_queue.assert_not_called() + + +@pytest.mark.asyncio +async def test_power_on(mock_madvr, mock_send_magic_packet): + mock_madvr.msg_dict = {"mac_address": "00:11:22:33:44:55"} + mock_madvr.stop_commands_flag = MagicMock() + mock_madvr.stop_commands_flag.is_set.return_value = False + + await mock_madvr.power_on() + + mock_send_magic_packet.assert_called_once_with("00:11:22:33:44:55", logger=mock_madvr.logger) + + +@pytest.mark.asyncio +async def test_power_off(mock_madvr): + mock_madvr._construct_command.return_value = (b"PowerOff\r", "enum_type") + + await mock_madvr.power_off() + + mock_madvr.stop.assert_called_once() + assert mock_madvr.powered_off_recently is True + mock_madvr._construct_command.assert_called_once_with(["PowerOff"]) + mock_madvr._write_with_timeout.assert_called_once_with(b"PowerOff\r") + mock_madvr.close_connection.assert_called_once() + + +@pytest.mark.asyncio +async def test_power_off_standby(mock_madvr): + mock_madvr._construct_command.return_value = (b"Standby\r", "enum_type") + + await mock_madvr.power_off(standby=True) + + mock_madvr.stop.assert_called_once() + assert mock_madvr.powered_off_recently is True + mock_madvr._construct_command.assert_called_once_with(["Standby"]) + mock_madvr._write_with_timeout.assert_called_once_with(b"Standby\r") + mock_madvr.close_connection.assert_called_once() + + +# Add more tests as needed for other methods and edge cases