From fcbbba661d8635d8da34600156b9ff4b3644b59d Mon Sep 17 00:00:00 2001 From: epenet Date: Mon, 7 Mar 2022 12:29:03 +0100 Subject: [PATCH] Fix issue with extra events in open routine --- samsungtvws/async_connection.py | 15 +++++++++++---- samsungtvws/connection.py | 15 +++++++++++---- tests/test_async_remote.py | 12 ++++++++++++ tests/test_remote.py | 11 +++++++++++ 4 files changed, 45 insertions(+), 8 deletions(-) diff --git a/samsungtvws/async_connection.py b/samsungtvws/async_connection.py index 98722cf..462b3ef 100644 --- a/samsungtvws/async_connection.py +++ b/samsungtvws/async_connection.py @@ -32,7 +32,7 @@ from . import connection, exceptions, helper from .command import SamsungTVCommand, SamsungTVSleepCommand -from .event import MS_CHANNEL_CONNECT_EVENT +from .event import ED_EDENTV_UPDATE_EVENT, MS_CHANNEL_CONNECT_EVENT _LOGGING = logging.getLogger(__name__) @@ -68,13 +68,20 @@ async def open(self) -> WebSocketClientProtocol: connect_kwargs["ssl"] = ssl_context connection = await connect(url, open_timeout=self.timeout, **connect_kwargs) - response = helper.process_api_response(await connection.recv()) - self._check_for_token(response) + event: Optional[str] = None + while event is None or event == ED_EDENTV_UPDATE_EVENT: + data = await connection.recv() + response = helper.process_api_response(data) + event = response.get("event", "*") + assert event + self._websocket_event(event, response) - if response["event"] != MS_CHANNEL_CONNECT_EVENT: + if event != MS_CHANNEL_CONNECT_EVENT: await self.close() raise exceptions.ConnectionFailure(response) + self._check_for_token(response) + self.connection = connection return connection diff --git a/samsungtvws/connection.py b/samsungtvws/connection.py index 630a0b5..3cd3349 100644 --- a/samsungtvws/connection.py +++ b/samsungtvws/connection.py @@ -32,7 +32,7 @@ from . import exceptions, helper from .command import SamsungTVCommand, SamsungTVSleepCommand -from .event import MS_CHANNEL_CONNECT_EVENT, MS_ERROR_EVENT +from .event import ED_EDENTV_UPDATE_EVENT, MS_CHANNEL_CONNECT_EVENT, MS_ERROR_EVENT _LOGGING = logging.getLogger(__name__) @@ -163,13 +163,20 @@ def open(self) -> websocket.WebSocket: connection="Connection: Upgrade", ) - response = helper.process_api_response(connection.recv()) - self._check_for_token(response) + event: Optional[str] = None + while event is None or event == ED_EDENTV_UPDATE_EVENT: + data = connection.recv() + response = helper.process_api_response(data) + event = response.get("event", "*") + assert event + self._websocket_event(event, response) - if response["event"] != MS_CHANNEL_CONNECT_EVENT: + if event != MS_CHANNEL_CONNECT_EVENT: self.close() raise exceptions.ConnectionFailure(response) + self._check_for_token(response) + self.connection = connection return connection diff --git a/tests/test_async_remote.py b/tests/test_async_remote.py index 3c92f99..250fa89 100644 --- a/tests/test_async_remote.py +++ b/tests/test_async_remote.py @@ -41,6 +41,18 @@ async def test_connect(async_connection: Mock) -> None: assert tv.token == 123456789 +@pytest.mark.asyncio +async def test_connect_with_extra_event(async_connection: Mock) -> None: + """Ensure simple data can be parsed.""" + async_connection.recv = Mock( + side_effect=[ED_EDENTV_UPDATE_FUTURE, MS_CHANNEL_CONNECT_FUTURE] + ) + async_connection.send = Mock(return_value=NONE_FUTURE) + tv = SamsungTVWSAsyncRemote("127.0.0.1") + await tv.start_listening() + assert tv.token == 123456789 + + @pytest.mark.asyncio async def test_connection_failure(async_connection: Mock) -> None: """Ensure simple data can be parsed.""" diff --git a/tests/test_remote.py b/tests/test_remote.py index 03dd1c5..efc8b8c 100644 --- a/tests/test_remote.py +++ b/tests/test_remote.py @@ -8,6 +8,7 @@ from .const import ( ED_APPS_LAUNCH_SAMPLE, + ED_EDENTV_UPDATE_SAMPLE, ED_INSTALLED_APP_SAMPLE, MS_CHANNEL_CONNECT_SAMPLE, MS_ERROR_SAMPLE, @@ -22,6 +23,16 @@ def test_connect(connection: Mock) -> None: assert tv.token == 123456789 +def test_connect_with_extra_event(connection: Mock) -> None: + """Ensure simple data can be parsed.""" + connection.recv = Mock( + side_effect=[ED_EDENTV_UPDATE_SAMPLE, MS_CHANNEL_CONNECT_SAMPLE] + ) + tv = SamsungTVWS("127.0.0.1") + tv.open() + assert tv.token == 123456789 + + def test_connection_failure(connection: Mock) -> None: """Ensure simple data can be parsed.""" connection.recv = Mock(side_effect=[MS_ERROR_SAMPLE])