Merge pull request #41 from dbarrosop/develop

Added support to specify timeout

pyeapi/client.py

@@ -375,7 +375,7 @@ def make_connection(transport, **kwargs):
     return klass(**kwargs)
 
 def connect(transport=None, host='localhost', username='admin',
-            password='', port=None):
+            password='', port=None, timeout=60):
     """ Creates a connection using the supplied settings
 
     This function will create a connection to an Arista EOS node using
@@ -401,7 +401,7 @@ def connect(transport=None, host='localhost', username='admin',
     """
     transport = transport or DEFAULT_TRANSPORT
     return make_connection(transport, host=host, username=username,
-                           password=password, port=port)
+                           password=password, port=port, timeout=timeout)
 
 
 class Node(object):
@@ -745,7 +745,3 @@ def connect_to(name):
                          port=kwargs.get('port'))
     node = Node(connection, **kwargs)
     return node
-
-
-
-

pyeapi/eapilib.py

@@ -61,11 +61,11 @@
 DEFAULT_UNIX_SOCKET = '/var/run/command-api.sock'
 
 
-def https_connection_factory(path, host, port, context=None):
+def https_connection_factory(path, host, port, context=None, timeout=60):
     # ignore ssl context for python versions before 2.7.9
     if sys.hexversion < 34015728:
-        return HttpsConnection(path, host, port)
-    return HttpsConnection(path, host, port, context=context)
+        return HttpsConnection(path, host, port, timeout=timeout)
+    return HttpsConnection(path, host, port, context=context, timeout=timeout)
 
 class EapiError(Exception):
     """Base exception class for all exceptions generated by eapilib
@@ -156,9 +156,10 @@ def __init__(self, connection_type, message, commands=None):
 
 class SocketConnection(HTTPConnection):
 
-    def __init__(self, path):
+    def __init__(self, path, timeout=60):
         HTTPConnection.__init__(self, 'localhost')
         self.path = path
+        self.timeout = timeout
 
     def __str__(self):
         return 'unix:%s' % self.path
@@ -168,6 +169,7 @@ def __repr__(self):
 
     def connect(self):
         self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        self.sock.settimeout(self.timeout)
         self.sock.connect(self.path)
 
 class HttpConnection(HTTPConnection):
@@ -470,30 +472,30 @@ def execute(self, commands, encoding='json', **kwargs):
             raise
 
 class SocketEapiConnection(EapiConnection):
-    def __init__(self, path=None, **kwargs):
+    def __init__(self, path=None, timeout=60, **kwargs):
         super(SocketEapiConnection, self).__init__()
         path = path or DEFAULT_UNIX_SOCKET
-        self.transport = SocketConnection(path)
+        self.transport = SocketConnection(path, timeout)
 
 class HttpLocalEapiConnection(EapiConnection):
-    def __init__(self, port=None, path=None, **kwargs):
+    def __init__(self, port=None, path=None, timeout=60, **kwargs):
         super(HttpLocalEapiConnection, self).__init__()
         port = port or DEFAULT_HTTP_LOCAL_PORT
         path = path or DEFAULT_HTTP_PATH
-        self.transport = HttpConnection(path, 'localhost', port)
+        self.transport = HttpConnection(path, 'localhost', port, timeout=timeout)
 
 class HttpEapiConnection(EapiConnection):
     def __init__(self, host, port=None, path=None, username=None,
-                 password=None, **kwargs):
+                 password=None, timeout=60, **kwargs):
         super(HttpEapiConnection, self).__init__()
         port = port or DEFAULT_HTTP_PORT
         path = path or DEFAULT_HTTP_PATH
-        self.transport = HttpConnection(path, host, port)
+        self.transport = HttpConnection(path, host, port, timeout=timeout)
         self.authentication(username, password)
 
 class HttpsEapiConnection(EapiConnection):
     def __init__(self, host, port=None, path=None, username=None,
-                 password=None, context=None, **kwargs):
+                 password=None, context=None, timeout=60, **kwargs):
         super(HttpsEapiConnection, self).__init__()
         port = port or DEFAULT_HTTPS_PORT
         path = path or DEFAULT_HTTP_PATH
@@ -503,7 +505,7 @@ def __init__(self, host, port=None, path=None, username=None,
         if context is None and not enforce_verification:
             context = self.disable_certificate_verification()
 
-        self.transport = https_connection_factory(path, host, port, context)
+        self.transport = https_connection_factory(path, host, port, context, timeout)
         self.authentication(username, password)
 
     def disable_certificate_verification(self):
@@ -518,6 +520,3 @@ def disable_certificate_verification(self):
         # temporary until a proper fix is implemented.
         if hasattr(ssl, '_create_unverified_context'):
             return ssl._create_unverified_context()
-
-
-

test/unit/test_client.py

@@ -194,7 +194,7 @@ def test_hosts_for_tag_returns_names(self):
     def test_connect_types(self, connection):
         transports = list(pyeapi.client.TRANSPORTS.keys())
         kwargs = dict(host='localhost', username='admin', password='',
-                      port=None)
+                      port=None, timeout=60)
 
         for transport in transports:
             pyeapi.client.connect(transport)
@@ -250,7 +250,7 @@ def test_connect_default_type(self):
         with patch.dict(pyeapi.client.TRANSPORTS, {'https': transport}):
             pyeapi.client.connect()
             kwargs = dict(host='localhost', username='admin', password='',
-                          port=None)
+                          port=None, timeout=60)
             transport.assert_called_once_with(**kwargs)
 
     def test_connect_to_with_config(self):
@@ -260,7 +260,7 @@ def test_connect_to_with_config(self):
             pyeapi.client.load_config(filename=conf)
             pyeapi.client.connect_to('test1')
             kwargs = dict(host='192.168.1.16', username='eapi',
-                          password='password', port=None)
+                          password='password', port=None, timeout=60)
             transport.assert_called_once_with(**kwargs)