From 4fa760cd21f5bd0d572a3e90f737ad1f20f0f906 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Sat, 10 Aug 2019 20:30:24 +0100
Subject: [PATCH] issue #613: add tests for all the weird shutdown methods

---
 docs/internals.rst   |   5 +++
 mitogen/parent.py    |   3 +-
 tests/router_test.py | 105 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 112 insertions(+), 1 deletion(-)

diff --git a/docs/internals.rst b/docs/internals.rst
index d062b6d9..40ea33df 100644
--- a/docs/internals.rst
+++ b/docs/internals.rst
@@ -321,6 +321,11 @@ These signals are used internally by Mitogen.
       - ``disconnect``
       - Fired on the Broker thread during shutdown (???)
 
+    * - :py:class:`mitogen.parent.Process`
+      - ``exit``
+      - Fired when :class:`mitogen.parent.Reaper` detects subprocess has fully
+        exitted.
+
     * - :py:class:`mitogen.core.Broker`
       - ``shutdown``
       - Fired after Broker.shutdown() is called.
diff --git a/mitogen/parent.py b/mitogen/parent.py
index 7a94c2b0..983df829 100644
--- a/mitogen/parent.py
+++ b/mitogen/parent.py
@@ -2331,7 +2331,7 @@ class Router(mitogen.core.Router):
         directly connected.
         """
         stream = self.stream_by_id(context)
-        if stream.protocol.remote_id != context.context_id:
+        if stream is None or stream.protocol.remote_id != context.context_id:
             return
 
         l = mitogen.core.Latch()
@@ -2589,6 +2589,7 @@ class Reaper(object):
         status = self.proc.poll()
         if status is not None:
             LOG.debug('%r: %s', self.proc, returncode_to_str(status))
+            mitogen.core.fire(self.proc, 'exit')
             self._remove_timer()
             return
 
diff --git a/tests/router_test.py b/tests/router_test.py
index dba56c9b..a9ad5ae1 100644
--- a/tests/router_test.py
+++ b/tests/router_test.py
@@ -1,3 +1,5 @@
+import errno
+import os
 import sys
 import time
 import zlib
@@ -425,5 +427,108 @@ class EgressIdsTest(testlib.RouterMixin, testlib.TestCase):
         ]))
 
 
+class ShutdownTest(testlib.RouterMixin, testlib.TestCase):
+    # 613: tests for all the weird shutdown() variants we ended up with.
+
+    def test_shutdown_wait_false(self):
+        l1 = self.router.local()
+        pid = l1.call(os.getpid)
+
+        conn = self.router.stream_by_id(l1.context_id).conn
+        exitted = mitogen.core.Latch()
+        mitogen.core.listen(conn.proc, 'exit', exitted.put)
+
+        l1.shutdown(wait=False)
+        exitted.get()
+
+        e = self.assertRaises(OSError,
+            lambda: os.waitpid(pid, 0))
+        self.assertEquals(e.args[0], errno.ECHILD)
+
+        e = self.assertRaises(mitogen.core.ChannelError,
+            lambda: l1.call(os.getpid))
+        self.assertEquals(e.args[0], mitogen.core.Router.no_route_msg % (
+            l1.context_id,
+            mitogen.context_id,
+        ))
+
+    def test_shutdown_wait_true(self):
+        l1 = self.router.local()
+        pid = l1.call(os.getpid)
+
+        conn = self.router.stream_by_id(l1.context_id).conn
+        exitted = mitogen.core.Latch()
+        mitogen.core.listen(conn.proc, 'exit', exitted.put)
+
+        l1.shutdown(wait=True)
+        exitted.get()
+
+        e = self.assertRaises(OSError,
+            lambda: os.waitpid(pid, 0))
+        self.assertEquals(e.args[0], errno.ECHILD)
+
+        e = self.assertRaises(mitogen.core.ChannelError,
+            lambda: l1.call(os.getpid))
+        self.assertEquals(e.args[0], mitogen.core.Router.no_route_msg % (
+            l1.context_id,
+            mitogen.context_id,
+        ))
+
+    def test_disconnect_invalid_context(self):
+        self.router.disconnect(
+            mitogen.core.Context(self.router, 1234)
+        )
+
+    def test_disconnect_valid_context(self):
+        l1 = self.router.local()
+        pid = l1.call(os.getpid)
+
+        strm = self.router.stream_by_id(l1.context_id)
+
+        exitted = mitogen.core.Latch()
+        mitogen.core.listen(strm.conn.proc, 'exit', exitted.put)
+        self.router.disconnect_stream(strm)
+        exitted.get()
+
+        e = self.assertRaises(OSError,
+            lambda: os.waitpid(pid, 0))
+        self.assertEquals(e.args[0], errno.ECHILD)
+
+        e = self.assertRaises(mitogen.core.ChannelError,
+            lambda: l1.call(os.getpid))
+        self.assertEquals(e.args[0], mitogen.core.Router.no_route_msg % (
+            l1.context_id,
+            mitogen.context_id,
+        ))
+
+    def test_disconnet_all(self):
+        l1 = self.router.local()
+        l2 = self.router.local()
+
+        pids = [l1.call(os.getpid), l2.call(os.getpid)]
+
+        exitted = mitogen.core.Latch()
+        for ctx in l1, l2:
+            strm = self.router.stream_by_id(ctx.context_id)
+            mitogen.core.listen(strm.conn.proc, 'exit', exitted.put)
+
+        self.router.disconnect_all()
+        exitted.get()
+        exitted.get()
+
+        for pid in pids:
+            e = self.assertRaises(OSError,
+                lambda: os.waitpid(pid, 0))
+            self.assertEquals(e.args[0], errno.ECHILD)
+
+        for ctx in l1, l2:
+            e = self.assertRaises(mitogen.core.ChannelError,
+                lambda: ctx.call(os.getpid))
+            self.assertEquals(e.args[0], mitogen.core.Router.no_route_msg % (
+                ctx.context_id,
+                mitogen.context_id,
+            ))
+
+
 if __name__ == '__main__':
     unittest2.main()