From d527f28c2f9a60809b98916c22bef33fbbdc83b4 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Tue, 5 Dec 2017 02:03:03 +0530
Subject: [PATCH] WIP: exit status changes/hacks for fakessh

---
 mitogen/fakessh.py    | 39 ++++++++++++++++++++++++++++++++++-----
 tests/fakessh_test.py | 17 +++++++++++++++--
 2 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/mitogen/fakessh.py b/mitogen/fakessh.py
index 25f35149..a286e2d9 100644
--- a/mitogen/fakessh.py
+++ b/mitogen/fakessh.py
@@ -163,12 +163,39 @@ class Process(object):
         self.stdin = mitogen.core.Sender(dest, arg[1])
         self.router.broker.start_receive(self.pump)
 
+        if self.proc:
+            pmon = mitogen.master.ProcessMonitor.instance()
+            pmon.add(self.proc.pid, self._on_proc_exit)
+
+    exitted = False
+
     def _on_exit(self, msg, arg):
-        LOG.debug('on_exit: proc = %r', self.proc)
+        LOG.debug('on_exit: proc = %r, arg = %r', self.proc, arg)
         if self.proc:
-            self.proc.terminate()
+            try:
+                self.proc.terminate()
+            except OSError:
+                pass
         else:
-            self.router.broker.shutdown()
+            #self.router.broker.shutdown()
+            msg = (
+                mitogen.core.Message.pickled(
+                    ('sys', None, 'exit', (arg or 0,), {}),
+                    handle=mitogen.core.CALL_FUNCTION,
+                    src_id=mitogen.parent_ids[0],
+                    dst_id=mitogen.context_id,
+                )
+            )
+            print 'HAHA'
+            print msg
+            self.router.route(
+                mitogen.core.Message.pickled(
+                    ('sys', None, 'exit', (arg or 0,), {}),
+                    handle=mitogen.core.CALL_FUNCTION,
+                    src_id=mitogen.parent_ids[0],
+                    dst_id=mitogen.context_id,
+                )
+            )
 
     def _on_pump_receive(self, s):
         IOLOG.info('%r._on_pump_receive(len %d)', self, len(s))
@@ -205,8 +232,10 @@ def _start_slave(src_id, cmdline, router):
     """
     LOG.debug('_start_slave(%r, %r)', router, cmdline)
 
-    proc = subprocess.Popen(
-        cmdline,
+    # Must be constructed from main thread due to signal use.
+    pmon = mitogen.master.ProcessMonitor.instance()
+
+    proc = subprocess.Popen(cmdline,
         # SSH server always uses user's shell.
         shell=True,
         # SSH server always executes new commands in the user's HOME.
diff --git a/tests/fakessh_test.py b/tests/fakessh_test.py
index 6e28be1d..3665a0a8 100644
--- a/tests/fakessh_test.py
+++ b/tests/fakessh_test.py
@@ -9,6 +9,19 @@ import mitogen.fakessh
 import testlib
 
 
+import logging
+logging.basicConfig(level=logging.DEBUG)
+class ExitStatusTest(testlib.RouterMixin, unittest2.TestCase):
+    # Ensure child's exit status is propagated to the parent.
+    def test_exit_0(self):
+        local = self.router.local()
+
+        exit_status = mitogen.fakessh.run(local, self.router, [
+            'bash', '-c', 'exec ssh foo exit 0'
+        ])
+        self.assertEquals(0, exit_status)
+
+
 class RsyncTest(testlib.DockerMixin, unittest2.TestCase):
     def test_rsync_from_master(self):
         context = self.docker_ssh_any()
@@ -17,7 +30,7 @@ class RsyncTest(testlib.DockerMixin, unittest2.TestCase):
             context.call(shutil.rmtree, '/tmp/data')
 
         return_code = mitogen.fakessh.run(context, self.router, [
-            'rsync', '--progress', '-vvva',
+            'rsync', '-a',
             testlib.data_path('.'), 'target:/tmp/data'
         ])
 
@@ -49,7 +62,7 @@ class RsyncTest(testlib.DockerMixin, unittest2.TestCase):
             webapp_acct.call(shutil.rmtree, dest_path)
 
         return_code = pubkey_acct.call(mitogen.fakessh.run, webapp_acct, args=[
-            'rsync', '--progress', '-vvva', '.ssh/', 'target:' + dest_path
+            'rsync', '-a', '.ssh/', 'target:' + dest_path
         ])
 
         self.assertEqual(return_code, 0)