From c281490f931d1fc8f3a078b3f6b1ce8439a62c29 Mon Sep 17 00:00:00 2001
From: test <test@test.com>
Date: Thu, 30 Nov 2023 16:49:28 +0100
Subject: [PATCH] Better options

---
 spinnman/spalloc/spalloc_client.py | 21 +++++++++++----------
 spinnman/spalloc/spalloc_job.py    | 13 +++++--------
 2 files changed, 16 insertions(+), 18 deletions(-)

diff --git a/spinnman/spalloc/spalloc_client.py b/spinnman/spalloc/spalloc_client.py
index f440de6f3..9ecab5b66 100644
--- a/spinnman/spalloc/spalloc_client.py
+++ b/spinnman/spalloc/spalloc_client.py
@@ -615,23 +615,24 @@ def open_udp_listener_connection(self) -> UDPConnection:
 
     @overrides(SpallocJob.wait_for_state_change)
     def wait_for_state_change(self, old_state: SpallocState,
-                              timeout: Optional[int] = None) -> SpallocState:
-        while old_state != SpallocState.DESTROYED:
-            obj = self._get(self._url, wait="true", timeout=timeout).json()
+                              n_retries: Optional[int] = None) -> SpallocState:
+        retries = 0
+        while (old_state != SpallocState.DESTROYED and 
+               (n_retries is None or retries < n_retries)):
+            retries += 1
+            obj = self._get(self._url, wait="true", timeout=None).json()
             s = SpallocState[obj["state"]]
             if s != old_state or s == SpallocState.DESTROYED:
                 return s
         return old_state
 
     @overrides(SpallocJob.wait_until_ready)
-    def wait_until_ready(self, timeout: Optional[int] = None,
-                         n_retries: Optional[int] = None):
+    def wait_until_ready(self, n_retries: Optional[int] = None):
         state = self.get_state()
-        retries = 0
-        while (state != SpallocState.READY and
-               (n_retries is None or retries < n_retries)):
-            retries += 1
-            state = self.wait_for_state_change(state, timeout=timeout)
+        while state != SpallocState.READY:
+            state = self.wait_for_state_change(state, n_retries=n_retries)
+            if state == SpallocState.READY or n_retries is not None:
+                return
             if state == SpallocState.DESTROYED:
                 raise SpallocException("job was unexpectedly destroyed")
 
diff --git a/spinnman/spalloc/spalloc_job.py b/spinnman/spalloc/spalloc_job.py
index 7cc8caac4..e3ca15010 100644
--- a/spinnman/spalloc/spalloc_job.py
+++ b/spinnman/spalloc/spalloc_job.py
@@ -141,15 +141,15 @@ def create_transceiver(self) -> Transceiver:
 
     @abstractmethod
     def wait_for_state_change(self, old_state: SpallocState,
-                              timeout: Optional[int] = None) -> SpallocState:
+                              n_retries: Optional[int] = None) -> SpallocState:
         """
         Wait until the allocation is not in the given old state.
 
         :param SpallocState old_state:
             The state that we are looking to change out of.
-        :param timeout:
-            The time to wait, or None to wait forever
-        :type timeout: int or None
+        :param n_retries:
+            The number of times to retry or None to try forever
+        :type n_retries: int or None
         :return: The state that the allocation is now in.
 
             .. note::
@@ -159,13 +159,10 @@ def wait_for_state_change(self, old_state: SpallocState,
         raise NotImplementedError()
 
     @abstractmethod
-    def wait_until_ready(self, timeout: Optional[int] = None,
-                         n_retries: Optional[int] = None):
+    def wait_until_ready(self, n_retries: Optional[int] = None):
         """
         Wait until the allocation is in the ``READY`` state.
 
-        :param timeout: The timeout or None to wait forever
-        :type timeout: int or None
         :param n_retries:
             The number of times to retry, or None to retry forever
         :type n_retries: int or None