From 3277324aca38397b448d7f67b1d5dc9d4cc62ab8 Mon Sep 17 00:00:00 2001
From: Jason Guo <jsg353@cornell.edu>
Date: Tue, 28 Jun 2022 17:46:46 -0400
Subject: [PATCH 01/18] Add starting disconnection and reconnection code (in
 progress)

---
 socs/agent/prologix_interface.py              | 22 +++++++++++++++
 socs/testing/device_emulator.py               |  7 +++++
 .../test_scpi_psu_agent_integration.py        | 27 +++++++++++++++++++
 3 files changed, 56 insertions(+)

diff --git a/socs/agent/prologix_interface.py b/socs/agent/prologix_interface.py
index 68072512c..4aec505d7 100644
--- a/socs/agent/prologix_interface.py
+++ b/socs/agent/prologix_interface.py
@@ -1,5 +1,6 @@
 import time
 import socket
+import select
 
 
 class PrologixInterface:
@@ -21,12 +22,33 @@ def configure(self):
         self.write('++auto 1')
         self.write('++addr ' + str(self.gpibAddr))
 
+    def connection_check(self):
+        try:
+            ready_to_read, ready_to_write, in_error = \
+                select.select([self.sock,], [self.sock,], [], 5)
+        except select.error:
+            self.sock.shutdown(2)    # 0 = done receiving, 1 = done sending, 2 = both
+            self.sock.close()
+            # connection error event here, maybe reconnect
+            print('prologix interface connection error')
+        return ready_to_read, ready_to_write
+
+    def connection_check_read(self):
+        ready_to_read, _ = self.connection_check()
+        assert len(ready_to_read) > 0
+
+    def connection_check_write(self):
+        _, ready_to_write = self.connection_check()
+        assert len(ready_to_write) > 0
+
     def write(self, msg):
+        self.connection_check_write()
         message = msg + '\n'
         self.sock.sendall(message.encode())
         time.sleep(0.1)  # Don't send messages too quickly
 
     def read(self):
+        self.connection_check_read()
         return self.sock.recv(128).decode().strip()
 
     def version(self):
diff --git a/socs/testing/device_emulator.py b/socs/testing/device_emulator.py
index dbd6be86e..bf046cb55 100644
--- a/socs/testing/device_emulator.py
+++ b/socs/testing/device_emulator.py
@@ -192,6 +192,13 @@ def _read_serial(self):
     def __del__(self):
         self.shutdown()
 
+    def disconnect_reconnect(self, timeout, port):
+        print(f"---Disconnecting tcp relay for {timeout} seconds---")
+        self.shutdown()
+        time.sleep(timeout)
+        self.create_tcp_relay(port)
+        print(f"---Re-created tcp relay on port {port}---")
+
     def shutdown(self):
         """Shutdown communication on the configured relay. This will stop any
         attempt to read communication on the relay, as well as shutdown the relay
diff --git a/tests/integration/test_scpi_psu_agent_integration.py b/tests/integration/test_scpi_psu_agent_integration.py
index 41dd39e88..e7a50e7b5 100644
--- a/tests/integration/test_scpi_psu_agent_integration.py
+++ b/tests/integration/test_scpi_psu_agent_integration.py
@@ -1,3 +1,4 @@
+import time
 import pytest
 
 import ocs
@@ -71,6 +72,7 @@ def test_scpi_psu_set_voltage(wait_for_crossbar, gpib_emu, run_agent, client):
 @pytest.mark.integtest
 def test_scpi_psu_monitor_output(wait_for_crossbar, gpib_emu, run_agent, client):
     responses = {
+        "*idn?": "Keithley instruments, 2230G-30-1, 9203269, 1.16-1.04",
         "MEAS:VOLT? CH1": "3.14",
         "MEAS:CURR? CH1": "6.28",
         "MEAS:VOLT? CH2": "2.72",
@@ -84,3 +86,28 @@ def test_scpi_psu_monitor_output(wait_for_crossbar, gpib_emu, run_agent, client)
     resp = client.monitor_output.start(test_mode=True, wait=0)
     resp = client.monitor_output.wait()
     check_resp_success(resp)
+
+
+# @pytest.mark.integtest
+# def test_scpi_psu_monitor_output_disconnect(
+#     wait_for_crossbar, gpib_emu, run_agent, client
+# ):
+#     responses = {
+#         "*idn?": "Keithley instruments, 2230G-30-1, 9203269, 1.16-1.04",
+#         "MEAS:VOLT? CH1": "3.14",
+#         "MEAS:CURR? CH1": "6.28",
+#         "MEAS:VOLT? CH2": "2.72",
+#         "MEAS:CURR? CH2": "5.44",
+#         "MEAS:VOLT? CH3": "1.23",
+#         "MEAS:CURR? CH3": "2.46",
+#     }
+#     gpib_emu.define_responses(responses)
+
+#     client.init()
+#     resp = client.monitor_output.start(test_mode=False, wait=0)
+#     time.sleep(5)
+#     gpib_emu.disconnect_reconnect(timeout=1, port=1234)
+#     time.sleep(10)
+#     client.init()
+#     resp = client.monitor_output.start(test_mode=True)
+#     check_resp_success(resp)

From 0a91921f79bff2fa51ecb047cb5a4f1fb65b3e88 Mon Sep 17 00:00:00 2001
From: Jason Guo <jsg353@cornell.edu>
Date: Wed, 29 Jun 2022 15:05:29 -0400
Subject: [PATCH 02/18] Successful reconnect response on 'graceful'
 socket.close()

---
 socs/agent/prologix_interface.py              | 25 +++++++++-
 socs/testing/device_emulator.py               |  4 +-
 .../test_scpi_psu_agent_integration.py        | 46 +++++++++----------
 3 files changed, 48 insertions(+), 27 deletions(-)

diff --git a/socs/agent/prologix_interface.py b/socs/agent/prologix_interface.py
index 4aec505d7..03a510a3f 100644
--- a/socs/agent/prologix_interface.py
+++ b/socs/agent/prologix_interface.py
@@ -31,6 +31,7 @@ def connection_check(self):
             self.sock.close()
             # connection error event here, maybe reconnect
             print('prologix interface connection error')
+            assert False, "select.error exception"
         return ready_to_read, ready_to_write
 
     def connection_check_read(self):
@@ -44,12 +45,32 @@ def connection_check_write(self):
     def write(self, msg):
         self.connection_check_write()
         message = msg + '\n'
-        self.sock.sendall(message.encode())
+        try:
+            self.sock.sendall(message.encode())
+        except socket.error:
+            print("socket.error exception on socket write -- disconnect!")
+            self.disconnect_handler()
         time.sleep(0.1)  # Don't send messages too quickly
 
     def read(self):
         self.connection_check_read()
-        return self.sock.recv(128).decode().strip()
+        data = self.sock.recv(128)
+        if not data:
+            print("received no data from socket -- disconnect!")
+            self.disconnect_handler()
+        return data.decode().strip()
+
+    def disconnect_handler(self):
+        for i in range(5):
+            try:
+                self.conn_socket()
+                self.configure()
+                print(f"Successfully reconnected on attempt #{i}")
+                return
+            except Exception as e:
+                print(f"Reconnect attempt #{i} failed with: {e}")
+                time.sleep(1)
+        assert False, "Could not reconnect"
 
     def version(self):
         self.write('++ver')
diff --git a/socs/testing/device_emulator.py b/socs/testing/device_emulator.py
index bf046cb55..811dea0b9 100644
--- a/socs/testing/device_emulator.py
+++ b/socs/testing/device_emulator.py
@@ -193,11 +193,11 @@ def __del__(self):
         self.shutdown()
 
     def disconnect_reconnect(self, timeout, port):
-        print(f"---Disconnecting tcp relay for {timeout} seconds---")
+        print(f"<<< Disconnecting tcp relay for {timeout} seconds >>>")
         self.shutdown()
         time.sleep(timeout)
+        print(f"<<< Re-creating tcp relay on port {port} >>>")
         self.create_tcp_relay(port)
-        print(f"---Re-created tcp relay on port {port}---")
 
     def shutdown(self):
         """Shutdown communication on the configured relay. This will stop any
