Allow setting a timeout for Client.waitForServer()

Change-Id: I614de364a668b1ae01ad361254fd4afcdfe48051

2 files changed, 20 insertions, 1 deletions

diff --git a/gear/__init__.py b/gear/__init__.py
index b19d386..f820aa2 100644
--- a/gear/__init__.py
+++ b/gear/__init__.py
@@ -1186,15 +1186,28 @@ class BaseClient(BaseClientServer):
         finally:
             self.connections_condition.release()
 
-    def waitForServer(self):
+    def _checkTimeout(self, start_time, timeout):
+        if time.time() - start_time > timeout:
+            raise TimeoutError()
+
+    def waitForServer(self, timeout=None):
         """Wait for at least one server to be connected.
 
         Block until at least one gearman server is connected.
+
+        :arg numeric timeout: Number of seconds to wait for a connection.
+            If None, wait forever (default: no timeout).
+        :raises TimeoutError: If the timeout is reached before any server
+            connects.
         """
+
         connected = False
+        start_time = time.time()
         while self.running:
             self.connections_condition.acquire()
             while self.running and not self.active_connections:
+                if timeout is not None:
+                    self._checkTimeout(start_time, timeout)
                 self.log.debug("Waiting for at least one active connection")
                 self.connections_condition.wait(timeout=1)
             if self.active_connections:
diff --git a/gear/tests/test_gear.py b/gear/tests/test_gear.py
index d90d99e..c16df04 100644
--- a/gear/tests/test_gear.py
+++ b/gear/tests/test_gear.py
@@ -203,6 +203,12 @@ class TestServerConnection(tests.BaseTestCase):
 
 class TestClient(tests.BaseTestCase):
 
+    def test_wait_for_server_timeout(self):
+        client = gear.Client('client')
+        client.addServer('127.0.0.1', 0)
+        self.assertRaises(gear.TimeoutError,
+                          client.waitForServer, timeout=1)
+
     def test_handleStatusRes_1(self):
         client = gear.Client()