Remove duplicate code in SamsungTV bridge (#68839)

homeassistant/components/samsungtv/bridge.py

@@ -6,7 +6,7 @@ import asyncio
 from asyncio.exceptions import TimeoutError as AsyncioTimeoutError
 from collections.abc import Callable, Iterable, Mapping
 import contextlib
-from typing import Any, cast
+from typing import Any, Generic, TypeVar, cast
 
 from samsungctl import Remote
 from samsungctl.exceptions import AccessDenied, ConnectionClosed, UnhandledResponse
@@ -74,6 +74,9 @@ ENCRYPTED_MODEL_USES_POWER = {"JU6400", "JU641D"}
 
 REST_EXCEPTIONS = (HttpApiError, AsyncioTimeoutError, ResponseError)
 
+_TRemote = TypeVar("_TRemote", SamsungTVWSAsyncRemote, SamsungTVEncryptedWSAsyncRemote)
+_TCommand = TypeVar("_TCommand", SamsungTVCommand, SamsungTVEncryptedCommand)
+
 
 def mac_from_device_info(info: dict[str, Any]) -> str | None:
     """Extract the mac address from the device info."""
@@ -193,9 +196,15 @@ class SamsungTVBridge(ABC):
     async def async_send_keys(self, keys: list[str]) -> None:
         """Send a list of keys to the tv."""
 
-    @abstractmethod
     async def async_power_off(self) -> None:
         """Send power off command to remote and close."""
+        await self._async_send_power_off()
+        # Force closing of remote session to provide instant UI feedback
+        await self.async_close_remote()
+
+    @abstractmethod
+    async def _async_send_power_off(self) -> None:
+        """Send power off command."""
 
     @abstractmethod
     async def async_close_remote(self) -> None:
@@ -334,11 +343,9 @@ class SamsungTVLegacyBridge(SamsungTVBridge):
             # Different reasons, e.g. hostname not resolveable
             pass
 
-    async def async_power_off(self) -> None:
+    async def _async_send_power_off(self) -> None:
         """Send power off command to remote."""
         await self.async_send_keys(["KEY_POWEROFF"])
-        # Force closing of remote session to provide instant UI feedback
-        await self.async_close_remote()
 
     async def async_close_remote(self) -> None:
         """Close remote object."""
@@ -355,8 +362,79 @@ class SamsungTVLegacyBridge(SamsungTVBridge):
             LOGGER.debug("Could not establish connection")
 
 
-class SamsungTVWSBridge(SamsungTVBridge):
-    """The Bridge for WebSocket TVs."""
+class SamsungTVWSBaseBridge(SamsungTVBridge, Generic[_TRemote, _TCommand]):
+    """The Bridge for WebSocket TVs (v1/v2)."""
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        method: str,
+        host: str,
+        port: int | None = None,
+    ) -> None:
+        """Initialize Bridge."""
+        super().__init__(hass, method, host, port)
+        self._remote: _TRemote | None = None
+        self._remote_lock = asyncio.Lock()
+
+    async def async_is_on(self) -> bool:
+        """Tells if the TV is on."""
+        LOGGER.debug("Checking if TV %s is on using websocket", self.host)
+        if remote := await self._async_get_remote():
+            return remote.is_alive()  # type: ignore[no-any-return]
+        return False
+
+    async def _async_send_commands(self, commands: list[_TCommand]) -> None:
+        """Send the commands using websocket protocol."""
+        try:
+            # recreate connection if connection was dead
+            retry_count = 1
+            for _ in range(retry_count + 1):
+                try:
+                    if remote := await self._async_get_remote():
+                        await remote.send_commands(commands)
+                    break
+                except (
+                    BrokenPipeError,
+                    WebSocketException,
+                ):
+                    # BrokenPipe can occur when the commands is sent to fast
+                    # WebSocketException can occur when timed out
+                    self._remote = None
+        except OSError:
+            # Different reasons, e.g. hostname not resolveable
+            pass
+
+    async def _async_get_remote(self) -> _TRemote | None:
+        """Create or return a remote control instance."""
+        if (remote := self._remote) and remote.is_alive():
+            # If we have one then try to use it
+            return remote  # type: ignore[no-any-return]
+
+        async with self._remote_lock:
+            # If we don't have one make sure we do it under the lock
+            # so we don't make two do due a race to get the remote
+            return await self._async_get_remote_under_lock()
+
+    @abstractmethod
+    async def _async_get_remote_under_lock(self) -> _TRemote | None:
+        """Create or return a remote control instance."""
+
+    async def async_close_remote(self) -> None:
+        """Close remote object."""
+        try:
+            if self._remote is not None:
+                # Close the current remote connection
+                await self._remote.close()
+            self._remote = None
+        except OSError as err:
+            LOGGER.debug("Error closing connection to %s: %s", self.host, err)
+
+
+class SamsungTVWSBridge(
+    SamsungTVWSBaseBridge[SamsungTVWSAsyncRemote, SamsungTVCommand]
+):
+    """The Bridge for WebSocket TVs (v2)."""
 
     def __init__(
         self,
@@ -372,8 +450,6 @@ class SamsungTVWSBridge(SamsungTVBridge):
             self.token = entry_data.get(CONF_TOKEN)
         self._rest_api: SamsungTVAsyncRest | None = None
         self._device_info: dict[str, Any] | None = None
-        self._remote: SamsungTVWSAsyncRemote | None = None
-        self._remote_lock = asyncio.Lock()
 
     def _get_device_spec(self, key: str) -> Any | None:
         """Check if a flag exists in latest device info."""
@@ -389,10 +465,7 @@ class SamsungTVWSBridge(SamsungTVBridge):
             info = await self.async_device_info()
             return info is not None and info["device"]["PowerState"] == "on"
 
-        LOGGER.debug("Checking if TV %s is on using websocket", self.host)
-        if remote := await self._async_get_remote():
-            return remote.is_alive()
-        return False
+        return await super().async_is_on()
 
     async def async_try_connect(self) -> str:
         """Try to connect to the Websocket TV."""
@@ -477,38 +550,6 @@ class SamsungTVWSBridge(SamsungTVBridge):
         """Send a list of keys using websocket protocol."""
         await self._async_send_commands([SendRemoteKey.click(key) for key in keys])
 
-    async def _async_send_commands(self, commands: list[SamsungTVCommand]) -> None:
-        """Send the commands using websocket protocol."""
-        try:
-            # recreate connection if connection was dead
-            retry_count = 1
-            for _ in range(retry_count + 1):
-                try:
-                    if remote := await self._async_get_remote():
-                        await remote.send_commands(commands)
-                    break
-                except (
-                    BrokenPipeError,
-                    WebSocketException,
-                ):
-                    # BrokenPipe can occur when the commands is sent to fast
-                    # WebSocketException can occur when timed out
-                    self._remote = None
-        except OSError:
-            # Different reasons, e.g. hostname not resolveable
-            pass
-
-    async def _async_get_remote(self) -> SamsungTVWSAsyncRemote | None:
-        """Create or return a remote control instance."""
-        if (remote := self._remote) and remote.is_alive():
-            # If we have one then try to use it
-            return remote
-
-        async with self._remote_lock:
-            # If we don't have one make sure we do it under the lock
-            # so we don't make two do due a race to get the remote
-            return await self._async_get_remote_under_lock()
-
     async def _async_get_remote_under_lock(self) -> SamsungTVWSAsyncRemote | None:
         """Create or return a remote control instance."""
         if self._remote is None or not self._remote.is_alive():
@@ -595,28 +636,18 @@ class SamsungTVWSBridge(SamsungTVBridge):
                     }
                 )
 
