From 047bde887e80df64e293eee27ff6da3b832af2e8 Mon Sep 17 00:00:00 2001
From: atse <atse@m-labs.hk>
Date: Fri, 5 Jul 2024 17:21:39 +0800
Subject: [PATCH] Integrate WrappedClient into Thermostat model

---
 pytec/pytec/gui/model/thermostat.py | 15 +++------------
 pytec/tec_qt.py                     | 15 ++++++++-------
 2 files changed, 11 insertions(+), 19 deletions(-)

diff --git a/pytec/pytec/gui/model/thermostat.py b/pytec/pytec/gui/model/thermostat.py
index 38a832c..1671603 100644
--- a/pytec/pytec/gui/model/thermostat.py
+++ b/pytec/pytec/gui/model/thermostat.py
@@ -1,4 +1,3 @@
-from pytec.aioclient import Client
 from PyQt6.QtCore import pyqtSignal, QObject, pyqtSlot
 from qasync import asyncSlot
 from pytec.gui.model.property import Property, PropertyMeta
@@ -6,17 +5,6 @@ import asyncio
 import logging
 
 
-class WrappedClient(QObject, Client):
-    connection_error = pyqtSignal()
-
-    async def _read_line(self):
-        try:
-            return await super()._read_line()
-        except (Exception, TimeoutError, asyncio.exceptions.TimeoutError):
-            logging.error("Client connection error, disconnecting", exc_info=True)
-            self.connection_error.emit()
-
-
 class Thermostat(QObject, metaclass=PropertyMeta):
     hw_rev = Property(dict)
     fan = Property(dict)
@@ -27,6 +15,7 @@ class Thermostat(QObject, metaclass=PropertyMeta):
     interval = Property(list)
     report = Property(list)
     info_box_trigger = pyqtSignal(str, str)
+    connection_error = pyqtSignal()
 
     def __init__(self, parent, client, update_s):
         self._update_s = update_s
@@ -49,6 +38,8 @@ class Thermostat(QObject, metaclass=PropertyMeta):
                         "Encountered an error while polling for information from Thermostat.",
                         exc_info=True,
                     )
+                    self.connection_error.emit()
+                    return
                 self._update_params_task = asyncio.create_task(self.update_params())
             await asyncio.sleep(self._update_s)
 
diff --git a/pytec/tec_qt.py b/pytec/tec_qt.py
index 96665af..9163d2e 100755
--- a/pytec/tec_qt.py
+++ b/pytec/tec_qt.py
@@ -7,7 +7,8 @@ from pytec.gui.view.live_plot_view import LiveDataPlotter
 from pytec.gui.view.ctrl_panel import CtrlPanel
 from pytec.gui.view.info_box import InfoBox
 from pytec.gui.model.pid_autotuner import PIDAutoTuner
-from pytec.gui.model.thermostat import WrappedClient, Thermostat
+from pytec.gui.model.thermostat import Thermostat
+from pytec.aioclient import Client
 import json
 from autotune import PIDAutotuneState
 from qasync import asyncSlot, asyncClose
@@ -63,7 +64,11 @@ class MainWindow(QtWidgets.QMainWindow):
         self.hw_rev_data = None
         self.info_box = InfoBox()
 
-        self.client = WrappedClient(self)
+        self.client = Client()
+
+        self.thermostat = Thermostat(
+            self, self.client, self.report_refresh_spin.value()
+        )
 
         def handle_connection_error():
             self.info_box.display_info_box(
@@ -72,11 +77,7 @@ class MainWindow(QtWidgets.QMainWindow):
 
             self.bail()
 
-        self.client.connection_error.connect(handle_connection_error)
-
-        self.thermostat = Thermostat(
-            self, self.client, self.report_refresh_spin.value()
-        )
+        self.thermostat.connection_error.connect(handle_connection_error)
 
         self.client.connection_error.connect(self.thermostat.timed_out)
         self.client.connection_error.connect(self.bail)