From fa9def91f2e421be75d787d39f3d28b45f59c066 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Tue, 3 Oct 2017 15:54:55 +0530
Subject: [PATCH] Implement test_aborted_on_local_context_disconnect

---
 mitogen/core.py             |  8 +++++---
 tests/call_function_test.py | 14 ++++++++------
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/mitogen/core.py b/mitogen/core.py
index 19dca457..1d39f894 100644
--- a/mitogen/core.py
+++ b/mitogen/core.py
@@ -74,7 +74,8 @@ def _unpickle_call_error(s):
 
 
 class ChannelError(Error):
-    pass
+    remote_msg = 'Channel closed by remote end.'
+    local_msg = 'Channel closed by local end.'
 
 
 class StreamError(Error):
@@ -320,12 +321,12 @@ class Receiver(object):
         IOLOG.debug('%r.get() got %r', self, msg)
 
         if msg == _DEAD:
-            raise ChannelError('Channel closed by local end.')
+            raise ChannelError(ChannelError.local_msg)
 
         # Must occur off the broker thread.
         data = msg.unpickle()
         if data == _DEAD and self.raise_channelerror:
-            raise ChannelError('Channel closed by remote end.')
+            raise ChannelError(ChannelError.remote_msg)
 
         if isinstance(data, CallError):
             raise data
@@ -795,6 +796,7 @@ class Waker(BasicStream):
         Write a byte to the self-pipe, causing the IO multiplexer to wake up.
         Nothing is written if the current thread is the IO multiplexer thread.
         """
+        IOLOG.debug('%r.wake() [fd=%r]', self, self.transmit_side.fd)
         if threading.currentThread() != self._broker._thread and \
            self.transmit_side.fd:
             os.write(self.transmit_side.fd, ' ')
diff --git a/tests/call_function_test.py b/tests/call_function_test.py
index a3b81b2b..3dec3dc2 100644
--- a/tests/call_function_test.py
+++ b/tests/call_function_test.py
@@ -61,12 +61,14 @@ class CallFunctionTest(testlib.RouterMixin, testlib.TestCase):
     def test_returns_dead(self):
         assert mitogen.core._DEAD == self.local.call(func_returns_dead)
 
-    def test_aborted_on_context_disconnect(self):
-        assert 0, 'todo'
-
-    def test_aborted_on_context_hang_deadline(self):
-        # related: how to treat context after a function call hangs
-        assert 0, 'todo'
+    def test_aborted_on_local_context_disconnect(self):
+        stream = self.router._stream_by_id[self.local.context_id]
+        self.broker.stop_receive(stream)
+        recv = self.local.call_async(time.sleep, 120)
+        self.broker.defer(stream.on_disconnect, self.broker)
+        exc = self.assertRaises(mitogen.core.ChannelError,
+            lambda: recv.get())
+        self.assertEquals(exc[0], mitogen.core.ChannelError.local_msg)
 
     def test_aborted_on_local_broker_shutdown(self):
         assert 0, 'todo'