Use timeout value in connect, handle socket errors.

Author: rshipp

nut2.py

@@ -94,19 +94,23 @@ def _connect(self):
         """
         logging.debug("Connecting to host")
 
-        self._srv_handler = telnetlib.Telnet(self._host, self._port)
-
-        if self._login != None:
-            self._srv_handler.write("USERNAME %s\n" % self._login)
-            result = self._srv_handler.read_until("\n", self._timeout)
-            if not result == "OK\n":
-                raise PyNUTError(result.replace("\n", ""))
-
-        if self._password != None:
-            self._srv_handler.write("PASSWORD %s\n" % self._password)
-            result = self._srv_handler.read_until("\n", self._timeout)
-            if not result == "OK\n":
-                raise PyNUTError(result.replace("\n", ""))
+        try:
+            self._srv_handler = telnetlib.Telnet(self._host, self._port,
+                    timeout=self._timeout)
+
+            if self._login != None:
+                self._srv_handler.write("USERNAME %s\n" % self._login)
+                result = self._srv_handler.read_until("\n", self._timeout)
+                if not result == "OK\n":
+                    raise PyNUTError(result.replace("\n", ""))
+
+            if self._password != None:
+                self._srv_handler.write("PASSWORD %s\n" % self._password)
+                result = self._srv_handler.read_until("\n", self._timeout)
+                if not result == "OK\n":
+                    raise PyNUTError(result.replace("\n", ""))
+        except telnetlib.socket.error:
+            raise PyNUTError("Socket error.")
 
     def description(self, ups):
         """Returns the description for a given UPS."""

tests/mockserver.py

@@ -2,7 +2,7 @@
 
 class MockServer(object):
     def __init__(self, host=None, port=None, broken=True, ok=True,
-            broken_username=False):
+            broken_username=False, timeout=None):
         self.valid = "test"
         self.valid_desc = '"Test UPS 1"'
         self.broken = broken