diff --git a/tests/integration/test_scpi_psu_agent_integration.py b/tests/integration/test_scpi_psu_agent_integration.py
index e7a50e7b5..f07c372ef 100644
--- a/tests/integration/test_scpi_psu_agent_integration.py
+++ b/tests/integration/test_scpi_psu_agent_integration.py
@@ -88,26 +88,26 @@ def test_scpi_psu_monitor_output(wait_for_crossbar, gpib_emu, run_agent, client)
     check_resp_success(resp)
 
 
-# @pytest.mark.integtest
-# def test_scpi_psu_monitor_output_disconnect(
-#     wait_for_crossbar, gpib_emu, run_agent, client
-# ):
-#     responses = {
-#         "*idn?": "Keithley instruments, 2230G-30-1, 9203269, 1.16-1.04",
-#         "MEAS:VOLT? CH1": "3.14",
-#         "MEAS:CURR? CH1": "6.28",
-#         "MEAS:VOLT? CH2": "2.72",
-#         "MEAS:CURR? CH2": "5.44",
-#         "MEAS:VOLT? CH3": "1.23",
-#         "MEAS:CURR? CH3": "2.46",
-#     }
-#     gpib_emu.define_responses(responses)
-
-#     client.init()
-#     resp = client.monitor_output.start(test_mode=False, wait=0)
-#     time.sleep(5)
-#     gpib_emu.disconnect_reconnect(timeout=1, port=1234)
-#     time.sleep(10)
-#     client.init()
-#     resp = client.monitor_output.start(test_mode=True)
-#     check_resp_success(resp)
+@pytest.mark.integtest
+def test_scpi_psu_monitor_output_disconnect(
+    wait_for_crossbar, gpib_emu, run_agent, client
+):
+    responses = {
+        "*idn?": "Keithley instruments, 2230G-30-1, 9203269, 1.16-1.04",
+        "MEAS:VOLT? CH1": "3.14",
+        "MEAS:CURR? CH1": "6.28",
+        "MEAS:VOLT? CH2": "2.72",
+        "MEAS:CURR? CH2": "5.44",
+        "MEAS:VOLT? CH3": "1.23",
+        "MEAS:CURR? CH3": "2.46",
+    }
+    gpib_emu.define_responses(responses)
+
+    client.init()
+    resp = client.monitor_output.start(test_mode=False, wait=0)
+    time.sleep(5)
+    gpib_emu.disconnect_reconnect(timeout=3, port=1234)
+    # time.sleep(10)
+    # client.init()
+    # resp = client.monitor_output.start(test_mode=True)
+    # check_resp_success(resp)

From d2b47b75c9f20e99d54776eb62200c8a3f66c091 Mon Sep 17 00:00:00 2001
From: Jason Guo <jsg353@cornell.edu>
Date: Wed, 29 Jun 2022 15:17:45 -0400
Subject: [PATCH 03/18] Clean up exception handling

---
 socs/agent/prologix_interface.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/socs/agent/prologix_interface.py b/socs/agent/prologix_interface.py
index 03a510a3f..0e64b6a33 100644
--- a/socs/agent/prologix_interface.py
+++ b/socs/agent/prologix_interface.py
@@ -47,8 +47,8 @@ def write(self, msg):
         message = msg + '\n'
         try:
             self.sock.sendall(message.encode())
-        except socket.error:
-            print("socket.error exception on socket write -- disconnect!")
+        except socket.error as e:
+            print(f"socket write failed (disconnect?): {e}")
             self.disconnect_handler()
         time.sleep(0.1)  # Don't send messages too quickly
 
@@ -56,7 +56,7 @@ def read(self):
         self.connection_check_read()
         data = self.sock.recv(128)
         if not data:
-            print("received no data from socket -- disconnect!")
+            print("received no data from socket (disconnect?)")
             self.disconnect_handler()
         return data.decode().strip()
 
@@ -67,7 +67,7 @@ def disconnect_handler(self):
                 self.configure()
                 print(f"Successfully reconnected on attempt #{i}")
                 return
-            except Exception as e:
+            except socket.error as e:
                 print(f"Reconnect attempt #{i} failed with: {e}")
                 time.sleep(1)
         assert False, "Could not reconnect"

From f847b863d3f310f73c82c8f29c97f03a2f822d46 Mon Sep 17 00:00:00 2001
From: Jason Guo <jsg353@cornell.edu>
Date: Wed, 29 Jun 2022 16:50:59 -0400
Subject: [PATCH 04/18] Reconsolidate connection checks

---
 socs/agent/prologix_interface.py | 31 ++++++++++++++-----------------
 1 file changed, 14 insertions(+), 17 deletions(-)

diff --git a/socs/agent/prologix_interface.py b/socs/agent/prologix_interface.py
index 0e64b6a33..ecaa27957 100644
--- a/socs/agent/prologix_interface.py
+++ b/socs/agent/prologix_interface.py
@@ -22,41 +22,38 @@ def configure(self):
         self.write('++auto 1')
         self.write('++addr ' + str(self.gpibAddr))
 
-    def connection_check(self):
+    def connection_check(self, op):
+        assert op in ['read', 'write'], "'op' must be 'read' or 'write'"
         try:
             ready_to_read, ready_to_write, in_error = \
                 select.select([self.sock,], [self.sock,], [], 5)
         except select.error:
-            self.sock.shutdown(2)    # 0 = done receiving, 1 = done sending, 2 = both
+            self.sock.shutdown(2)
             self.sock.close()
-            # connection error event here, maybe reconnect
-            print('prologix interface connection error')
-            assert False, "select.error exception"
-        return ready_to_read, ready_to_write
-
-    def connection_check_read(self):
-        ready_to_read, _ = self.connection_check()
-        assert len(ready_to_read) > 0
+            print("Prologix interface connection error")
+            self.disconnect_handler()
+            self.connection_check(op)  # need to test on real hardware
+        if op == 'read':
+            assert len(ready_to_read) > 0
+        elif op == 'write':
+            assert len(ready_to_write) > 0
 
-    def connection_check_write(self):
-        _, ready_to_write = self.connection_check()
-        assert len(ready_to_write) > 0
 
     def write(self, msg):
-        self.connection_check_write()
+        self.connection_check('write')
         message = msg + '\n'
         try:
             self.sock.sendall(message.encode())
         except socket.error as e:
-            print(f"socket write failed (disconnect?): {e}")
+            print(f"Socket write failed (disconnect?): {e}")
             self.disconnect_handler()
         time.sleep(0.1)  # Don't send messages too quickly
 
     def read(self):
-        self.connection_check_read()
+        self.connection_check('read')
         data = self.sock.recv(128)
         if not data:
-            print("received no data from socket (disconnect?)")
+            print("Received no data from socket (disconnect?)")
             self.disconnect_handler()
         return data.decode().strip()
 

From d8b6014b8e36d2291db06c6ecb629cf773572f88 Mon Sep 17 00:00:00 2001
From: Jason Guo <jsg353@cornell.edu>
Date: Wed, 29 Jun 2022 16:52:54 -0400
Subject: [PATCH 05/18] Should return on recursive call

