Implement async websocket in samsungtv

homeassistant/components/samsungtv/bridge.py

@@ -2,16 +2,25 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from asyncio import Future
 from asyncio.exceptions import TimeoutError as AsyncioTimeoutError
 import contextlib
-from typing import Any
+from typing import Any, cast
 
+import async_timeout
 from samsungctl import Remote
 from samsungctl.exceptions import AccessDenied, ConnectionClosed, UnhandledResponse
-from samsungtvws import SamsungTVWS
+from samsungtvws.async_connection import SamsungTVWSAsyncConnection
 from samsungtvws.async_rest import SamsungTVAsyncRest
-from samsungtvws.exceptions import ConnectionFailure, HttpApiError
-from websocket import WebSocketException
+from samsungtvws.event import (
+    ED_INSTALLED_APP_EVENT,
+    MS_ERROR_EVENT,
+    parse_installed_app,
+    parse_ms_error,
+)
+from samsungtvws.exceptions import ConnectionFailure, HttpApiError, MessageError
+from samsungtvws.remote import REMOTE_ENDPOINT, ChannelEmitCommand, SendRemoteKey
+from websockets.exceptions import WebSocketException
 
 from homeassistant.const import (
     CONF_HOST,
@@ -298,7 +307,8 @@ def __init__(
         self.token = token
         self._rest_api: SamsungTVAsyncRest | None = None
         self._app_list: dict[str, str] | None = None
-        self._remote: SamsungTVWS | None = None
+        self._remote: SamsungTVWSAsyncConnection | None = None
+        self._app_list_futures: set[Future[dict[str, Any]]] = set()
 
     async def async_mac_from_device(self) -> str | None:
         """Try to fetch the mac address of the TV."""
@@ -307,36 +317,40 @@ async def async_mac_from_device(self) -> str | None:
 
     async def async_get_app_list(self) -> dict[str, str] | None:
         """Get installed app list."""
-        return await self.hass.async_add_executor_job(self._get_app_list)
+        if self._app_list is not None:
+            return self._app_list
 
-    def _get_app_list(self) -> dict[str, str] | None:
-        """Get installed app list."""
-        if self._app_list is None:
-            if remote := self._get_remote():
-                raw_app_list: list[dict[str, str]] = remote.app_list()
-                self._app_list = {
-                    app["name"]: app["appId"]
-                    for app in sorted(raw_app_list, key=lambda app: app["name"])
-                }
+        if not (remote := await self._async_get_remote()):
+            return None
 
+        app_list_future: Future[dict[str, Any]] = Future()
+        self._app_list_futures.add(app_list_future)
+        await remote.send_command(ChannelEmitCommand.get_installed_app())
+
+        try:
+            async with async_timeout.timeout(TIMEOUT_WEBSOCKET):
+                response = await app_list_future
+        except (AsyncioTimeoutError, MessageError) as err:
+            self._app_list = {}
+            LOGGER.debug("Fetching app list failed: %s", err)
+        else:
+            self._app_list = {
+                app["name"]: app["appId"]
+                for app in sorted(
+                    parse_installed_app(response),
+                    key=lambda app: cast(str, app["name"]),
+                )
+            }
+            LOGGER.debug("Generated app list: %s", self._app_list)
         return self._app_list
 
     async def async_is_on(self) -> bool:
         """Tells if the TV is on."""
-        return await self.hass.async_add_executor_job(self._is_on)
-
-    def _is_on(self) -> bool:
-        """Tells if the TV is on."""
-        if self._remote is not None:
-            self._close_remote()
-
-        return self._get_remote() is not None
+        if remote := await self._async_get_remote():
+            return remote.is_alive()
+        return False
 
     async def async_try_connect(self) -> str:
-        """Try to connect to the Websocket TV."""
-        return await self.hass.async_add_executor_job(self._try_connect)
-
-    def _try_connect(self) -> str:
         """Try to connect to the Websocket TV."""
         for self.port in WEBSOCKET_PORTS:
             config = {
@@ -351,14 +365,15 @@ def _try_connect(self) -> str:
             result = None
             try:
                 LOGGER.debug("Try config: %s", config)
-                with SamsungTVWS(
+                async with SamsungTVWSAsyncConnection(
                     host=self.host,
+                    endpoint=REMOTE_ENDPOINT,
                     port=self.port,
                     token=self.token,
-                    timeout=config[CONF_TIMEOUT],
-                    name=config[CONF_NAME],
+                    timeout=TIMEOUT_REQUEST,
+                    name=VALUE_CONF_NAME,
                 ) as remote:
-                    remote.open()
+                    await remote.open()
                     self.token = remote.token
                     LOGGER.debug("Working config: %s", config)
                     return RESULT_SUCCESS
@@ -367,7 +382,7 @@ def _try_connect(self) -> str:
                     "Working but unsupported config: %s, error: %s", config, err
                 )
                 result = RESULT_NOT_SUPPORTED
-            except (OSError, ConnectionFailure) as err:
+            except (OSError, AsyncioTimeoutError, ConnectionFailure) as err:
                 LOGGER.debug("Failing config: %s, error: %s", config, err)
         # pylint: disable=useless-else-on-loop
         else:
@@ -397,10 +412,6 @@ async def async_device_info(self) -> dict[str, Any] | None:
         return None
 
     async def async_send_key(self, key: str, key_type: str | None = None) -> None:
-        """Send the key using websocket protocol."""
-        await self.hass.async_add_executor_job(self._send_key, key, key_type)
-
-    def _send_key(self, key: str, key_type: str | None = None) -> None:
         """Send the key using websocket protocol."""
         if key == "KEY_POWEROFF":
             key = "KEY_POWER"
@@ -409,70 +420,89 @@ def _send_key(self, key: str, key_type: str | None = None) -> None:
             retry_count = 1
             for _ in range(retry_count + 1):
                 try:
-                    if remote := self._get_remote():
+                    if remote := await self._async_get_remote():
                         if key_type == "run_app":
-                            remote.run_app(key)
+                            await remote.send_command(
+                                ChannelEmitCommand.launch_app(key)
+                            )
                         else:
-                            remote.send_key(key)
+                            await remote.send_command(SendRemoteKey.click(key))
                     break
                 except (
                     BrokenPipeError,
                     WebSocketException,
-                ):
+                ) as err:
+                    LOGGER.warning(
+                        "Failed to send key to %s: %s",
+                        self.host,
+                        err,
+                        stack_info=True,
+                    )
                     # BrokenPipe can occur when the commands is sent to fast
                     # WebSocketException can occur when timed out
                     self._remote = None
-        except OSError:
+        except OSError as err:
             # Different reasons, e.g. hostname not resolveable
-            pass
+            LOGGER.warning(
+                "Failed to send key to %s: %s",
+                self.host,
+                err,
+                stack_info=True,
+            )
 
-    def _get_remote(self) -> SamsungTVWS:
+    async def _async_get_remote(self) -> SamsungTVWSAsyncConnection | None:
         """Create or return a remote control instance."""
-        if self._remote is None:
+        if not ((remote := self._remote) and remote.is_alive()):
             # We need to create a new instance to reconnect.
             try:
-                LOGGER.debug(
-                    "Create SamsungTVWSBridge for %s (%s)", CONF_NAME, self.host
-                )
-                self._remote = SamsungTVWS(
+                LOGGER.debug("Creating SamsungTVWSBridge for %s", self.host)
+                assert self.port
+                remote = SamsungTVWSAsyncConnection(
                     host=self.host,
+                    endpoint=REMOTE_ENDPOINT,
                     port=self.port,
                     token=self.token,
                     timeout=TIMEOUT_WEBSOCKET,
                     name=VALUE_CONF_NAME,
                 )
-                self._remote.open()
+                await remote.start_listening(self._websocket_event)
             # This is only happening when the auth was switched to DENY
             # A removed auth will lead to socket timeout because waiting for auth popup is just an open socket
             except ConnectionFailure as err:
-                LOGGER.debug("ConnectionFailure %s", err.__repr__())
+                LOGGER.debug("ConnectionFailure %s", err)
                 self._notify_reauth_callback()
-            except (WebSocketException, OSError) as err:
-                LOGGER.debug("WebSocketException, OSError %s", err.__repr__())
-                self._remote = None
+            except (WebSocketException, AsyncioTimeoutError, OSError) as err:
+                LOGGER.debug("WebSocketException, OSError %s", err)
             else:
-                LOGGER.debug(
-                    "Created SamsungTVWSBridge for %s (%s)", CONF_NAME, self.host
-                )
-                if self.token != self._remote.token:
+                LOGGER.info("Created SamsungTVWSBridge for %s", self.host)
+                self._remote = remote
+                if self.token != remote.token:
                     LOGGER.debug(
                         "SamsungTVWSBridge has provided a new token %s",
-                        self._remote.token,
+                        remote.token,
                     )
-                    self.token = self._remote.token
+                    self.token = remote.token
                     self._notify_new_token_callback()
         return self._remote
 
-    async def async_close_remote(self) -> None:
-        """Close remote object."""
-        await self.hass.async_add_executor_job(self._close_remote)
+    async def _websocket_event(self, event: str, response: dict[str, Any]) -> None:
+        """Handle websocket event."""
+        LOGGER.debug("SamsungTVWS websocket event: %s", response)
+        if event == ED_INSTALLED_APP_EVENT:
+            while self._app_list_futures:
+                self._app_list_futures.pop().set_result(response)
+            return
+        if event == MS_ERROR_EVENT:
+            error = parse_ms_error(response)
+            while self._app_list_futures:
+                self._app_list_futures.pop().set_exception(error)
 
-    def _close_remote(self) -> None:
+    async def async_close_remote(self) -> None:
         """Close remote object."""
         try:
             if self._remote is not None:
                 # Close the current remote connection
-                self._remote.close()
+                await self._remote.close()
             self._remote = None
         except OSError:
             LOGGER.debug("Could not establish connection")

tests/components/samsungtv/conftest.py

@@ -1,14 +1,51 @@
 """Fixtures for Samsung TV."""
+import asyncio
 from datetime import datetime
-from unittest.mock import Mock, patch
+import json
+from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
 from samsungctl import Remote
-from samsungtvws import SamsungTVWS
+from samsungtvws.async_connection import SamsungTVWSAsyncConnection
+from samsungtvws.command import SamsungTVCommand
+from samsungtvws.event import ED_INSTALLED_APP_EVENT
 
 import homeassistant.util.dt as dt_util
 
-from .const import SAMPLE_APP_LIST
+from tests.common import load_fixture
+
+
+class MockAsyncWebsocket:
+    """Class to mock a websockets."""
+
+    def __init__(self):
+        """Initialise MockAsyncWebsocket."""
+        self.app_list_return = json.loads(
+            load_fixture("samsungtv/ed_installedApp_get.json")
+        )
+        self.closed = True
+        self.callback = None
+
+    async def start_listening(self, callback):
+        """Mock successful start_listening."""
+        self.callback = callback
+        self.closed = False
+
+    async def close(self):
+        """Mock successful close."""
+        self.callback = None
+        self.closed = True
+
+    async def send_command(self, command: SamsungTVCommand):
+        """Mock status based on failure mode."""
+        if command.method == "ms.channel.emit":
+            if command.params == {
+                "event": "ed.installedApp.get",
+                "to": "host",
+            }:
+                asyncio.ensure_future(
+                    self.callback(ED_INSTALLED_APP_EVENT, self.app_list_return)
+                )
 
 
 @pytest.fixture(autouse=True)
@@ -52,16 +89,25 @@ def rest_api_fixture() -> Mock:
         yield rest_api_class.return_value
 
 
+@pytest.fixture(name="ws_connection")
+def ws_connection_fixture() -> MockAsyncWebsocket:
+    """Patch the samsungtvws SamsungTVWS."""
+    ws_connection = MockAsyncWebsocket()
+    yield ws_connection
+
+
 @pytest.fixture(name="remotews")
-def remotews_fixture() -> Mock:
+def remotews_fixture(ws_connection: MockAsyncWebsocket) -> Mock:
     """Patch the samsungtvws SamsungTVWS."""
     with patch(
-        "homeassistant.components.samsungtv.bridge.SamsungTVWS"
+        "homeassistant.components.samsungtv.bridge.SamsungTVWSAsyncConnection",
     ) as remotews_class:
-        remotews = Mock(SamsungTVWS)
-        remotews.__enter__ = Mock(return_value=remotews)
-        remotews.__exit__ = Mock()
-        remotews.app_list.return_value = SAMPLE_APP_LIST
+        remotews = Mock(SamsungTVWSAsyncConnection)
+        remotews.__aenter__ = AsyncMock(return_value=remotews)
+        remotews.__aexit__ = AsyncMock()
+        remotews.connection = ws_connection
+        remotews.start_listening.side_effect = ws_connection.start_listening
+        remotews.send_command.side_effect = ws_connection.send_command
         remotews.token = "FAKE_TOKEN"
         remotews_class.return_value = remotews
         yield remotews

tests/components/samsungtv/fixtures/ed_installedApp_get.json

@@ -0,0 +1,22 @@
+{
+    "data": {
+      "data": [
+        {
+          "appId": "111299001912",
+          "app_type": 2,
+          "icon": "/opt/share/webappservice/apps_icon/FirstScreen/111299001912/250x250.png",
+          "is_lock": 0,
+          "name": "YouTube"
+        },
+        {
+          "appId": "3201608010191",
+          "app_type": 2,
+          "icon": "/opt/share/webappservice/apps_icon/FirstScreen/3201608010191/250x250.png",
+          "is_lock": 0,
+          "name": "Deezer"
+        }
+      ]
+    },
+    "event": "ed.installedApp.get",
+    "from": "host"
+  }