connection closed/opend async callbacks

pyvlx/connection.py

@@ -76,13 +76,14 @@ def connection_lost(self, exc: object) -> None:
 class Connection:
     """Class for handling TCP connection."""
 
-    def __init__(self, loop: asyncio.AbstractEventLoop, config: Config, connection_closed_cb: Callable[[], None] = None):
+    def __init__(self, loop: asyncio.AbstractEventLoop, config: Config):
         """Init TCP connection."""
         self.loop = loop
         self.config = config
         self.transport: Optional[asyncio.Transport] = None
         self.frame_received_cbs: List[CallbackType] = []
-        self.connection_closed_cb: Callable[[], None] = connection_closed_cb
+        self.connection_closed_cbs: List[Callable[[], Coroutine]] = []
+        self.connection_opened_cbs: List[Callable[[], Coroutine]] = []
         self.connected = False
         self.connection_counter = 0
 
@@ -96,8 +97,9 @@ def disconnect(self) -> None:
             self.transport.close()
             self.transport = None
         self.connected = False
-        if self.connection_closed_cb is not None:
-            self.connection_closed_cb()
+        for connection_closed_cb in self.connection_closed_cbs:
+            # pylint: disable=not-callable
+            self.loop.create_task(connection_closed_cb())
 
     async def connect(self) -> None:
         """Connect to gateway via SSL."""
@@ -114,6 +116,9 @@ async def connect(self) -> None:
         PYVLXLOG.debug(
             "Amount of connections since last HA start: %s", self.connection_counter
         )
+        for connection_opened_cb in self.connection_opened_cbs:
+            # pylint: disable=not-callable
+            await self.loop.create_task(connection_opened_cb())
 
     def register_frame_received_cb(self, callback: CallbackType) -> None:
         """Register frame received callback."""
@@ -123,6 +128,22 @@ def unregister_frame_received_cb(self, callback: CallbackType) -> None:
         """Unregister frame received callback."""
         self.frame_received_cbs.remove(callback)
 
+    def register_connection_closed_cb(self, callback: Callable[[], Coroutine]) -> None:
+        """Register frame received callback."""
+        self.connection_closed_cbs.append(callback)
+
+    def unregister_connection_closed_cb(self, callback: Callable[[], Coroutine]) -> None:
+        """Unregister frame received callback."""
+        self.connection_closed_cbs.remove(callback)
+
+    def register_connection_opened_cb(self, callback: Callable[[], Coroutine]) -> None:
+        """Register frame received callback."""
+        self.connection_opened_cbs.append(callback)
+
+    def unregister_connection_opened_cb(self, callback: Callable[[], Coroutine]) -> None:
+        """Unregister frame received callback."""
+        self.connection_opened_cbs.remove(callback)
+
     def write(self, frame: FrameBase) -> None:
         """Write frame to Bus."""
         if not isinstance(frame, FrameBase):

pyvlx/pyvlx.py

@@ -36,7 +36,8 @@ def __init__(
         """Initialize PyVLX class."""
         self.loop = loop or asyncio.get_event_loop()
         self.config = Config(self, path, host, password)
-        self.connection = Connection(loop=self.loop, config=self.config, connection_closed_cb=self.connection_closed_cb)
+        self.connection = Connection(loop=self.loop, config=self.config)
+        self.connection.register_connection_closed_cb(self.connection_closed_cb)
         self.heartbeat = Heartbeat(
             pyvlx=self,
             interval=heartbeat_interval,
@@ -121,8 +122,8 @@ async def get_limitation(self, node_id: int) -> None:
         limit = get_limitation.GetLimitation(self, node_id)
         await limit.do_api_call()
 
-    def connection_closed_cb(self) -> None:
+    async def connection_closed_cb(self) -> None:
         """Callback when connection to KLF 200 is closed."""
         PYVLXLOG.debug("Connecting to KLF 200 was closed")
         for node in self.nodes:
-            self.loop.create_task(node.after_update())
+            await self.loop.create_task(node.after_update())