From a89c20e54ef4371bd93faef4355c6e8b832805fe Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Sun, 17 Sep 2017 20:33:43 +0530
Subject: [PATCH] Fixup some more tests.

---
 tests/data/self_contained_program.py | 14 +++-----
 tests/responder_test.py              | 48 +++++++++++++++++-----------
 tests/testlib.py                     |  6 +++-
 3 files changed, 39 insertions(+), 29 deletions(-)

diff --git a/tests/data/self_contained_program.py b/tests/data/self_contained_program.py
index 35a42a3e..52be442e 100644
--- a/tests/data/self_contained_program.py
+++ b/tests/data/self_contained_program.py
@@ -9,14 +9,10 @@ def repr_stuff():
     return repr([__name__, 50])
 
 
-def main():
-    broker = mitogen.master.Broker()
-    try:
-        context = mitogen.master.connect(broker)
-        print context.call(repr_stuff)
-    finally:
-        broker.shutdown()
-        broker.join()
+def main(router):
+    context = router.local()
+    print context.call(repr_stuff)
 
 if __name__ == '__main__' and mitogen.is_master:
-    main()
+    import mitogen.utils
+    mitogen.utils.run_with_router(main)
diff --git a/tests/responder_test.py b/tests/responder_test.py
index 6e6aede2..80729a21 100644
--- a/tests/responder_test.py
+++ b/tests/responder_test.py
@@ -11,17 +11,17 @@ import plain_old_module
 import simple_pkg.a
 
 
-class GoodModulesTest(testlib.BrokerMixin, unittest.TestCase):
+class GoodModulesTest(testlib.RouterMixin, unittest.TestCase):
     def test_plain_old_module(self):
         # The simplest case: a top-level module with no interesting imports or
         # package machinery damage.
-        context = mitogen.master.connect(self.broker)
+        context = self.router.local()
         self.assertEquals(256, context.call(plain_old_module.pow, 2, 8))
 
     def test_simple_pkg(self):
         # Ensure success of a simple package containing two submodules, one of
         # which imports the other.
-        context = mitogen.master.connect(self.broker)
+        context = self.router.local()
         self.assertEquals(3,
             context.call(simple_pkg.a.subtract_one_add_two, 2))
 
@@ -38,15 +38,20 @@ class BrokenModulesTest(unittest.TestCase):
         # Ensure we don't crash in the case of a module legitimately being
         # unavailable. Should never happen in the real world.
 
-        context = mock.Mock()
-        responder = mitogen.master.ModuleResponder(context)
-        responder.get_module((50, 'non_existent_module'))
-        self.assertEquals(1, len(context.enqueue.mock_calls))
+        router = mock.Mock()
+        responder = mitogen.master.ModuleResponder(router)
+        responder._on_get_module(
+            mitogen.core.Message(
+                data='non_existent_module',
+                reply_to=50,
+            )
+        )
+        self.assertEquals(1, len(router.route.mock_calls))
 
-        call = context.enqueue.mock_calls[0]
-        reply_to, data = call[1]
-        self.assertEquals(50, reply_to)
-        self.assertTrue(data is None)
+        call = router.route.mock_calls[0]
+        msg, = call[1]
+        self.assertEquals(50, msg.handle)
+        self.assertTrue(msg.unpickle() is None)
 
     def test_ansible_six_messed_up_path(self):
         # The copy of six.py shipped with Ansible appears in a package whose
@@ -56,12 +61,17 @@ class BrokenModulesTest(unittest.TestCase):
         # cause an attempt to request ansible.compat.six._six from the master.
         import six_brokenpkg
 
-        context = mock.Mock()
-        responder = mitogen.master.ModuleResponder(context)
-        responder.get_module((50, 'six_brokenpkg._six'))
-        self.assertEquals(1, len(context.enqueue.mock_calls))
+        router = mock.Mock()
+        responder = mitogen.master.ModuleResponder(router)
+        responder._on_get_module(
+            mitogen.core.Message(
+                data='six_brokenpkg._six',
+                reply_to=50,
+            )
+        )
+        self.assertEquals(1, len(router.route.mock_calls))
 
-        call = context.enqueue.mock_calls[0]
-        reply_to, data = call[1]
-        self.assertEquals(50, reply_to)
-        self.assertTrue(isinstance(data, tuple))
+        call = router.route.mock_calls[0]
+        msg, = call[1]
+        self.assertEquals(50, msg.handle)
+        self.assertTrue(isinstance(msg.unpickle(), tuple))
diff --git a/tests/testlib.py b/tests/testlib.py
index 9d49b5a2..5f3f62d1 100644
--- a/tests/testlib.py
+++ b/tests/testlib.py
@@ -21,7 +21,11 @@ def set_debug():
 
 
 def data_path(suffix):
-    return os.path.join(DATA_DIR, suffix)
+    path = os.path.join(DATA_DIR, suffix)
+    if path.endswith('.key'):
+        # SSH is funny about private key permissions.
+        os.chmod(path, 0600)
+    return path
 
 
 class DockerizedSshDaemon(object):