From 4e51f516a6159504a658c09095389be23c80e0bf Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Sat, 16 Sep 2017 15:27:30 +0530
Subject: [PATCH] Add call_function_test, fix various exception bugs.

---
 mitogen/core.py             |  4 +--
 mitogen/master.py           |  5 +--
 tests/call_function_test.py | 67 +++++++++++++++++++++++++++++++++++++
 3 files changed, 72 insertions(+), 4 deletions(-)
 create mode 100644 tests/call_function_test.py

diff --git a/mitogen/core.py b/mitogen/core.py
index 9db9f773..9e05fb80 100644
--- a/mitogen/core.py
+++ b/mitogen/core.py
@@ -51,7 +51,7 @@ class Error(Exception):
     def __init__(self, fmt, *args):
         if args:
             fmt %= args
-        Exception.__init__(self, fmt % args)
+        Exception.__init__(self, fmt)
 
 
 class CallError(Error):
@@ -544,7 +544,7 @@ class Stream(BasicStream):
         if len(self._input_buf) < self.HEADER_LEN:
             return False
 
-        msg = Message()
+        msg = self.message_class()
         (msg.dst_id, msg.src_id,
          msg.handle, msg.reply_to, msg_len) = struct.unpack(
             self.HEADER_FMT,
diff --git a/mitogen/master.py b/mitogen/master.py
index 237fc680..c694dfc7 100644
--- a/mitogen/master.py
+++ b/mitogen/master.py
@@ -325,8 +325,9 @@ class Message(mitogen.core.Message):
         `StreamError` if the module is not whitelisted."""
         if (module_name, class_name) not in PERMITTED_CLASSES:
             raise mitogen.core.StreamError(
-                '%r attempted to unpickle %r in module %r',
-                self._context, class_name, module_name)
+                'attempted to unpickle %r in module %r',
+                class_name, module_name
+            )
         return getattr(sys.modules[module_name], class_name)
 
 
diff --git a/tests/call_function_test.py b/tests/call_function_test.py
new file mode 100644
index 00000000..c28a4c94
--- /dev/null
+++ b/tests/call_function_test.py
@@ -0,0 +1,67 @@
+
+import unittest
+import mitogen.core
+import mitogen.master
+
+
+class CrazyType(object):
+    pass
+
+
+def function_that_fails():
+    raise ValueError('exception text')
+
+
+def func_with_bad_return_value():
+    return CrazyType()
+
+
+def func_returns_dead():
+    return mitogen.core._DEAD
+
+
+class CallFunctionTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        super(CallFunctionTest, cls).setUpClass()
+        cls.broker = mitogen.master.Broker()
+        cls.router = mitogen.master.Router(cls.broker)
+        cls.local = cls.router.local()
+
+    @classmethod
+    def tearDownClass(cls):
+        super(CallFunctionTest, cls).tearDownClass()
+        cls.broker.shutdown()
+        cls.broker.join()
+
+    def test_call_crashes(self):
+        try:
+            self.local.call(function_that_fails)
+            assert 0, 'call didnt fail'
+        except mitogen.core.CallError, e:
+            pass
+
+        s = str(e)
+        etype, _, s = s.partition(': ')
+        assert etype == 'exceptions.ValueError'
+
+        msg, _, s = s.partition('\n')
+        assert msg == 'exception text'
+
+        # Traceback
+        assert len(s) > 0
+
+    def test_bad_return_value(self):
+        try:
+            self.local.call(func_with_bad_return_value)
+            assert 0, 'call didnt fail'
+        except mitogen.core.StreamError, e:
+            pass
+
+        assert e[0] == (
+            "attempted to unpickle 'CrazyType' "
+            "in module 'call_function_test'"
+        )
+
+    def test_call_returns_dead(self):
+        assert mitogen.core._DEAD == self.local.call(func_returns_dead)