---
 socs/agent/prologix_interface.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/socs/agent/prologix_interface.py b/socs/agent/prologix_interface.py
index ecaa27957..15713d4be 100644
--- a/socs/agent/prologix_interface.py
+++ b/socs/agent/prologix_interface.py
@@ -32,7 +32,7 @@ def connection_check(self, op):
             self.sock.close()
             print("Prologix interface connection error")
             self.disconnect_handler()
-            self.connection_check(op)  # need to test on real hardware
+            return self.connection_check(op)  # need to test on real hardware
         if op == 'read':
             assert len(ready_to_read) > 0
         elif op == 'write':

From 779668d07f0808123a9f341909d6ab3a28e118f0 Mon Sep 17 00:00:00 2001
From: Jason Guo <jsg353@cornell.edu>
Date: Thu, 30 Jun 2022 15:29:08 -0400
Subject: [PATCH 06/18] Successful monitor_output recovery due to gracefully
 closed socket -- need to test behavior on hardware

---
 agents/scpi_psu/scpi_psu_agent.py                    | 12 +++++++++---
 socs/agent/prologix_interface.py                     | 10 +++++++---
 tests/integration/test_scpi_psu_agent_integration.py |  2 +-
 3 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/agents/scpi_psu/scpi_psu_agent.py b/agents/scpi_psu/scpi_psu_agent.py
index f7bc74400..e6692df20 100644
--- a/agents/scpi_psu/scpi_psu_agent.py
+++ b/agents/scpi_psu/scpi_psu_agent.py
@@ -79,9 +79,15 @@ def monitor_output(self, session, params=None):
                         'data': {}
                     }
 
-                    for chan in channels:
-                        data['data']["Voltage_{}".format(chan)] = self.psu.get_volt(chan)
-                        data['data']["Current_{}".format(chan)] = self.psu.get_curr(chan)
+                    try:
+                        for chan in channels:
+                            data['data']["Voltage_{}".format(chan)] = self.psu.get_volt(chan)
+                            data['data']["Current_{}".format(chan)] = self.psu.get_curr(chan)
+                    except AssertionError as e:
+                        print(
+                            f"Failed to get volt/curr in monitor_output, skipping this iteration: {e}"
+                        )
+                        continue
 
                     # self.log.info(str(data))
                     # print(data)
diff --git a/socs/agent/prologix_interface.py b/socs/agent/prologix_interface.py
index 15713d4be..c5368f727 100644
--- a/socs/agent/prologix_interface.py
+++ b/socs/agent/prologix_interface.py
@@ -32,11 +32,12 @@ def connection_check(self, op):
             self.sock.close()
             print("Prologix interface connection error")
             self.disconnect_handler()
-            return self.connection_check(op)  # need to test on real hardware
+            self.connection_check(op)  # need to test on real hardware
+            return
         if op == 'read':
-            assert len(ready_to_read) > 0
+            assert len(ready_to_read) > 0, "No sockets ready for reading"
         elif op == 'write':
-            assert len(ready_to_write) > 0
+            assert len(ready_to_write) > 0, "No sockets ready for writing"
 
 
     def write(self, msg):
@@ -47,6 +48,8 @@ def write(self, msg):
         except socket.error as e:
             print(f"Socket write failed (disconnect?): {e}")
             self.disconnect_handler()
+            self.write(msg)
+            return
         time.sleep(0.1)  # Don't send messages too quickly
 
     def read(self):
@@ -55,6 +58,7 @@ def read(self):
         if not data:
             print("Received no data from socket (disconnect?)")
             self.disconnect_handler()
+            return self.read()
         return data.decode().strip()
 
     def disconnect_handler(self):
