From 2d0684a8b39a752fa6bdffe6201edd274a96e73c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20S=C3=A1nchez-Gallego?= <gallegoj@uw.edu>
Date: Sat, 11 Jan 2025 13:11:07 +0000
Subject: [PATCH] Add timeouts to valve operations

---
 CHANGELOG.md                  |  1 +
 src/lvmcryo/handlers/valve.py | 18 +++++++++++-------
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2633add..7f57ab5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,7 @@
 ### ✨ Improved
 
 * Try to close as many LN2 valves as possible before failing
+* Add timeouts to valve operations.
 
 
 ## 0.3.7 - 2025-01-09
diff --git a/src/lvmcryo/handlers/valve.py b/src/lvmcryo/handlers/valve.py
index 7bc61c8..c3cc528 100644
--- a/src/lvmcryo/handlers/valve.py
+++ b/src/lvmcryo/handlers/valve.py
@@ -19,6 +19,7 @@
 
 from lvmopstools.clu import CluClient
 from lvmopstools.retrier import Retrier
+from lvmopstools.utils import with_timeout
 
 from lvmcryo.config import get_internal_config
 from lvmcryo.handlers.thermistor import ThermistorHandler
@@ -93,7 +94,7 @@ async def valve_on_off(
 
     if on is True and isinstance(timeout, (int, float)) and use_script is True:
         # First we need to get the outlet number.
-        info = await outlet_info(actor, outlet_name)
+        info = await with_timeout(outlet_info(actor, outlet_name), timeout=5)
         id_ = info["id"]
 
         command_string = f"scripts run cycle_with_timeout {id_} {timeout}"
@@ -148,11 +149,14 @@ async def close_all_valves(config: Configuration | None = None, dry_run: bool =
 
     await asyncio.gather(
         *[
-            valve_on_off(
-                valve_info[valve]["actor"],
-                valve_info[valve]["outlet"],
-                False,
-                dry_run=dry_run,
+            with_timeout(
+                valve_on_off(
+                    valve_info[valve]["actor"],
+                    valve_info[valve]["outlet"],
+                    False,
+                    dry_run=dry_run,
+                ),
+                timeout=10,
             )
             for valve in valve_info
         ]
@@ -210,7 +214,7 @@ async def check(self):
         """Check the connection to the NPS."""
 
         try:
-            info = await outlet_info(self.actor, self.outlet)
+            info = await with_timeout(outlet_info(self.actor, self.outlet), timeout=5)
             assert isinstance(info, dict), "Invalid outlet info."
             assert not info["state"], "Valve is already open."
         except Exception as err: