diff --git a/mitogen/core.py b/mitogen/core.py index 46be09f8..b631ad23 100644 --- a/mitogen/core.py +++ b/mitogen/core.py @@ -284,6 +284,7 @@ def _queue_interruptible_get(queue, timeout=None, block=True): class Receiver(object): notify = None + raise_channelerror = True def __init__(self, router, handle=None, persist=True, respondent=None): self.router = router @@ -319,7 +320,7 @@ class Receiver(object): # Must occur off the broker thread. data = msg.unpickle() - if data == _DEAD: + if data == _DEAD and self.raise_channelerror: raise ChannelError('Channel closed by remote end.') if isinstance(data, CallError): diff --git a/mitogen/master.py b/mitogen/master.py index 0a5eeca9..a4207a65 100644 --- a/mitogen/master.py +++ b/mitogen/master.py @@ -807,12 +807,14 @@ class Context(mitogen.core.Context): else: klass = None - return self.send_async( + recv = self.send_async( mitogen.core.Message.pickled( (fn.__module__, klass, fn.__name__, args, kwargs), handle=mitogen.core.CALL_FUNCTION, ) ) + recv.raise_channelerror = False + return recv def call(self, fn, *args, **kwargs): return self.call_async(fn, *args, **kwargs).get_data() diff --git a/tests/call_function_test.py b/tests/call_function_test.py index 6dd8d584..7f977a74 100644 --- a/tests/call_function_test.py +++ b/tests/call_function_test.py @@ -28,19 +28,10 @@ def func_accepts_returns_context(context): return context -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() +class CallFunctionTest(testlib.RouterMixin, unittest.TestCase): + def setUp(self): + super(CallFunctionTest, self).setUp() + self.local = self.router.local() def test_succeeds(self): assert 3 == self.local.call(function_that_adds_numbers, 1, 2) @@ -69,7 +60,7 @@ class CallFunctionTest(unittest.TestCase): except mitogen.core.StreamError, e: pass - assert e[0] == "cannot unpickle 'call_function_test'/'CrazyType'" + self.assertEquals(e[0], "cannot unpickle '__main__'/'CrazyType'") def test_returns_dead(self): assert mitogen.core._DEAD == self.local.call(func_returns_dead)