From 17666f3b5b7143f39aa164b1d189e80ef4b9b4b1 Mon Sep 17 00:00:00 2001
From: Brian Brondel <BBrondel@lsst.org>
Date: Tue, 15 Oct 2024 13:59:47 -0700
Subject: [PATCH] Fix test reliability

---
 python/lsst/ts/atbuilding/csc/building_csc.py    |  7 ++++---
 python/lsst/ts/atbuilding/csc/mock_controller.py | 16 +++++++++++-----
 tests/test_atbuilding_csc.py                     | 14 +++++++++-----
 3 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/python/lsst/ts/atbuilding/csc/building_csc.py b/python/lsst/ts/atbuilding/csc/building_csc.py
index 94ba61f..753736d 100644
--- a/python/lsst/ts/atbuilding/csc/building_csc.py
+++ b/python/lsst/ts/atbuilding/csc/building_csc.py
@@ -310,9 +310,7 @@ async def run_command(self, command: str) -> None:
         command : str
             The command string to send to the server.
         """
-        await asyncio.wait_for(
-            self.client.write_str(command + "\r\n"), timeout=TCP_TIMEOUT
-        )
+        await asyncio.wait_for(self.client.write_str(command), timeout=TCP_TIMEOUT)
 
         # Wait for a response
         command_name = command.split()[0]
@@ -357,6 +355,9 @@ async def listen_for_messages(self):
                     # Response queues provide the response back to the
                     # command that sent them.
                     await self.response_queue[command].put(message_json)
+            except asyncio.IncompleteReadError:
+                # Incomplete read implies disconnect
+                break
             except asyncio.CancelledError:
                 # Cancelled listen_task for disconnect
                 break
diff --git a/python/lsst/ts/atbuilding/csc/mock_controller.py b/python/lsst/ts/atbuilding/csc/mock_controller.py
index 3c3151d..f33edfa 100644
--- a/python/lsst/ts/atbuilding/csc/mock_controller.py
+++ b/python/lsst/ts/atbuilding/csc/mock_controller.py
@@ -71,17 +71,23 @@ def __init__(self, port: int, log: logging.Logger):
         self.TELEMETRY_INTERVAL = 10
 
         super().__init__(
-            port=port, log=log, connect_callback=self.on_connect, terminator=b"\r"
+            port=port,
+            log=log,
+            connect_callback=self.on_connect,
         )
 
     async def respond(self, message: str) -> None:
-        await self.write_str(message + "\r\n")
+        await self.write_str(message)
 
     async def read_and_dispatch(self) -> None:
         """Read, parse and execute a command, and send a response."""
-        data = await self.read_str()
-        data = data.strip()
-        if not data:
+        try:
+            data = await self.read_str()
+            data = data.strip()
+            if not data:
+                return
+        except asyncio.IncompleteReadError:
+            self.log.exception("IncompleteReadError while reading message from CSC")
             return
 
         self.log.debug(f"Received command: {data!r}")
diff --git a/tests/test_atbuilding_csc.py b/tests/test_atbuilding_csc.py
index 3f5da26..f31df9f 100644
--- a/tests/test_atbuilding_csc.py
+++ b/tests/test_atbuilding_csc.py
@@ -18,6 +18,7 @@
 #
 # You should have received a copy of the GNU General Public License
 
+import asyncio
 import unittest
 
 from lsst.ts import salobj
@@ -39,10 +40,11 @@ async def test_open_one_vent(self):
             initial_state=salobj.State.ENABLED, config_dir=None, simulation_mode=1
         ):
             await self.remote.cmd_openVentGate.set_start(gate=[0, -1, -1, -1])
+            await asyncio.sleep(1)
             await self.assert_next_sample(
                 topic=self.remote.evt_ventGateState,
                 state=[VentGateState.OPENED] + [VentGateState.CLOSED] * 3,
-                flush=True,
+                flush=False,
             )
 
     async def test_open_vents(self):
@@ -87,10 +89,11 @@ async def test_close_vents(self):
                 flush=True,
             )
             await self.remote.cmd_closeVentGate.set_start(gate=[0, 1, 2, 3])
+            await asyncio.sleep(1)
             await self.assert_next_sample(
                 topic=self.remote.evt_ventGateState,
                 state=[VentGateState.CLOSED] * 4,
-                flush=True,
+                flush=False,
             )
 
     async def test_reset_extraction_fan_drive(self):
@@ -172,13 +175,13 @@ async def test_drive_fault_code(self):
             await self.assert_next_sample(
                 topic=self.remote.evt_extractionFanDriveFaultCode,
                 state=22,  # default value
-                flush=True,
+                flush=False,
             )
             self.csc.mock_ctrl.fault_codes = [123] * 8
             await self.assert_next_sample(
                 topic=self.remote.evt_extractionFanDriveFaultCode,
                 state=123,
-                flush=True,
+                flush=False,
             )
 
     async def test_drive_state(self):
@@ -187,10 +190,11 @@ async def test_drive_state(self):
             initial_state=salobj.State.ENABLED, config_dir=None, simulation_mode=1
         ):
             self.csc.mock_ctrl.fan_drive_state = FanDriveState.OPERATING
+            await asyncio.sleep(1)
             await self.assert_next_sample(
                 topic=self.remote.evt_extractionFanDriveState,
                 state=FanDriveState.OPERATING,
-                flush=True,
+                flush=False,
             )