From 97d26e3c65b798438e7e884e0cedf87ca4a7c9ec Mon Sep 17 00:00:00 2001
From: atse <atse@m-labs.hk>
Date: Mon, 9 Sep 2024 11:41:17 +0800
Subject: [PATCH] Update thermostat state from controller code

---
 pytec/pytec/gui/model/thermostat.py          |  7 +--
 pytec/pytec/gui/view/thermostat_ctrl_menu.py |  5 +-
 pytec/tec_qt.py                              | 49 +++++++++++++-------
 3 files changed, 36 insertions(+), 25 deletions(-)

diff --git a/pytec/pytec/gui/model/thermostat.py b/pytec/pytec/gui/model/thermostat.py
index deae06d..26f3aca 100644
--- a/pytec/pytec/gui/model/thermostat.py
+++ b/pytec/pytec/gui/model/thermostat.py
@@ -40,14 +40,9 @@ class Thermostat(QObject, metaclass=PropertyMeta):
         self.connection_state = ThermostatConnectionState.DISCONNECTED
 
     async def start_session(self, host, port):
-        self.connection_state = ThermostatConnectionState.CONNECTING
-
         await self._client.connect(host, port)
         self.hw_rev = await self._client.hw_rev()
 
-        self.connection_state = ThermostatConnectionState.CONNECTED
-        self.start_watching()
-
     async def run(self):
         self._update_params_task = asyncio.create_task(self.update_params())
         while True:
@@ -60,6 +55,7 @@ class Thermostat(QObject, metaclass=PropertyMeta):
                         exc_info=True,
                     )
                     await self.end_session()
+                    self.connection_state = ThermostatConnectionState.DISCONNECTED
                     self.connection_error.emit()
                     return
                 self._update_params_task = asyncio.create_task(self.update_params())
@@ -131,7 +127,6 @@ class Thermostat(QObject, metaclass=PropertyMeta):
                 self.disconnect_cb()
 
         await self._client.disconnect()
-        self.connection_state = ThermostatConnectionState.DISCONNECTED
 
     async def set_ipv4(self, ipv4):
         await self._client.set_param("ipv4", ipv4)
diff --git a/pytec/pytec/gui/view/thermostat_ctrl_menu.py b/pytec/pytec/gui/view/thermostat_ctrl_menu.py
index 78b88bc..9cdf8f2 100644
--- a/pytec/pytec/gui/view/thermostat_ctrl_menu.py
+++ b/pytec/pytec/gui/view/thermostat_ctrl_menu.py
@@ -188,6 +188,7 @@ class ThermostatCtrlMenu(QtWidgets.QMenu):
 
         await self._thermostat.reset()
         await self._thermostat.end_session()
+        self._thermostat.connection_state = ThermostatConnectionState.DISCONNECTED
 
     @asyncSlot(bool)
     async def dfu_request(self, _):
@@ -195,6 +196,7 @@ class ThermostatCtrlMenu(QtWidgets.QMenu):
 
         await self._thermostat.dfu()
         await self._thermostat.end_session()
+        self._thermostat.connection_state = ThermostatConnectionState.DISCONNECTED
 
     @asyncSlot(bool)
     async def net_settings_request(self, _):
@@ -209,4 +211,5 @@ class ThermostatCtrlMenu(QtWidgets.QMenu):
         assert self._thermostat.connected()
 
         await self._thermostat.set_ipv4(ipv4_settings)
-        await self._thermostat.end_session()
\ No newline at end of file
+        await self._thermostat.end_session()
+        self._thermostat.connection_state = ThermostatConnectionState.DISCONNECTED
\ No newline at end of file
diff --git a/pytec/tec_qt.py b/pytec/tec_qt.py
index 2d2fc6f..e5267f0 100755
--- a/pytec/tec_qt.py
+++ b/pytec/tec_qt.py
@@ -128,6 +128,7 @@ class MainWindow(QtWidgets.QMainWindow):
     async def closeEvent(self, _event):
         try:
             await self.thermostat.end_session()
+            self.thermostat.connection_state = ThermostatConnectionState.DISCONNECTED
         except:
             pass
 
@@ -181,27 +182,39 @@ class MainWindow(QtWidgets.QMainWindow):
 
     @asyncSlot()
     async def on_connect_btn_clicked(self):
-        if (self._connecting_task is None) and (not self.thermostat.connected()):
-            self._connecting_task = asyncio.create_task(
-                self.thermostat.start_session(
-                    host=self.conn_menu.host_set_line.text(),
-                    port=self.conn_menu.port_set_spin.value(),
+        match self.thermostat.connection_state:
+            case ThermostatConnectionState.DISCONNECTED:
+                self._connecting_task = asyncio.create_task(
+                    self.thermostat.start_session(
+                        host=self.conn_menu.host_set_line.text(),
+                        port=self.conn_menu.port_set_spin.value(),
+                    )
+                )
+                self.thermostat.connection_state = ThermostatConnectionState.CONNECTING
+                try:
+                    await self._connecting_task
+                except (OSError, asyncio.CancelledError) as exc:
+                    await self.thermostat.end_session()
+                    if isinstance(exc, asyncio.CancelledError):
+                        return
+                    raise
+                finally:
+                    self._connecting_task = None
+                self.thermostat.connection_state = ThermostatConnectionState.CONNECTED
+                self.thermostat.start_watching()
+
+            case ThermostatConnectionState.CONNECTING:
+                self._connecting_task.cancel()
+                self.thermostat.connection_state = (
+                    ThermostatConnectionState.DISCONNECTED
                 )
-            )
-            try:
-                await self._connecting_task
-            except (OSError, asyncio.CancelledError) as exc:
-                await self.thermostat.end_session()
-                if isinstance(exc, asyncio.CancelledError):
-                    return
-                raise
-            finally:
                 self._connecting_task = None
 
-        elif self._connecting_task is not None:
-            self._connecting_task.cancel()
-        else:
-            await self.thermostat.end_session()
+            case ThermostatConnectionState.CONNECTED:
+                await self.thermostat.end_session()
+                self.thermostat.connection_state = (
+                    ThermostatConnectionState.DISCONNECTED
+                )
 
     @asyncSlot(int)
     async def on_report_box_stateChanged(self, enabled):