-    async def async_power_off(self) -> None:
+    async def _async_send_power_off(self) -> None:
         """Send power off command to remote."""
         if self._get_device_spec("FrameTVSupport") == "true":
             await self._async_send_commands(SendRemoteKey.hold("KEY_POWER", 3))
         else:
             await self._async_send_commands([SendRemoteKey.click("KEY_POWER")])
-        # Force closing of remote session to provide instant UI feedback
-        await self.async_close_remote()
-
-    async def async_close_remote(self) -> None:
-        """Close remote object."""
-        try:
-            if self._remote is not None:
-                # Close the current remote connection
-                await self._remote.close()
-            self._remote = None
-        except OSError as err:
-            LOGGER.debug("Error closing connection to %s: %s", self.host, err)
 
 
-class SamsungTVEncryptedBridge(SamsungTVBridge):
-    """The Bridge for Encrypted WebSocket TVs (J/H models)."""
+class SamsungTVEncryptedBridge(
+    SamsungTVWSBaseBridge[SamsungTVEncryptedWSAsyncRemote, SamsungTVEncryptedCommand]
+):
+    """The Bridge for Encrypted WebSocket TVs (v1 - J/H models)."""
 
     def __init__(
         self,
@@ -640,15 +671,6 @@ class SamsungTVEncryptedBridge(SamsungTVBridge):
 
         self._rest_api_port: int | None = None
         self._device_info: dict[str, Any] | None = None
-        self._remote: SamsungTVEncryptedWSAsyncRemote | None = None
-        self._remote_lock = asyncio.Lock()
-
-    async def async_is_on(self) -> bool:
-        """Tells if the TV is on."""
-        LOGGER.debug("Checking if TV %s is on using websocket", self.host)
-        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."""
@@ -715,40 +737,6 @@ class SamsungTVEncryptedBridge(SamsungTVBridge):
             [SendEncryptedRemoteKey.click(key) for key in keys]
         )
 
-    async def _async_send_commands(
-        self, commands: list[SamsungTVEncryptedCommand]
-    ) -> None:
-        """Send the commands using websocket protocol."""
-        try:
-            # recreate connection if connection was dead
-            retry_count = 1
-            for _ in range(retry_count + 1):
-                try:
-                    if remote := await self._async_get_remote():
-                        await remote.send_commands(commands)
-                    break
-                except (
-                    BrokenPipeError,
-                    WebSocketException,
-                ):
-                    # BrokenPipe can occur when the commands is sent to fast
-                    # WebSocketException can occur when timed out
-                    self._remote = None
-        except OSError:
-            # Different reasons, e.g. hostname not resolveable
-            pass
-
-    async def _async_get_remote(self) -> SamsungTVEncryptedWSAsyncRemote | None:
-        """Create or return a remote control instance."""
-        if (remote := self._remote) and remote.is_alive():
-            # If we have one then try to use it
-            return remote
-
-        async with self._remote_lock:
-            # If we don't have one make sure we do it under the lock
-            # so we don't make two do due a race to get the remote
-            return await self._async_get_remote_under_lock()
-
     async def _async_get_remote_under_lock(
         self,
     ) -> SamsungTVEncryptedWSAsyncRemote | None:
@@ -774,7 +762,7 @@ class SamsungTVEncryptedBridge(SamsungTVBridge):
                 LOGGER.debug("Created SamsungTVEncryptedBridge for %s", self.host)
         return self._remote
 
-    async def async_power_off(self) -> None:
+    async def _async_send_power_off(self) -> None:
         """Send power off command to remote."""
         power_off_commands: list[SamsungTVEncryptedCommand] = []
         if self._short_model in ENCRYPTED_MODEL_USES_POWER_OFF:
@@ -792,15 +780,3 @@ class SamsungTVEncryptedBridge(SamsungTVBridge):
             power_off_commands.append(SendEncryptedRemoteKey.click("KEY_POWEROFF"))
             power_off_commands.append(SendEncryptedRemoteKey.click("KEY_POWER"))
         await self._async_send_commands(power_off_commands)
-        # Force closing of remote session to provide instant UI feedback
-        await self.async_close_remote()
-
-    async def async_close_remote(self) -> None:
-        """Close remote object."""
-        try:
-            if self._remote is not None:
-                # Close the current remote connection
-                await self._remote.close()
-            self._remote = None
-        except OSError as err:
-            LOGGER.debug("Error closing connection to %s: %s", self.host, err)