Improve error handling and add its tests. fixes #12

Author: repeatedly

msgpackrpc/future.py

@@ -39,7 +39,10 @@ def get(self):
                 if self._error_handler is not None:
                     self._error_handler(self._error)
                 else:
-                    raise error.RPCError(self._error)
+                    if isinstance(self._error, error.RPCError):
+                        raise self._error
+                    else:
+                        raise error.RPCError(self._error)
             else:
                 return self._result
 

msgpackrpc/session.py

@@ -3,6 +3,7 @@
 from msgpackrpc.future import Future
 from msgpackrpc.transport import tcp
 from msgpackrpc.compat import iteritems
+from msgpackrpc.error import TimeoutError
 
 
 class Session(object):
@@ -109,7 +110,7 @@ def step_timeout(self):
         self._loop.stop()
         for timeout in timeouts:
             future = self._request_table.pop(timeout)
-            future.set_error("Request timed out")
+            future.set_error(TimeoutError("Request timed out"))
         self._loop.start()
 
 

msgpackrpc/transport/tcp.py

@@ -3,7 +3,7 @@
 from tornado.iostream import IOStream
 
 import msgpackrpc.message
-from msgpackrpc.error import RPCError
+from msgpackrpc.error import RPCError, TransportError
 
 
 class BaseSocket(object):
@@ -120,7 +120,7 @@ def on_connect_failed(self, sock):
         else:
             self._connecting = 0
             self._pending = []
-            self._session.on_connect_failed("Retry connection over the limit")
+            self._session.on_connect_failed(TransportError("Retry connection over the limit"))
 
     def on_close(self, sock):
         # Avoid calling self.on_connect_failed after self.close called.

test/test_msgpackrpc.py

@@ -1,3 +1,4 @@
+from time import sleep
 import threading
 try:
     import unittest2 as unittest
@@ -10,6 +11,7 @@
 
 
 class TestMessagePackRPC(unittest.TestCase):
+    ENABLE_TIMEOUT_TEST = False
 
     class TestArg:
         ''' this class must know completely how to deserialize '''
@@ -52,6 +54,9 @@ def add_arg(self, arg0, arg1):
         def raise_error(self):
             raise Exception('error')
 
+        def long_exec(self):
+            sleep(3)
+            return 'finish!'
 
     def setUp(self):
         self._address = msgpackrpc.Address('localhost', helper.unused_port())
@@ -141,5 +146,30 @@ def test_unknown_method(self):
             message = e.args[0]
             self.assertEqual(message, "'unknown' method not found", "Error message mismatched")
 
+    def test_connect_failed(self):
+        client = self.setup_env();
+        client = msgpackrpc.Client(msgpackrpc.Address('localhost', self._address.port - 10), unpack_encoding='utf-8')
+        self.assertRaises(error.TransportError, lambda: client.call('hello'))
+
+    def test_timeout(self):
+        client = self.setup_env();
+
+        if self.__class__.ENABLE_TIMEOUT_TEST:
+            self.assertEqual(client.call('long_exec'), 'finish!', "'long_exec' result is incorrect")
+
+            client = msgpackrpc.Client(self._address, timeout=1, unpack_encoding='utf-8')
+            self.assertRaises(error.TimeoutError, lambda: client.call('long_exec'))
+        else:
+            print("Skip test_timeout")
+
+
 if __name__ == '__main__':
+    import sys
+
+    try:
+        sys.argv.remove('--timeout-test')
+        TestMessagePackRPC.ENABLE_TIMEOUT_TEST = True
+    except:
+        pass
+
     unittest.main()