diff --git a/tests/integration/test_scpi_psu_agent_integration.py b/tests/integration/test_scpi_psu_agent_integration.py
index f07c372ef..1ad084a85 100644
--- a/tests/integration/test_scpi_psu_agent_integration.py
+++ b/tests/integration/test_scpi_psu_agent_integration.py
@@ -107,7 +107,7 @@ def test_scpi_psu_monitor_output_disconnect(
     resp = client.monitor_output.start(test_mode=False, wait=0)
     time.sleep(5)
     gpib_emu.disconnect_reconnect(timeout=3, port=1234)
-    # time.sleep(10)
+    time.sleep(10)
     # client.init()
     # resp = client.monitor_output.start(test_mode=True)
     # check_resp_success(resp)

From e49ba0573a6449f9572c2f15fe355ab7fbc98f09 Mon Sep 17 00:00:00 2001
From: Jason Guo <jsg353@cornell.edu>
Date: Thu, 30 Jun 2022 15:56:55 -0400
Subject: [PATCH 07/18] Change behavior to always raise exception if connection
 recovered mid-read, since that read attempt should not be satisfiable

---
 agents/scpi_psu/scpi_psu_agent.py | 2 +-
 socs/agent/prologix_interface.py  | 9 ++++++++-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/agents/scpi_psu/scpi_psu_agent.py b/agents/scpi_psu/scpi_psu_agent.py
index e6692df20..07a602658 100644
--- a/agents/scpi_psu/scpi_psu_agent.py
+++ b/agents/scpi_psu/scpi_psu_agent.py
@@ -83,7 +83,7 @@ def monitor_output(self, session, params=None):
                         for chan in channels:
                             data['data']["Voltage_{}".format(chan)] = self.psu.get_volt(chan)
                             data['data']["Current_{}".format(chan)] = self.psu.get_curr(chan)
-                    except AssertionError as e:
+                    except ConnectionResetError as e:
                         print(
                             f"Failed to get volt/curr in monitor_output, skipping this iteration: {e}"
                         )
diff --git a/socs/agent/prologix_interface.py b/socs/agent/prologix_interface.py
index c5368f727..1567db2ff 100644
--- a/socs/agent/prologix_interface.py
+++ b/socs/agent/prologix_interface.py
@@ -48,6 +48,8 @@ def write(self, msg):
         except socket.error as e:
             print(f"Socket write failed (disconnect?): {e}")
             self.disconnect_handler()
+            # still write immediately after reconnect,
+            # may not be desirable in certain use cases
             self.write(msg)
             return
         time.sleep(0.1)  # Don't send messages too quickly
@@ -58,7 +60,12 @@ def read(self):
         if not data:
             print("Received no data from socket (disconnect?)")
             self.disconnect_handler()
-            return self.read()
+            # reading from socket immediately after reconnect
+            # should timeout or give irrelevant data,
+            # so raise exception and let caller handle it
+            raise ConnectionResetError(
+                "Recovered connection during read attempt -- this read cannot be satisfied"
+            )
         return data.decode().strip()
 
     def disconnect_handler(self):

From 55c125006a22301af1feef0355a257dcf053c643 Mon Sep 17 00:00:00 2001
From: Jason Guo <jsg353@cornell.edu>
Date: Thu, 30 Jun 2022 17:13:58 -0400
Subject: [PATCH 08/18] Complete monitor_output automatic connection recovery
 test (specifically for graceful socket shutdown)

---
 tests/integration/test_scpi_psu_agent_integration.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/tests/integration/test_scpi_psu_agent_integration.py b/tests/integration/test_scpi_psu_agent_integration.py
index 1ad084a85..edb6a8891 100644
--- a/tests/integration/test_scpi_psu_agent_integration.py
+++ b/tests/integration/test_scpi_psu_agent_integration.py
@@ -107,7 +107,8 @@ def test_scpi_psu_monitor_output_disconnect(
     resp = client.monitor_output.start(test_mode=False, wait=0)
     time.sleep(5)
     gpib_emu.disconnect_reconnect(timeout=3, port=1234)
-    time.sleep(10)
-    # client.init()
-    # resp = client.monitor_output.start(test_mode=True)
-    # check_resp_success(resp)
+    time.sleep(5)
+    resp = client.monitor_output.stop()
+    time.sleep(1)
+    resp = client.monitor_output.status()
+    check_resp_success(resp)

From 4f7c34be828e0aa414212147b4170f00644d4d6b Mon Sep 17 00:00:00 2001
From: Jason Guo <jsg353@cornell.edu>
Date: Wed, 6 Jul 2022 17:23:34 -0400
Subject: [PATCH 09/18] Implement reconnect into LS372 driver (needs testing)

---
 socs/Lakeshore/Lakeshore372.py                | 68 +++++++++++++++++--
 .../test_ls372_agent_integration.py           | 20 +++++-
 2 files changed, 82 insertions(+), 6 deletions(-)

diff --git a/socs/Lakeshore/Lakeshore372.py b/socs/Lakeshore/Lakeshore372.py
index 9b353a6c1..503f035ce 100644
--- a/socs/Lakeshore/Lakeshore372.py
+++ b/socs/Lakeshore/Lakeshore372.py
@@ -2,6 +2,8 @@
 
 import sys
 import socket
+import select
+import time
 import numpy as np
 
 # Lookup keys for command parameters.
@@ -215,6 +217,8 @@ def __init__(self, ip, timeout=10, num_channels=16):
         self.com = _establish_socket_connection(ip, timeout)
         self.num_channels = num_channels
 
+        self.reset = lambda: self.__init__(ip, timeout, num_channels)
+
         self.id = self.get_id()
         self.autoscan = self.get_autoscan()
         # Enable all channels
@@ -234,6 +238,62 @@ def __init__(self, ip, timeout=10, num_channels=16):
         self.sample_heater = Heater(self, 0)
         self.still_heater = Heater(self, 2)
 
+    def connection_check(self, op):
+        assert op in ['read', 'write'], "'op' must be 'read' or 'write'"
+        select_lists = ([self.com,], [], []) if op == 'read' else ([], [self.com,], [])
+        try:
+            ready_to_read, ready_to_write, in_error = \
+                select.select(*select_lists, 5)
+        except select.error:
+            self.com.shutdown(2)
+            self.com.close()
+            print("Lakeshore372 connection error")
+            self.disconnect_handler()
+            self.connection_check(op)  # need to test on real hardware
+            return
+        if op == 'read':
+            assert len(ready_to_read) > 0, "No sockets ready for reading"
+        elif op == 'write':
+            assert len(ready_to_write) > 0, "No sockets ready for writing"
+
+    def disconnect_handler(self):
+        for i in range(5):
+            try:
+                self.reset()
+                print(f"Successfully reconnected on attempt #{i}")
+                return
+            except socket.error as e:
+                print(f"Reconnect attempt #{i} failed with: {e}")
+                time.sleep(1)
+        assert False, "Could not reconnect"
+
+    def write(self, message):
+        self.connection_check('write')
+        msg_str = f'{message}\r\n'.encode()
+        try:
+            self.com.send(msg_str)
+        except socket.error as e:
+            print(f"Socket write failed (disconnect?): {e}")
+            self.disconnect_handler()
+            # still write immediately after reconnect,
+            # may not be desirable in certain use cases
+            self.write(message)
+
+    def read(self):
+        self.connection_check('read')
+        data = self.com.recv(4096)
+        if not data:
+            print("Received no data from socket (disconnect?)")
+            self.disconnect_handler()
+            # reading from socket immediately after reconnect
+            # should timeout or give irrelevant data,
+            # so raise exception and let caller handle it
+            raise ConnectionResetError(
+                "Recovered connection during read attempt -- this read cannot be satisfied"
+            )
+        resp = str(data, 'utf-8').strip()
+        return resp
+
     def msg(self, message):
         """Send message to the Lakeshore 372 over ethernet.
 
@@ -252,14 +312,12 @@ def msg(self, message):
             Response string from the Lakeshore, if any. Else, an empty string.
 
         """
-        msg_str = f'{message}\r\n'.encode()
-
         if '?' in message:
-            self.com.send(msg_str)
+            self.write(message)
             # Try once, if we timeout, try again. Usually gets around single event glitches.
             for attempt in range(2):
                 try:
-                    resp = str(self.com.recv(4096), 'utf-8').strip()
+                    resp = self.read()
                     break
                 except socket.timeout:
                     print("Warning: Caught timeout waiting for response to '%s', trying again "
@@ -268,7 +326,7 @@ def msg(self, message):
                         raise RuntimeError('Query response to Lakeshore timed out after two '
                                            'attempts. Check connection.')
         else:
-            self.com.send(msg_str)
+            self.write(message)
             resp = ''
 
         return resp
diff --git a/tests/integration/test_ls372_agent_integration.py b/tests/integration/test_ls372_agent_integration.py
index 4996634aa..1131db534 100644
--- a/tests/integration/test_ls372_agent_integration.py
+++ b/tests/integration/test_ls372_agent_integration.py
@@ -1,5 +1,6 @@
 import os
 import pytest
+import time
 
 import ocs
 from ocs.base import OpCode
@@ -17,7 +18,8 @@
 
 run_agent = create_agent_runner_fixture(
     '../agents/lakeshore372/LS372_agent.py',
-    'ls372')
+    'ls372',
+    args=["--log-dir", "./logs/"])
 client = create_client_fixture('LSASIM')
 wait_for_crossbar = create_crossbar_fixture()
 
@@ -224,3 +226,19 @@ def test_ls372_get_still_output(wait_for_crossbar, emulator, run_agent, client):
     resp = client.get_still_output()
     assert resp.status == ocs.OK
     assert resp.session['op_code'] == OpCode.SUCCEEDED.value
+
+# @pytest.mark.integtest
+# def test_ls372_disconnect_reconnect(wait_for_crossbar, emulator, run_agent, client):
+#     client.init_lakeshore()
+#     resp = client.acq.start(sample_heater=False, run_once=False)
+#     assert resp.status == ocs.OK
+#     assert resp.session['op_code'] == OpCode.STARTING.value
+
+#     time.sleep(5)
+#     emulator.disconnect_reconnect(timeout=3, port=7777)
+#     time.sleep(5)
+
+#     # client.acq.wait()
+#     # resp = client.acq.status()
+#     # assert resp.status == ocs.OK
+#     # assert resp.session['op_code'] == OpCode.SUCCEEDED.value

From 40322281fee5cdad6aca348d006949f35127861c Mon Sep 17 00:00:00 2001
From: Jason Guo <jsg353@cornell.edu>
Date: Wed, 6 Jul 2022 17:28:54 -0400
Subject: [PATCH 10/18] Properly utilize select timeout behavior by minimally
 specifying select socket lists

---
 socs/agent/prologix_interface.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/socs/agent/prologix_interface.py b/socs/agent/prologix_interface.py
index 1567db2ff..9d8891272 100644
--- a/socs/agent/prologix_interface.py
+++ b/socs/agent/prologix_interface.py
@@ -24,9 +24,10 @@ def configure(self):
 
     def connection_check(self, op):
         assert op in ['read', 'write'], "'op' must be 'read' or 'write'"
+        select_lists = ([self.com,], [], []) if op == 'read' else ([], [self.com,], [])
         try:
             ready_to_read, ready_to_write, in_error = \
-                select.select([self.sock,], [self.sock,], [], 5)
+                select.select(*select_lists, 5)
         except select.error:
             self.sock.shutdown(2)
             self.sock.close()

From 5153813f40eb57b9c3d7d43fde21da373eed57ad Mon Sep 17 00:00:00 2001
From: Jason Guo <jsg353@cornell.edu>
Date: Wed, 6 Jul 2022 17:31:53 -0400
Subject: [PATCH 11/18] Fix self.com -> self.sock

---
 socs/agent/prologix_interface.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/socs/agent/prologix_interface.py b/socs/agent/prologix_interface.py
index 9d8891272..d204dd91f 100644
--- a/socs/agent/prologix_interface.py
+++ b/socs/agent/prologix_interface.py
@@ -24,7 +24,7 @@ def configure(self):
 
     def connection_check(self, op):
         assert op in ['read', 'write'], "'op' must be 'read' or 'write'"
-        select_lists = ([self.com,], [], []) if op == 'read' else ([], [self.com,], [])
+        select_lists = ([self.sock,], [], []) if op == 'read' else ([], [self.sock,], [])
         try:
             ready_to_read, ready_to_write, in_error = \
                 select.select(*select_lists, 5)

From a3d4276bbba1b666cff4bc0aad5a2fda9fd9e235 Mon Sep 17 00:00:00 2001
From: Jason Guo <jsg353@cornell.edu>
Date: Thu, 7 Jul 2022 13:58:22 -0400
Subject: [PATCH 12/18] Successful acq recovery from disconnect

---
 agents/lakeshore372/LS372_agent.py            | 222 +++++++++---------
 socs/Lakeshore/Lakeshore372.py                |   2 +
 .../test_ls372_agent_integration.py           |  30 +--
 3 files changed, 130 insertions(+), 124 deletions(-)

diff --git a/agents/lakeshore372/LS372_agent.py b/agents/lakeshore372/LS372_agent.py
index b92ed4123..40010f479 100644
--- a/agents/lakeshore372/LS372_agent.py
+++ b/agents/lakeshore372/LS372_agent.py
@@ -280,122 +280,126 @@ def acq(self, session, params=None):
                                       f"currently held by {self._lock.job}.")
                         continue
 
-                if self.fake_data:
-                    data = {
-                        'timestamp': time.time(),
-                        'block_name': 'fake-data',
-                        'data': {}
-                    }
-                    for therm in self.thermometers:
-                        reading = np.random.normal(self.res, 20)
-                        data['data'][therm] = reading
-                    time.sleep(.1)
-
-                else:
-                    active_channel = self.module.get_active_channel()
-
-                    # The 372 reports the last updated measurement repeatedly
-                    # during the "pause change time", this results in several
-                    # stale datapoints being recorded. To get around this we
-                    # query the pause time and skip data collection during it
-                    # if the channel has changed (as it would if autoscan is
-                    # enabled.)
-                    if previous_channel != active_channel:
-                        if previous_channel is not None:
-                            pause_time = active_channel.get_pause()
-                            self.log.debug("Pause time for {c}: {p}",
-                                           c=active_channel.channel_num,
-                                           p=pause_time)
-
-                            dwell_time = active_channel.get_dwell()
-                            self.log.debug("User set dwell_time_delay: {p}",
-                                           p=self.dwell_time_delay)
-
-                            # Check user set dwell time isn't too long
-                            if self.dwell_time_delay > dwell_time:
-                                self.log.warn("WARNING: User set dwell_time_delay of "
-                                              + "{delay} s is larger than channel "
-                                              + "dwell time of {chan_time} s. If "
-                                              + "you are autoscanning this will "
-                                              + "cause no data to be collected. "
-                                              + "Reducing dwell time delay to {s} s.",
-                                              delay=self.dwell_time_delay,
-                                              chan_time=dwell_time,
-                                              s=dwell_time - 1)
-                                total_time = pause_time + dwell_time - 1
-                            else:
-                                total_time = pause_time + self.dwell_time_delay
-
-                            for i in range(total_time):
-                                self.log.debug("Sleeping for {t} more seconds...",
-                                               t=total_time - i)
-                                time.sleep(1)
-
-                        # Track the last channel we measured
-                        previous_channel = self.module.get_active_channel()
-
-                    current_time = time.time()
-                    data = {
-                        'timestamp': current_time,
-                        'block_name': active_channel.name,
-                        'data': {}
-                    }
-
-                    # Collect both temperature and resistance values from each Channel
-                    channel_str = active_channel.name.replace(' ', '_')
-                    temp_reading = self.module.get_temp(unit='kelvin',
-                                                        chan=active_channel.channel_num)
-                    res_reading = self.module.get_temp(unit='ohms',
-                                                       chan=active_channel.channel_num)
-
-                    # For data feed
-                    data['data'][channel_str + '_T'] = temp_reading
-                    data['data'][channel_str + '_R'] = res_reading
-                    session.app.publish_to_feed('temperatures', data)
-                    self.log.debug("{data}", data=session.data)
-
-                    # For session.data
-                    field_dict = {channel_str: {"T": temp_reading,
-                                                "R": res_reading,
-                                                "timestamp": current_time}}
-                    session.data['fields'].update(field_dict)
-
-                    # Also queries control channel if enabled
-                    if self.control_chan_enabled:
-                        temp = self.module.get_temp(unit='kelvin', chan=0)
-                        res = self.module.get_temp(unit='ohms', chan=0)
-                        cur_time = time.time()
+                try:
+                    if self.fake_data:
                         data = {
                             'timestamp': time.time(),
-                            'block_name': 'control_chan',
-                            'data': {
-                                'control_T': temp,
-                                'control_R': res
-                            }
+                            'block_name': 'fake-data',
+                            'data': {}
+                        }
+                        for therm in self.thermometers:
+                            reading = np.random.normal(self.res, 20)
+                            data['data'][therm] = reading
+                        time.sleep(.1)
+
+                    else:
+                        active_channel = self.module.get_active_channel()
+
+                        # The 372 reports the last updated measurement repeatedly
+                        # during the "pause change time", this results in several
+                        # stale datapoints being recorded. To get around this we
+                        # query the pause time and skip data collection during it
+                        # if the channel has changed (as it would if autoscan is
+                        # enabled.)
+                        if previous_channel != active_channel:
+                            if previous_channel is not None:
+                                pause_time = active_channel.get_pause()
+                                self.log.debug("Pause time for {c}: {p}",
+                                            c=active_channel.channel_num,
+                                            p=pause_time)
+
+                                dwell_time = active_channel.get_dwell()
+                                self.log.debug("User set dwell_time_delay: {p}",
+                                            p=self.dwell_time_delay)
+
+                                # Check user set dwell time isn't too long
+                                if self.dwell_time_delay > dwell_time:
+                                    self.log.warn("WARNING: User set dwell_time_delay of "
+                                                + "{delay} s is larger than channel "
+                                                + "dwell time of {chan_time} s. If "
+                                                + "you are autoscanning this will "
+                                                + "cause no data to be collected. "
+                                                + "Reducing dwell time delay to {s} s.",
+                                                delay=self.dwell_time_delay,
+                                                chan_time=dwell_time,
+                                                s=dwell_time - 1)
+                                    total_time = pause_time + dwell_time - 1
+                                else:
+                                    total_time = pause_time + self.dwell_time_delay
+
+                                for i in range(total_time):
+                                    self.log.debug("Sleeping for {t} more seconds...",
+                                                t=total_time - i)
+                                    time.sleep(1)
+
+                            # Track the last channel we measured
+                            previous_channel = self.module.get_active_channel()
+
+                        current_time = time.time()
+                        data = {
+                            'timestamp': current_time,
+                            'block_name': active_channel.name,
+                            'data': {}
                         }
+
+                        # Collect both temperature and resistance values from each Channel
+                        channel_str = active_channel.name.replace(' ', '_')
+                        temp_reading = self.module.get_temp(unit='kelvin',
+                                                            chan=active_channel.channel_num)
+                        res_reading = self.module.get_temp(unit='ohms',
+                                                        chan=active_channel.channel_num)
+
+                        # For data feed
+                        data['data'][channel_str + '_T'] = temp_reading
+                        data['data'][channel_str + '_R'] = res_reading
                         session.app.publish_to_feed('temperatures', data)
                         self.log.debug("{data}", data=session.data)
-                        # Updates session data w/ control field
-                        session.data['fields'].update({
-                            'control': {
-                                'T': temp, 'R': res, 'timestamp': cur_time
+
+                        # For session.data
+                        field_dict = {channel_str: {"T": temp_reading,
+                                                    "R": res_reading,
+                                                    "timestamp": current_time}}
+                        session.data['fields'].update(field_dict)
+
+                        # Also queries control channel if enabled
+                        if self.control_chan_enabled:
+                            temp = self.module.get_temp(unit='kelvin', chan=0)
+                            res = self.module.get_temp(unit='ohms', chan=0)
+                            cur_time = time.time()
+                            data = {
+                                'timestamp': time.time(),
+                                'block_name': 'control_chan',
+                                'data': {
+                                    'control_T': temp,
+                                    'control_R': res
+                                }
                             }
-                        })
-
-                if params.get("sample_heater", False):
-                    # Sample Heater
-                    heater = self.module.sample_heater
-                    hout = heater.get_sample_heater_output()
-
-                    current_time = time.time()
-                    htr_data = {
-                        'timestamp': current_time,
-                        'block_name': "heaters",
-                        'data': {}
-                    }
-                    htr_data['data']['sample_heater_output'] = hout
+                            session.app.publish_to_feed('temperatures', data)
+                            self.log.debug("{data}", data=session.data)
+                            # Updates session data w/ control field
+                            session.data['fields'].update({
+                                'control': {
+                                    'T': temp, 'R': res, 'timestamp': cur_time
+                                }
+                            })
+
+                    if params.get("sample_heater", False):
+                        # Sample Heater
+                        heater = self.module.sample_heater
+                        hout = heater.get_sample_heater_output()
+
+                        current_time = time.time()
+                        htr_data = {
+                            'timestamp': current_time,
+                            'block_name': "heaters",
+                            'data': {}
+                        }
+                        htr_data['data']['sample_heater_output'] = hout
 
-                    session.app.publish_to_feed('temperatures', htr_data)
+                        session.app.publish_to_feed('temperatures', htr_data)
+                except ConnectionResetError as e:
+                    print(f"Skipping acq iteration due to: {e}")
+                    continue
 
                 if params['run_once']:
                     break
diff --git a/socs/Lakeshore/Lakeshore372.py b/socs/Lakeshore/Lakeshore372.py
index 503f035ce..1453f60c5 100644
--- a/socs/Lakeshore/Lakeshore372.py
+++ b/socs/Lakeshore/Lakeshore372.py
@@ -325,6 +325,8 @@ def msg(self, message):
                     if attempt == 1:
                         raise RuntimeError('Query response to Lakeshore timed out after two '
                                            'attempts. Check connection.')
+                except ConnectionResetError:
+                    raise
         else:
             self.write(message)
             resp = ''
diff --git a/tests/integration/test_ls372_agent_integration.py b/tests/integration/test_ls372_agent_integration.py
index 1131db534..f4167010d 100644
--- a/tests/integration/test_ls372_agent_integration.py
+++ b/tests/integration/test_ls372_agent_integration.py
@@ -227,18 +227,18 @@ def test_ls372_get_still_output(wait_for_crossbar, emulator, run_agent, client):
     assert resp.status == ocs.OK
     assert resp.session['op_code'] == OpCode.SUCCEEDED.value
 
-# @pytest.mark.integtest
-# def test_ls372_disconnect_reconnect(wait_for_crossbar, emulator, run_agent, client):
-#     client.init_lakeshore()
-#     resp = client.acq.start(sample_heater=False, run_once=False)
-#     assert resp.status == ocs.OK
-#     assert resp.session['op_code'] == OpCode.STARTING.value
-
-#     time.sleep(5)
-#     emulator.disconnect_reconnect(timeout=3, port=7777)
-#     time.sleep(5)
-
-#     # client.acq.wait()
-#     # resp = client.acq.status()
-#     # assert resp.status == ocs.OK
-#     # assert resp.session['op_code'] == OpCode.SUCCEEDED.value
+@pytest.mark.integtest
+def test_ls372_disconnect_reconnect(wait_for_crossbar, emulator, run_agent, client):
+    client.init_lakeshore()
+    resp = client.acq.start(sample_heater=False, run_once=False)
+    assert resp.status == ocs.OK
+    assert resp.session['op_code'] == OpCode.STARTING.value
+
+    time.sleep(5)
+    emulator.disconnect_reconnect(timeout=3, port=7777)
+    time.sleep(5)
+
+    # client.acq.wait()
+    # resp = client.acq.status()
+    # assert resp.status == ocs.OK
+    # assert resp.session['op_code'] == OpCode.SUCCEEDED.value

From 5a98e64b06dd684011d8fc332027b3cb1e331994 Mon Sep 17 00:00:00 2001
From: Jason Guo <jsg353@cornell.edu>
Date: Thu, 7 Jul 2022 15:37:02 -0400
Subject: [PATCH 13/18] Complete disconnect_reconnect integration test case

---
 tests/integration/test_ls372_agent_integration.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/tests/integration/test_ls372_agent_integration.py b/tests/integration/test_ls372_agent_integration.py
index f4167010d..2e6ea5c36 100644
--- a/tests/integration/test_ls372_agent_integration.py
+++ b/tests/integration/test_ls372_agent_integration.py
@@ -238,7 +238,8 @@ def test_ls372_disconnect_reconnect(wait_for_crossbar, emulator, run_agent, clie
     emulator.disconnect_reconnect(timeout=3, port=7777)
     time.sleep(5)
 
-    # client.acq.wait()
-    # resp = client.acq.status()
-    # assert resp.status == ocs.OK
-    # assert resp.session['op_code'] == OpCode.SUCCEEDED.value
+    client.acq.stop()
+    time.sleep(1)
+    resp = client.acq.status()
+    assert resp.status == ocs.OK
+    assert resp.session['op_code'] == OpCode.SUCCEEDED.value

From b7eb45646a4c10dd4257ff6f65f78ffb59cb7baa Mon Sep 17 00:00:00 2001
From: Jason Guo <jsg353@cornell.edu>
Date: Fri, 8 Jul 2022 20:20:51 +0000
Subject: [PATCH 14/18] Successful reconnect with hardware

---
 socs/Lakeshore/Lakeshore372.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/socs/Lakeshore/Lakeshore372.py b/socs/Lakeshore/Lakeshore372.py
index 1453f60c5..cd4cdc315 100644
--- a/socs/Lakeshore/Lakeshore372.py
+++ b/socs/Lakeshore/Lakeshore372.py
@@ -252,12 +252,18 @@ def connection_check(self, op):
             self.connection_check(op)  # need to test on real hardware
             return
         if op == 'read':
-            assert len(ready_to_read) > 0, "No sockets ready for reading"
+#            assert len(ready_to_read) > 0, "No sockets ready for reading"
+            if not len(ready_to_read) > 0:
+                self.disconnect_handler()
+                raise ConnectionResetError("No sockets ready for reading")
         elif op == 'write':
-            assert len(ready_to_write) > 0, "No sockets ready for writing"
+            if not len(ready_to_write) > 0:
+                self.disconnect_handler()
+                raise ConnectionResetError("No sockets ready for writing")
+#            assert len(ready_to_write) > 0, "No sockets ready for writing"
 
     def disconnect_handler(self):
-        for i in range(5):
+        for i in range(500):
             try:
                 self.reset()
                 print(f"Successfully reconnected on attempt #{i}")

From 05cb751709fc58a4861ab37597d16ef40541a847 Mon Sep 17 00:00:00 2001
From: Jason Guo <jsg353@cornell.edu>
Date: Fri, 8 Jul 2022 16:48:24 -0400
Subject: [PATCH 15/18] Clean up connection_check

---
 socs/Lakeshore/Lakeshore372.py | 29 +++++++++++++----------------
 1 file changed, 13 insertions(+), 16 deletions(-)

diff --git a/socs/Lakeshore/Lakeshore372.py b/socs/Lakeshore/Lakeshore372.py
index cd4cdc315..e9dbcd3f9 100644
--- a/socs/Lakeshore/Lakeshore372.py
+++ b/socs/Lakeshore/Lakeshore372.py
@@ -244,23 +244,20 @@ def connection_check(self, op):
         try:
             ready_to_read, ready_to_write, in_error = \
                 select.select(*select_lists, 5)
-        except select.error:
-            self.com.shutdown(2)
-            self.com.close()
-            print("Lakeshore372 connection error")
+        except select.error as e:
+            # self.com.shutdown(2)
+            # self.com.close()
+            # print("Lakeshore372 connection error")
+            # self.disconnect_handler()
+            # self.connection_check(op)  # need to test on real hardware
+            # return
+            raise Exception("Triggered select.error block unexpectedly") from e
+        if op == 'read' and not ready_to_read:
             self.disconnect_handler()
-            self.connection_check(op)  # need to test on real hardware
-            return
-        if op == 'read':
-#            assert len(ready_to_read) > 0, "No sockets ready for reading"
-            if not len(ready_to_read) > 0:
-                self.disconnect_handler()
-                raise ConnectionResetError("No sockets ready for reading")
-        elif op == 'write':
-            if not len(ready_to_write) > 0:
-                self.disconnect_handler()
-                raise ConnectionResetError("No sockets ready for writing")
-#            assert len(ready_to_write) > 0, "No sockets ready for writing"
+            raise ConnectionResetError("No sockets ready for reading")
+        elif op == 'write' and not ready_to_write:
+            self.disconnect_handler()
+            raise ConnectionResetError("No sockets ready for writing")
 
     def disconnect_handler(self):
         for i in range(500):

From 5f8318f3b69b3d1a2e1edb905cea87c00f5b5453 Mon Sep 17 00:00:00 2001
From: Jason Guo <jsg353@cornell.edu>
Date: Fri, 8 Jul 2022 17:19:49 -0400
Subject: [PATCH 16/18] Simplify disconnect_handler to always raise
 ConnectionResetError

---
 socs/Lakeshore/Lakeshore372.py | 34 ++++++++++++----------------------
 1 file changed, 12 insertions(+), 22 deletions(-)

diff --git a/socs/Lakeshore/Lakeshore372.py b/socs/Lakeshore/Lakeshore372.py
index e9dbcd3f9..292ceccf9 100644
--- a/socs/Lakeshore/Lakeshore372.py
+++ b/socs/Lakeshore/Lakeshore372.py
@@ -253,22 +253,23 @@ def connection_check(self, op):
             # return
             raise Exception("Triggered select.error block unexpectedly") from e
         if op == 'read' and not ready_to_read:
-            self.disconnect_handler()
-            raise ConnectionResetError("No sockets ready for reading")
+            self.disconnect_handler("No sockets ready for reading")
         elif op == 'write' and not ready_to_write:
-            self.disconnect_handler()
-            raise ConnectionResetError("No sockets ready for writing")
+            self.disconnect_handler("No sockets ready for writing")
 
-    def disconnect_handler(self):
-        for i in range(500):
+    def disconnect_handler(self, reset_reason):
+        max_attempts = 500
+        for i in range(max_attempts):
             try:
                 self.reset()
-                print(f"Successfully reconnected on attempt #{i}")
-                return
+                break
             except socket.error as e:
                 print(f"Reconnect attempt #{i} failed with: {e}")
+                if i == max_attempts - 1:
+                    assert False, "Could not reconnect"
                 time.sleep(1)
-        assert False, "Could not reconnect"
+        print(f"Successfully reconnected on attempt #{i}")
+        raise ConnectionResetError(reset_reason) # should be caught by agent
 
     def write(self, message):
         self.connection_check('write')
@@ -276,24 +277,13 @@ def write(self, message):
         try:
             self.com.send(msg_str)
         except socket.error as e:
-            print(f"Socket write failed (disconnect?): {e}")
-            self.disconnect_handler()
-            # still write immediately after reconnect,
-            # may not be desirable in certain use cases
-            self.write(message)
+            self.disconnect_handler(f"Socket write failed (disconnect?): {e}")
 
     def read(self):
         self.connection_check('read')
         data = self.com.recv(4096)
         if not data:
-            print("Received no data from socket (disconnect?)")
-            self.disconnect_handler()
-            # reading from socket immediately after reconnect
-            # should timeout or give irrelevant data,
-            # so raise exception and let caller handle it
-            raise ConnectionResetError(
-                "Recovered connection during read attempt -- this read cannot be satisfied"
-            )
+            self.disconnect_handler("Received no data from socket recv")
         resp = str(data, 'utf-8').strip()
         return resp
 

From 19cc255df92e532670875d0b32abee497d15b9dd Mon Sep 17 00:00:00 2001
From: Jason Guo <jsg353@cornell.edu>
Date: Mon, 11 Jul 2022 15:18:26 -0400
Subject: [PATCH 17/18] Update scpi reconnection handling to be simpler based
 on LS372 work

---
 socs/agent/prologix_interface.py | 52 ++++++++++++++------------------
 1 file changed, 22 insertions(+), 30 deletions(-)

diff --git a/socs/agent/prologix_interface.py b/socs/agent/prologix_interface.py
index d204dd91f..7a35bfe0a 100644
--- a/socs/agent/prologix_interface.py
+++ b/socs/agent/prologix_interface.py
@@ -28,17 +28,18 @@ def connection_check(self, op):
         try:
             ready_to_read, ready_to_write, in_error = \
                 select.select(*select_lists, 5)
-        except select.error:
-            self.sock.shutdown(2)
-            self.sock.close()
-            print("Prologix interface connection error")
-            self.disconnect_handler()
-            self.connection_check(op)  # need to test on real hardware
-            return
-        if op == 'read':
-            assert len(ready_to_read) > 0, "No sockets ready for reading"
-        elif op == 'write':
-            assert len(ready_to_write) > 0, "No sockets ready for writing"
+        except select.error as e:
+            # self.sock.shutdown(2)
+            # self.sock.close()
+            # print("Lakeshore372 connection error")
+            # self.disconnect_handler()
+            # self.connection_check(op)  # need to test on real hardware
+            # return
+            raise Exception("Triggered select.error block unexpectedly") from e
+        if op == 'read' and not ready_to_read:
+            self.disconnect_handler("No sockets ready for reading")
+        elif op == 'write' and not ready_to_write:
+            self.disconnect_handler("No sockets ready for writing")
 
 
     def write(self, msg):
@@ -47,39 +48,30 @@ def write(self, msg):
         try:
             self.sock.sendall(message.encode())
         except socket.error as e:
-            print(f"Socket write failed (disconnect?): {e}")
-            self.disconnect_handler()
-            # still write immediately after reconnect,
-            # may not be desirable in certain use cases
-            self.write(msg)
-            return
+            self.disconnect_handler(f"Socket write failed (disconnect?): {e}")
         time.sleep(0.1)  # Don't send messages too quickly
 
     def read(self):
         self.connection_check('read')
         data = self.sock.recv(128)
         if not data:
-            print("Received no data from socket (disconnect?)")
-            self.disconnect_handler()
-            # reading from socket immediately after reconnect
-            # should timeout or give irrelevant data,
-            # so raise exception and let caller handle it
-            raise ConnectionResetError(
-                "Recovered connection during read attempt -- this read cannot be satisfied"
-            )
+            self.disconnect_handler("Received no data from socket (disconnect?)")
         return data.decode().strip()
 
-    def disconnect_handler(self):
-        for i in range(5):
+    def disconnect_handler(self, reset_reason):
+        max_attempts = 500
+        for i in range(max_attempts):
             try:
                 self.conn_socket()
                 self.configure()
-                print(f"Successfully reconnected on attempt #{i}")
-                return
+                break
             except socket.error as e:
                 print(f"Reconnect attempt #{i} failed with: {e}")
+                if i == max_attempts - 1:
+                    assert False, "Could not reconnect"
                 time.sleep(1)
-        assert False, "Could not reconnect"
+        print(f"Successfully reconnected on attempt #{i}")
+        raise ConnectionResetError(reset_reason) # should be caught by agent
 
     def version(self):
         self.write('++ver')

From 4424a000723af236b0cd4131028a12f58ca149f0 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 11 Jul 2022 19:48:02 +0000
Subject: [PATCH 18/18] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 agents/lakeshore372/LS372_agent.py            | 26 +++++++++----------
 socs/Lakeshore/Lakeshore372.py                |  4 +--
 socs/agent/prologix_interface.py              |  5 ++--
 .../test_ls372_agent_integration.py           |  1 +
 4 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/agents/lakeshore372/LS372_agent.py b/agents/lakeshore372/LS372_agent.py
index 40010f479..399a2db87 100644
--- a/agents/lakeshore372/LS372_agent.py
+++ b/agents/lakeshore372/LS372_agent.py
@@ -305,31 +305,31 @@ def acq(self, session, params=None):
                             if previous_channel is not None:
                                 pause_time = active_channel.get_pause()
                                 self.log.debug("Pause time for {c}: {p}",
-                                            c=active_channel.channel_num,
-                                            p=pause_time)
+                                               c=active_channel.channel_num,
+                                               p=pause_time)
 
                                 dwell_time = active_channel.get_dwell()
                                 self.log.debug("User set dwell_time_delay: {p}",
-                                            p=self.dwell_time_delay)
+                                               p=self.dwell_time_delay)
 
                                 # Check user set dwell time isn't too long
                                 if self.dwell_time_delay > dwell_time:
                                     self.log.warn("WARNING: User set dwell_time_delay of "
-                                                + "{delay} s is larger than channel "
-                                                + "dwell time of {chan_time} s. If "
-                                                + "you are autoscanning this will "
-                                                + "cause no data to be collected. "
-                                                + "Reducing dwell time delay to {s} s.",
-                                                delay=self.dwell_time_delay,
-                                                chan_time=dwell_time,
-                                                s=dwell_time - 1)
+                                                  + "{delay} s is larger than channel "
+                                                  + "dwell time of {chan_time} s. If "
+                                                  + "you are autoscanning this will "
+                                                  + "cause no data to be collected. "
+                                                  + "Reducing dwell time delay to {s} s.",
+                                                  delay=self.dwell_time_delay,
+                                                  chan_time=dwell_time,
+                                                  s=dwell_time - 1)
                                     total_time = pause_time + dwell_time - 1
                                 else:
                                     total_time = pause_time + self.dwell_time_delay
 
                                 for i in range(total_time):
                                     self.log.debug("Sleeping for {t} more seconds...",
-                                                t=total_time - i)
+                                                   t=total_time - i)
                                     time.sleep(1)
 
                             # Track the last channel we measured
