Fixed a bug can not get the return value when None is returned from the server.

Author: kazupon

msgpackrpc/future.py

@@ -11,26 +11,37 @@ def __init__(self, loop, timeout, callback=None):
         self._loop = loop
         self._error = None
         self._result = None
+        self._set_flag = False
         self._timeout = timeout
         self._callback = callback
         self._error_handler = None
         self._result_handler = None
 
     def join(self):
-        while (self._error is None and self._result is None):
+        while (not self._set_flag):
             self._loop.start()
 
     def get(self):
         self.join()
+
+        assert self._set_flag == True
+        if not self._set_flag:
+            # TODO: should be designed error !!
+            raise error.RPCError(128)
+
         if self._result is not None:
             if self._result_handler is None:
                 return self._result
             else:
                 self._result_handler(self._result)
-
-        if self._error_handler is not None:
-            self._error_handler(self._error)
-        raise error.RPCError(self._error)
+        else:
+            if self._error is not None:
+                if self._error_handler is not None:
+                    self._error_handler(self._error)
+                else:
+                    raise error.RPCError(self._error)
+            else:
+                return self._result
 
     def set(self, error=None, result=None):
         self._error = error
@@ -44,16 +55,16 @@ def result(self):
         return self._result
 
     def set_result(self, result):
-        assert result is not None
         self.set(result=result)
+        self._set_flag = True
 
     @property
     def error(self):
         return self._error
 
     def set_error(self, error):
-        assert error is not None
         self.set(error=error)
+        self._set_flag = True
 
     def attach_callback(self, callback):
         self._callback = callback

test/test_msgpackrpc.py

@@ -15,6 +15,12 @@ def hello(self):
         def sum(self, x, y):
             return x + y
 
+        def nil(self):
+            return None
+
+        def raise_error(self):
+            raise Exception('error')
+
 
     def setUp(self):
         self._address = msgpackrpc.Address('localhost', helper.unused_port())
@@ -39,33 +45,46 @@ def tearDown(self):
 
     def test_call(self):
         client = self.setup_env();
-        result = client.call('hello')
-        self.assertEqual(result, "world", "'hello' result is incorrect")
 
-        result = client.call('sum', 1, 2)
-        self.assertEqual(result, 3, "'sum' result is incorrect")
+        result1 = client.call('hello')
+        result2 = client.call('sum', 1, 2)
+        result3 = client.call('nil')
+
+        self.assertEqual(result1, "world", "'hello' result is incorrect")
+        self.assertEqual(result2, 3, "'sum' result is incorrect")
+        self.assertIsNone(result3, "'nil' result is incorrect")
 
     def test_call_async(self):
         client = self.setup_env();
 
-        feture1 = client.call_async('hello')
-        feture2 = client.call_async('sum', 1, 2)
-        feture1.join()
-        feture2.join()
+        future1 = client.call_async('hello')
+        future2 = client.call_async('sum', 1, 2)
+        future3 = client.call_async('nil')
+        future1.join()
+        future2.join()
+        future3.join()
 
-        self.assertEqual(feture1.result, "world", "'hello' result is incorrect in call_async")
-        self.assertEqual(feture2.result, 3, "'sum' result is incorrect in call_async")
+        self.assertEqual(future1.result, "world", "'hello' result is incorrect in call_async")
+        self.assertEqual(future2.result, 3, "'sum' result is incorrect in call_async")
+        self.assertIsNone(future3.result, "'nil' result is incorrect in call_async")
 
     def test_notify(self):
         client = self.setup_env();
+
         result = True
         try:
             client.notify('hello')
             client.notify('sum', 1, 2)
+            client.notify('nil')
         except:
             result = False
+
         self.assertTrue(result)
 
+    def test_raise_error(self):
+        client = self.setup_env();
+        self.assertRaises(error.RPCError, lambda: client.call('raise_error'))
+
     def test_unknown_method(self):
         client = self.setup_env();
         self.assertRaises(error.RPCError, lambda: client.call('unknown', True))
@@ -76,6 +95,5 @@ def test_unknown_method(self):
             message = e.args[0]
             self.assertEqual(message, "'unknown' method not found", "Error message mismatched")
 
-
 if __name__ == '__main__':
     unittest.main()