Move HeartbeatController to a separate module (#968)

pychromecast/const.py

@@ -70,3 +70,5 @@
 MESSAGE_TYPE = "type"
 REQUEST_ID = "requestId"
 SESSION_ID = "sessionId"
+
+PLATFORM_DESTINATION_ID = "receiver-0"

pychromecast/controllers/heartbeat.py

@@ -0,0 +1,88 @@
+"""Controller to send and respond to heartbeat messages."""
+
+from __future__ import annotations
+
+import time
+
+from . import BaseController
+
+from ..const import MESSAGE_TYPE, PLATFORM_DESTINATION_ID
+from ..error import ControllerNotRegistered, NotConnected, PyChromecastStopped
+
+# pylint: disable-next=no-name-in-module
+from ..generated.cast_channel_pb2 import CastMessage
+
+NS_HEARTBEAT = "urn:x-cast:com.google.cast.tp.heartbeat"
+
+TYPE_PING = "PING"
+TYPE_PONG = "PONG"
+
+HB_PING_TIME = 10
+HB_PONG_TIME = 10
+
+
+class HeartbeatController(BaseController):
+    """Controller to send and respond to heartbeat messages."""
+
+    def __init__(self) -> None:
+        super().__init__(NS_HEARTBEAT, target_platform=True)
+        self.last_ping = 0.0
+        self.last_pong = time.time()
+
+    def receive_message(self, _message: CastMessage, data: dict) -> bool:
+        """
+        Called when a heartbeat message is received.
+
+        data is message.payload_utf8 interpreted as a JSON dict.
+        """
+        if self._socket_client is None:
+            raise ControllerNotRegistered
+
+        if self._socket_client.is_stopped:
+            return True
+
+        if data[MESSAGE_TYPE] == TYPE_PING:
+            try:
+                self._socket_client.send_message(
+                    PLATFORM_DESTINATION_ID,
+                    self.namespace,
+                    {MESSAGE_TYPE: TYPE_PONG},
+                    no_add_request_id=True,
+                )
+            except PyChromecastStopped:
+                self._socket_client.logger.debug(
+                    "Heartbeat error when sending response, "
+                    "Chromecast connection has stopped"
+                )
+
+            return True
+
+        if data[MESSAGE_TYPE] == TYPE_PONG:
+            self.reset()
+            return True
+
+        return False
+
+    def ping(self) -> None:
+        """Send a ping message."""
+        if self._socket_client is None:
+            raise ControllerNotRegistered
+
+        self.last_ping = time.time()
+        try:
+            self.send_message({MESSAGE_TYPE: TYPE_PING})
+        except NotConnected:
+            self._socket_client.logger.error(
+                "Chromecast is disconnected. Cannot ping until reconnected."
+            )
+
+    def reset(self) -> None:
+        """Reset expired counter."""
+        self.last_pong = time.time()
+
+    def is_expired(self) -> bool:
+        """Indicates if connection has expired."""
+        if time.time() - self.last_ping > HB_PING_TIME:
+            self.ping()
+
+        return (time.time() - self.last_pong) > HB_PING_TIME + HB_PONG_TIME

pychromecast/socket_client.py

@@ -24,8 +24,9 @@
 
 import zeroconf
 
-from .const import MESSAGE_TYPE, REQUEST_ID, SESSION_ID
+from .const import MESSAGE_TYPE, PLATFORM_DESTINATION_ID, REQUEST_ID, SESSION_ID
 from .controllers import BaseController, CallbackType
+from .controllers.heartbeat import NS_HEARTBEAT, HeartbeatController
 from .controllers.media import MediaController
 from .controllers.receiver import CastStatus, CastStatusListener, ReceiverController
 from .dial import get_host_from_service
@@ -42,12 +43,7 @@
 from .models import HostServiceInfo, MDNSServiceInfo
 
 NS_CONNECTION = "urn:x-cast:com.google.cast.tp.connection"
-NS_HEARTBEAT = "urn:x-cast:com.google.cast.tp.heartbeat"
 
-PLATFORM_DESTINATION_ID = "receiver-0"
-
-TYPE_PING = "PING"
-TYPE_PONG = "PONG"
 TYPE_CONNECT = "CONNECT"
 TYPE_CLOSE = "CLOSE"
 TYPE_LOAD = "LOAD"
@@ -65,8 +61,6 @@
 # The socket connection was lost and needs to be retried
 CONNECTION_STATUS_LOST = "LOST"
 
-HB_PING_TIME = 10
-HB_PONG_TIME = 10
 SELECT_TIMEOUT = 5.0
 TIMEOUT_TIME = 30.0
 RETRY_TIME = 5.0
@@ -1065,73 +1059,6 @@ def receive_message(self, message: CastMessage, data: dict) -> bool:
         return False
 
 
-class HeartbeatController(BaseController):
-    """Controller to respond to heartbeat messages."""
-
-    def __init__(self) -> None:
-        super().__init__(NS_HEARTBEAT, target_platform=True)
-        self.last_ping = 0.0
-        self.last_pong = time.time()
-
-    def receive_message(self, _message: CastMessage, data: dict) -> bool:
-        """
-        Called when a heartbeat message is received.
-
-        data is message.payload_utf8 interpreted as a JSON dict.
-        """
-        if self._socket_client is None:
-            raise ControllerNotRegistered
-
-        if self._socket_client.is_stopped:
-            return True
-
-        if data[MESSAGE_TYPE] == TYPE_PING:
-            try:
-                self._socket_client.send_message(
-                    PLATFORM_DESTINATION_ID,
-                    self.namespace,
-                    {MESSAGE_TYPE: TYPE_PONG},
-                    no_add_request_id=True,
-                )
-            except PyChromecastStopped:
-                self._socket_client.logger.debug(
-                    "Heartbeat error when sending response, "
-                    "Chromecast connection has stopped"
-                )
-
-            return True
-
-        if data[MESSAGE_TYPE] == TYPE_PONG:
-            self.reset()
-            return True
-
-        return False
-
-    def ping(self) -> None:
-        """Send a ping message."""
-        if self._socket_client is None:
-            raise ControllerNotRegistered
-
-        self.last_ping = time.time()
-        try:
-            self.send_message({MESSAGE_TYPE: TYPE_PING})
-        except NotConnected:
-            self._socket_client.logger.error(
-                "Chromecast is disconnected. Cannot ping until reconnected."
-            )
-
-    def reset(self) -> None:
-        """Reset expired counter."""
-        self.last_pong = time.time()
-
-    def is_expired(self) -> bool:
-        """Indicates if connection has expired."""
-        if time.time() - self.last_ping > HB_PING_TIME:
-            self.ping()
-
-        return (time.time() - self.last_pong) > HB_PING_TIME + HB_PONG_TIME
-
-
 def new_socket() -> socket.socket:
     """
     Create a new socket with OS-specific parameters