@@ -347,7 +347,7 @@ def acq(self, session, params=None):
                         temp_reading = self.module.get_temp(unit='kelvin',
                                                             chan=active_channel.channel_num)
                         res_reading = self.module.get_temp(unit='ohms',
-                                                        chan=active_channel.channel_num)
+                                                           chan=active_channel.channel_num)
 
                         # For data feed
                         data['data'][channel_str + '_T'] = temp_reading
diff --git a/socs/Lakeshore/Lakeshore372.py b/socs/Lakeshore/Lakeshore372.py
index 292ceccf9..c62575203 100644
--- a/socs/Lakeshore/Lakeshore372.py
+++ b/socs/Lakeshore/Lakeshore372.py
@@ -240,7 +240,7 @@ def __init__(self, ip, timeout=10, num_channels=16):
 
     def connection_check(self, op):
         assert op in ['read', 'write'], "'op' must be 'read' or 'write'"
-        select_lists = ([self.com,], [], []) if op == 'read' else ([], [self.com,], [])
+        select_lists = ([self.com, ], [], []) if op == 'read' else ([], [self.com, ], [])
         try:
             ready_to_read, ready_to_write, in_error = \
                 select.select(*select_lists, 5)
@@ -269,7 +269,7 @@ def disconnect_handler(self, reset_reason):
                     assert False, "Could not reconnect"
                 time.sleep(1)
         print(f"Successfully reconnected on attempt #{i}")
