Add call_function_test, fix various exception bugs.

pull/35/head
David Wilson 7 years ago
parent b4d4d8f041
commit 4e51f516a6

@ -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,

@ -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)

@ -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)
Loading…
Cancel
Save