From 300f8b2ff910a57e63b0d05f7332c29350623d15 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Tue, 23 Jul 2019 14:04:22 +0100
Subject: [PATCH] ansible: fixturize creation of MuxProcess

This relies on the previous commit resetting global variables.

Update clean_shutdown() to handle duplicate calls, due to tests
repeatedly installing it.
---
 ansible_mitogen/process.py             | 26 ++++++++---
 tests/ansible/tests/connection_test.py | 61 ++++++++++++++------------
 tests/testlib.py                       |  4 ++
 3 files changed, 58 insertions(+), 33 deletions(-)

diff --git a/ansible_mitogen/process.py b/ansible_mitogen/process.py
index 251acf30..a8827cb1 100644
--- a/ansible_mitogen/process.py
+++ b/ansible_mitogen/process.py
@@ -79,8 +79,15 @@ def clean_shutdown(sock):
     MuxProcess, debug logs may appear on the user's terminal *after* the prompt
     has been printed.
     """
-    sock.shutdown(socket.SHUT_WR)
+    try:
+        sock.shutdown(socket.SHUT_WR)
+    except socket.error:
+        # Already closed. This is possible when tests are running.
+        LOG.debug('clean_shutdown: ignoring duplicate call')
+        return
+
     sock.recv(1)
+    sock.close()
 
 
 def getenv_int(key, default=0):
@@ -154,8 +161,15 @@ class MuxProcess(object):
     #: forked WorkerProcesses to contact the MuxProcess
     unix_listener_path = None
 
-    #: Singleton.
-    _instance = None
+    @classmethod
+    def _reset(cls):
+        """
+        Used to clean up in unit tests.
+        """
+        assert cls.worker_sock is not None
+        cls.worker_sock.close()
+        cls.worker_sock = None
+        os.waitpid(cls.worker_pid, 0)
 
     @classmethod
     def start(cls, _init_logging=True):
@@ -178,7 +192,7 @@ class MuxProcess(object):
         mitogen.utils.setup_gil()
         cls.unix_listener_path = mitogen.unix.make_socket_path()
         cls.worker_sock, cls.child_sock = socket.socketpair()
-        atexit.register(lambda: clean_shutdown(cls.worker_sock))
+        atexit.register(clean_shutdown, cls.worker_sock)
         mitogen.core.set_cloexec(cls.worker_sock.fileno())
         mitogen.core.set_cloexec(cls.child_sock.fileno())
 
@@ -189,8 +203,8 @@ class MuxProcess(object):
             ansible_mitogen.logging.setup()
 
         cls.original_env = dict(os.environ)
-        cls.child_pid = os.fork()
-        if cls.child_pid:
+        cls.worker_pid = os.fork()
+        if cls.worker_pid:
             save_pid('controller')
             ansible_mitogen.logging.set_process_name('top')
             ansible_mitogen.affinity.policy.assign_controller()
diff --git a/tests/ansible/tests/connection_test.py b/tests/ansible/tests/connection_test.py
index 401cbe9e..d663ecc5 100644
--- a/tests/ansible/tests/connection_test.py
+++ b/tests/ansible/tests/connection_test.py
@@ -13,42 +13,29 @@ import ansible.errors
 import ansible.playbook.play_context
 
 import mitogen.core
+import mitogen.utils
+
 import ansible_mitogen.connection
 import ansible_mitogen.plugins.connection.mitogen_local
 import ansible_mitogen.process
-import testlib
-
 
-LOGGER_NAME = ansible_mitogen.target.LOG.name
-
-
-# TODO: fixtureize
-import mitogen.utils
-mitogen.utils.log_to_file()
-ansible_mitogen.process.MuxProcess.start(_init_logging=False)
-
-
-class OptionalIntTest(unittest2.TestCase):
-    func = staticmethod(ansible_mitogen.connection.optional_int)
-
-    def test_already_int(self):
-        self.assertEquals(0, self.func(0))
-        self.assertEquals(1, self.func(1))
-        self.assertEquals(-1, self.func(-1))
+import testlib
 
-    def test_is_string(self):
-        self.assertEquals(0, self.func("0"))
-        self.assertEquals(1, self.func("1"))
-        self.assertEquals(-1, self.func("-1"))
 
-    def test_is_none(self):
-        self.assertEquals(None, self.func(None))
+class MuxProcessMixin(object):
+    @classmethod
+    def setUpClass(cls):
+        #mitogen.utils.log_to_file()
+        ansible_mitogen.process.MuxProcess.start(_init_logging=False)
+        super(MuxProcessMixin, cls).setUpClass()
 
-    def test_is_junk(self):
-        self.assertEquals(None, self.func({1:2}))
+    @classmethod
+    def tearDownClass(cls):
+        super(MuxProcessMixin, cls).tearDownClass()
+        ansible_mitogen.process.MuxProcess._reset()
 
 
-class ConnectionMixin(object):
+class ConnectionMixin(MuxProcessMixin):
     klass = ansible_mitogen.plugins.connection.mitogen_local.Connection
 
     def make_connection(self):
@@ -70,6 +57,26 @@ class ConnectionMixin(object):
         super(ConnectionMixin, self).tearDown()
 
 
+class OptionalIntTest(unittest2.TestCase):
+    func = staticmethod(ansible_mitogen.connection.optional_int)
+
+    def test_already_int(self):
+        self.assertEquals(0, self.func(0))
+        self.assertEquals(1, self.func(1))
+        self.assertEquals(-1, self.func(-1))
+
+    def test_is_string(self):
+        self.assertEquals(0, self.func("0"))
+        self.assertEquals(1, self.func("1"))
+        self.assertEquals(-1, self.func("-1"))
+
+    def test_is_none(self):
+        self.assertEquals(None, self.func(None))
+
+    def test_is_junk(self):
+        self.assertEquals(None, self.func({1:2}))
+
+
 class PutDataTest(ConnectionMixin, unittest2.TestCase):
     def test_out_path(self):
         path = tempfile.mktemp(prefix='mitotest')
diff --git a/tests/testlib.py b/tests/testlib.py
index 7328ce09..762d0c6a 100644
--- a/tests/testlib.py
+++ b/tests/testlib.py
@@ -351,6 +351,10 @@ class TestCase(unittest2.TestCase):
             assert 0, "%s failed to reap subprocess %d (status %d)." % (
                 self, pid, status
             )
+
+        print()
+        print('Children of unit test process:')
+        os.system('ps uww --ppid ' + str(os.getpid()))
         assert 0, "%s leaked still-running subprocesses." % (self,)
 
     def tearDown(self):