-        raise ConnectionResetError(reset_reason) # should be caught by agent
+        raise ConnectionResetError(reset_reason)  # should be caught by agent
 
     def write(self, message):
         self.connection_check('write')
diff --git a/socs/agent/prologix_interface.py b/socs/agent/prologix_interface.py
index 7a35bfe0a..f9f60d835 100644
--- a/socs/agent/prologix_interface.py
+++ b/socs/agent/prologix_interface.py
@@ -24,7 +24,7 @@ def configure(self):
 
     def connection_check(self, op):
         assert op in ['read', 'write'], "'op' must be 'read' or 'write'"
-        select_lists = ([self.sock,], [], []) if op == 'read' else ([], [self.sock,], [])
+        select_lists = ([self.sock, ], [], []) if op == 'read' else ([], [self.sock, ], [])
         try:
             ready_to_read, ready_to_write, in_error = \
                 select.select(*select_lists, 5)
@@ -41,7 +41,6 @@ def connection_check(self, op):
         elif op == 'write' and not ready_to_write:
             self.disconnect_handler("No sockets ready for writing")
 
-
     def write(self, msg):
         self.connection_check('write')
         message = msg + '\n'
@@ -71,7 +70,7 @@ def disconnect_handler(self, reset_reason):
                     assert False, "Could not reconnect"
                 time.sleep(1)
         print(f"Successfully reconnected on attempt #{i}")
-        raise ConnectionResetError(reset_reason) # should be caught by agent
+        raise ConnectionResetError(reset_reason)  # should be caught by agent
 
     def version(self):
         self.write('++ver')
diff --git a/tests/integration/test_ls372_agent_integration.py b/tests/integration/test_ls372_agent_integration.py
index 2e6ea5c36..1c0c95e10 100644
--- a/tests/integration/test_ls372_agent_integration.py
+++ b/tests/integration/test_ls372_agent_integration.py
@@ -227,6 +227,7 @@ def test_ls372_get_still_output(wait_for_crossbar, emulator, run_agent, client):
     assert resp.status == ocs.OK
     assert resp.session['op_code'] == OpCode.SUCCEEDED.value
 
+
 @pytest.mark.integtest
 def test_ls372_disconnect_reconnect(wait_for_crossbar, emulator, run_agent, client):
     client.init_lakeshore()