From ce6fb05d8755fcd27c03ff8ea968aa877b796e12 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Mon, 16 Apr 2018 16:58:44 +0100
Subject: [PATCH] tests: 'fix' responder test.

Needs a complete rewrite, but this will do for now.
---
 mitogen/master.py                    |  2 +-
 tests/data/self_contained_program.py |  7 ++---
 tests/responder_test.py              | 41 ++++++++++++++++++----------
 3 files changed, 29 insertions(+), 21 deletions(-)

diff --git a/mitogen/master.py b/mitogen/master.py
index cb85f33e..9025810a 100644
--- a/mitogen/master.py
+++ b/mitogen/master.py
@@ -646,7 +646,7 @@ class ModuleResponder(object):
 
         except Exception:
             LOG.debug('While importing %r', fullname, exc_info=True)
-            msg.reply((fullname, None, None, None, []),
+            msg.reply((fullname, None, None, None, ()),
                       handle=mitogen.core.LOAD_MODULE)
 
 
diff --git a/tests/data/self_contained_program.py b/tests/data/self_contained_program.py
index 52be442e..9ef18acb 100644
--- a/tests/data/self_contained_program.py
+++ b/tests/data/self_contained_program.py
@@ -2,17 +2,14 @@
 I am a self-contained program!
 """
 
-import mitogen.master
+import mitogen
 
 
 def repr_stuff():
     return repr([__name__, 50])
 
 
+@mitogen.main()
 def main(router):
     context = router.local()
     print context.call(repr_stuff)
-
-if __name__ == '__main__' and mitogen.is_master:
-    import mitogen.utils
-    mitogen.utils.run_with_router(main)
diff --git a/tests/responder_test.py b/tests/responder_test.py
index b558b265..837beb3e 100644
--- a/tests/responder_test.py
+++ b/tests/responder_test.py
@@ -39,20 +39,26 @@ class BrokenModulesTest(unittest2.TestCase):
         # Ensure we don't crash in the case of a module legitimately being
         # unavailable. Should never happen in the real world.
 
+        stream = mock.Mock()
+        stream.sent_modules = set()
         router = mock.Mock()
-        responder = mitogen.master.ModuleResponder(router)
-        responder._on_get_module(
-            mitogen.core.Message(
-                data='non_existent_module',
-                reply_to=50,
-            )
+        router.stream_by_id = lambda n: stream
+
+        msg = mitogen.core.Message(
+            data='non_existent_module',
+            reply_to=50,
         )
+        msg.router = router
+
+        responder = mitogen.master.ModuleResponder(router)
+        responder._on_get_module(msg)
         self.assertEquals(1, len(router.route.mock_calls))
 
         call = router.route.mock_calls[0]
         msg, = call[1]
-        self.assertEquals(50, msg.handle)
-        self.assertIsNone(msg.unpickle())
+        self.assertEquals(mitogen.core.LOAD_MODULE, msg.handle)
+        self.assertEquals(('non_existent_module', None, None, None, ()),
+                          msg.unpickle())
 
     def test_ansible_six_messed_up_path(self):
         # The copy of six.py shipped with Ansible appears in a package whose
@@ -62,19 +68,24 @@ class BrokenModulesTest(unittest2.TestCase):
         # cause an attempt to request ansible.compat.six._six from the master.
         import six_brokenpkg
 
+        stream = mock.Mock()
+        stream.sent_modules = set()
         router = mock.Mock()
-        responder = mitogen.master.ModuleResponder(router)
-        responder._on_get_module(
-            mitogen.core.Message(
-                data='six_brokenpkg._six',
-                reply_to=50,
-            )
+        router.stream_by_id = lambda n: stream
+
+        msg = mitogen.core.Message(
+            data='six_brokenpkg._six',
+            reply_to=50,
         )
+        msg.router = router
+
+        responder = mitogen.master.ModuleResponder(router)
+        responder._on_get_module(msg)
         self.assertEquals(1, len(router.route.mock_calls))
 
         call = router.route.mock_calls[0]
         msg, = call[1]
-        self.assertEquals(50, msg.handle)
+        self.assertEquals(mitogen.core.LOAD_MODULE, msg.handle)
         self.assertIsInstance(msg.unpickle(), tuple)