Fix 2 call_function_test failures.

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

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

tests/call_function_test.py

@@ -28,19 +28,10 @@ def func_accepts_returns_context(context):
     return context
 
 
-class CallFunctionTest(unittest.TestCase):
+class CallFunctionTest(testlib.RouterMixin, unittest.TestCase):
-    @classmethod
+    def setUp(self):
-    def setUpClass(cls):
+        super(CallFunctionTest, self).setUp()
-        super(CallFunctionTest, cls).setUpClass()
+        self.local = self.router.local()
-        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_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)