From fbe30270584c2554e68fd59dc21bdd925ffd7b20 Mon Sep 17 00:00:00 2001
From: Benjamin Braatz <bb@bbraatz.eu>
Date: Mon, 31 Jul 2023 04:01:32 +0200
Subject: [PATCH] Another round of improvement.

---
 controlpi_plugins/wsclient.py | 36 +++++++++++++++++++----------------
 setup.py                      |  2 +-
 2 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/controlpi_plugins/wsclient.py b/controlpi_plugins/wsclient.py
index fd81680..686e3ee 100644
--- a/controlpi_plugins/wsclient.py
+++ b/controlpi_plugins/wsclient.py
@@ -196,14 +196,16 @@ class WSClient(BasePlugin):
         """Connect to wsserver and process messages from it."""
         async for websocket in connect(self.conf['url'],
                                        ping_interval=1, ping_timeout=5):
-            print(f"WSClient '{self.name}': Connection opened"
-                  f" to {self.conf['url']}.")
-            await self.bus.send(Message(self.name,
-                                        {'event':
-                                         'connection opened'}))
             try:
+                # Open connection:
+                print(f"WSClient '{self.name}': Connection opened"
+                      f" to {self.conf['url']}.")
+                await self.bus.send(Message(self.name,
+                                            {'event':
+                                             'connection opened'}))
+                # Configure name, MAC, and filter:
                 conf_command: Dict[str, Any] = \
-                        {'command': 'configure websocket', 'target': ''}
+                    {'command': 'configure websocket', 'target': ''}
                 if 'client' in self.conf:
                     conf_command['name'] = self._client
                 else:
@@ -216,22 +218,24 @@ class WSClient(BasePlugin):
                 conf_command['down filter'] = self._down_filter
                 json_command = json.dumps(conf_command)
                 await websocket.send(json_command)
+                print(f"WSClient '{self.name}': Connection configured"
+                      f" to {self.conf['url']} as '{conf_command['name']}'.")
                 await self.bus.send(Message(self.name,
                                             {'event':
                                              'connection configured'}))
+                # Read incoming messages in loop:
                 self._websocket = websocket
-                try:
-                    async for json_message in websocket:
-                        assert isinstance(json_message, str)
-                        await self._send(json_message)
-                except ConnectionClosed:
-                    self._websocket = None
-                    continue
-                await self.bus.send(Message(self.name,
-                                            {'event':
-                                             'connection closed'}))
+                async for json_message in websocket:
+                    assert isinstance(json_message, str)
+                    await self._send(json_message)
             except Exception as e:
                 print(f"WSClient '{self.name}': Exception in connection: {e}")
+            self._websocket = None
+            print(f"WSClient '{self.name}': Connection closed"
+                  f" to {self.conf['url']}.")
+            await self.bus.send(Message(self.name,
+                                        {'event':
+                                         'connection closed'}))
             await asyncio.sleep(1)
 
     async def _send(self, json_message: str) -> None:
diff --git a/setup.py b/setup.py
index 8d69102..3f6cfef 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@ with open("README.md", "r") as readme_file:
 
 setuptools.setup(
     name="controlpi-wsclient",
-    version="0.3.2",
+    version="0.3.3",
     author="Graph-IT GmbH",
     author_email="info@graph-it.com",
     description="ControlPi Plugin for Websocket Clients",
-- 
2.34.1