import mock import textwrap import subprocess import sys import unittest2 import mitogen.master import testlib import plain_old_module import simple_pkg.a class NeutralizeMainTest(testlib.RouterMixin, testlib.TestCase): klass = mitogen.master.ModuleResponder def call(self, *args, **kwargs): router = mock.Mock() return self.klass(router).neutralize_main(*args, **kwargs) def test_missing_exec_guard(self): path = testlib.data_path('main_with_no_exec_guard.py') args = [sys.executable, path] proc = subprocess.Popen(args, stderr=subprocess.PIPE) _, stderr = proc.communicate() self.assertEquals(1, proc.returncode) expect = self.klass.main_guard_msg % (path,) self.assertTrue(expect in stderr.decode()) HAS_MITOGEN_MAIN = mitogen.core.b( textwrap.dedent(""" herp derp def myprog(): pass @mitogen.main(maybe_some_option=True) def main(router): pass """) ) def test_mitogen_main(self): untouched = self.call("derp.py", self.HAS_MITOGEN_MAIN) self.assertEquals(untouched, self.HAS_MITOGEN_MAIN) HAS_EXEC_GUARD = mitogen.core.b( textwrap.dedent(""" herp derp def myprog(): pass def main(): pass if __name__ == '__main__': main() """) ) def test_exec_guard(self): touched = self.call("derp.py", self.HAS_EXEC_GUARD) bits = touched.decode().split() self.assertEquals(bits[-3:], ['def', 'main():', 'pass']) class GoodModulesTest(testlib.RouterMixin, testlib.TestCase): def test_plain_old_module(self): # The simplest case: a top-level module with no interesting imports or # package machinery damage. context = self.router.local() self.assertEquals(256, context.call(plain_old_module.pow, 2, 8)) self.assertEquals(1, self.router.responder.get_module_count) self.assertEquals(1, self.router.responder.good_load_module_count) self.assertLess(300, self.router.responder.good_load_module_size) def test_simple_pkg(self): # Ensure success of a simple package containing two submodules, one of # which imports the other. context = self.router.local() self.assertEquals(3, context.call(simple_pkg.a.subtract_one_add_two, 2)) self.assertEquals(2, self.router.responder.get_module_count) self.assertEquals(3, self.router.responder.good_load_module_count) self.assertEquals(0, self.router.responder.bad_load_module_count) self.assertLess(450, self.router.responder.good_load_module_size) def test_self_contained_program(self): # Ensure a program composed of a single script can be imported # successfully. args = [sys.executable, testlib.data_path('self_contained_program.py')] output = testlib.subprocess__check_output(args).decode() self.assertEquals(output, "['__main__', 50]\n") class BrokenModulesTest(testlib.TestCase): def test_obviously_missing(self): # 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() router.stream_by_id = lambda n: stream msg = mitogen.core.Message( data=mitogen.core.b('non_existent_module'), reply_to=50, ) msg.router = router responder = mitogen.master.ModuleResponder(router) responder._on_get_module(msg) self.assertEquals(1, len(router._async_route.mock_calls)) self.assertEquals(1, responder.get_module_count) self.assertEquals(0, responder.good_load_module_count) self.assertEquals(0, responder.good_load_module_size) self.assertEquals(1, responder.bad_load_module_count) call = router._async_route.mock_calls[0] msg, = call[1] 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 # __path__ subsequently ends up empty, which prevents pkgutil from # finding its submodules. After ansible.compat.six is initialized in # the parent, attempts to execute six/__init__.py on the slave will # 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() router.stream_by_id = lambda n: stream msg = mitogen.core.Message( data=mitogen.core.b('six_brokenpkg._six'), reply_to=50, ) msg.router = router responder = mitogen.master.ModuleResponder(router) responder._on_get_module(msg) self.assertEquals(1, len(router._async_route.mock_calls)) self.assertEquals(1, responder.get_module_count) self.assertEquals(0, responder.good_load_module_count) self.assertEquals(0, responder.good_load_module_size) self.assertEquals(1, responder.bad_load_module_count) call = router._async_route.mock_calls[0] msg, = call[1] self.assertEquals(mitogen.core.LOAD_MODULE, msg.handle) self.assertIsInstance(msg.unpickle(), tuple) class ForwardTest(testlib.RouterMixin, testlib.TestCase): def test_stats(self): # Forwarding stats broken because forwarding is broken. See #469. c1 = self.router.local() c2 = self.router.fork(via=c1) self.assertEquals(256, c2.call(plain_old_module.pow, 2, 8)) self.assertEquals(3, self.router.responder.get_module_count) self.assertEquals(3, self.router.responder.good_load_module_count) self.assertLess(23000, self.router.responder.good_load_module_size) class BlacklistTest(testlib.TestCase): @unittest2.skip('implement me') def test_whitelist_no_blacklist(self): assert 0 @unittest2.skip('implement me') def test_whitelist_has_blacklist(self): assert 0 @unittest2.skip('implement me') def test_blacklist_no_whitelist(self): assert 0 @unittest2.skip('implement me') def test_blacklist_has_whitelist(self): assert 0 if __name__ == '__main__': unittest2.main()