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)