import Queue import StringIO import logging import subprocess import time import unittest2 import testlib import mitogen.master import mitogen.utils def ping(): return True @mitogen.core.takes_router def return_router_max_message_size(router): return router.max_message_size def send_n_sized_reply(sender, n): sender.send(' ' * n) return 123 class SourceVerifyTest(testlib.RouterMixin, unittest2.TestCase): def setUp(self): super(SourceVerifyTest, self).setUp() # Create some children, ping them, and store what their messages look # like so we can mess with them later. self.child1 = self.router.fork() self.child1_msg = self.child1.call_async(ping).get() self.child1_stream = self.router._stream_by_id[self.child1.context_id] self.child2 = self.router.fork() self.child2_msg = self.child2.call_async(ping).get() self.child2_stream = self.router._stream_by_id[self.child2.context_id] def test_bad_auth_id(self): # Deliver a message locally from child2, but using child1's stream. log = testlib.LogCapturer() log.start() # Used to ensure the message was dropped rather than routed after the # error is logged. recv = mitogen.core.Receiver(self.router) self.child2_msg.handle = recv.handle self.broker.defer(self.router._async_route, self.child2_msg, stream=self.child1_stream) # Wait for IO loop to finish everything above. self.sync_with_broker() # Ensure message wasn't forwarded. self.assertTrue(recv.empty()) # Ensure error was logged. expect = 'bad auth_id: got %d via' % (self.child2_msg.auth_id,) self.assertTrue(expect in log.stop()) def test_bad_src_id(self): # Deliver a message locally from child2 with the correct auth_id, but # the wrong src_id. log = testlib.LogCapturer() log.start() # Used to ensure the message was dropped rather than routed after the # error is logged. recv = mitogen.core.Receiver(self.router) self.child2_msg.handle = recv.handle self.child2_msg.src_id = self.child1.context_id self.broker.defer(self.router._async_route, self.child2_msg, self.child2_stream) # Wait for IO loop to finish everything above. self.sync_with_broker() # Ensure message wasn't forwarded. self.assertTrue(recv.empty()) # Ensure error was lgoged. expect = 'bad src_id: got %d via' % (self.child1_msg.src_id,) self.assertTrue(expect in log.stop()) class PolicyTest(testlib.RouterMixin, testlib.TestCase): def test_allow_any(self): # This guy gets everything. recv = mitogen.core.Receiver(self.router) recv.to_sender().send(123) self.sync_with_broker() self.assertFalse(recv.empty()) self.assertEquals(123, recv.get().unpickle()) def test_refuse_all(self): # Deliver a message locally from child2 with the correct auth_id, but # the wrong src_id. log = testlib.LogCapturer() log.start() # This guy never gets anything. recv = mitogen.core.Receiver( router=self.router, policy=(lambda msg, stream: False), ) # This guy becomes the reply_to of our refused message. reply_target = mitogen.core.Receiver(self.router) # Send the message. self.router.route( mitogen.core.Message( dst_id=mitogen.context_id, handle=recv.handle, reply_to=reply_target.handle, ) ) # Wait for IO loop. self.sync_with_broker() # Verify log. expect = '%r: policy refused message: ' % (self.router,) self.assertTrue(expect in log.stop()) # Verify message was not delivered. self.assertTrue(recv.empty()) # Verify CallError received by reply_to target. e = self.assertRaises(mitogen.core.CallError, lambda: reply_target.get().unpickle()) self.assertEquals(e[0], self.router.refused_msg) class CrashTest(testlib.BrokerMixin, unittest2.TestCase): # This is testing both Broker's ability to crash nicely, and Router's # ability to respond to the crash event. klass = mitogen.master.Router def _naughty(self): raise ValueError('eek') def test_shutdown(self): router = self.klass(self.broker) sem = mitogen.core.Latch() router.add_handler(sem.put) log = testlib.LogCapturer('mitogen') log.start() # Force a crash and ensure it wakes up. self.broker._loop_once = self._naughty self.broker.defer(lambda: None) # sem should have received _DEAD. self.assertEquals(mitogen.core._DEAD, sem.get()) # Ensure it was logged. expect = '_broker_main() crashed' self.assertTrue(expect in log.stop()) class AddHandlerTest(unittest2.TestCase): klass = mitogen.master.Router def test_invoked_at_shutdown(self): router = self.klass() queue = Queue.Queue() handle = router.add_handler(queue.put) router.broker.shutdown() self.assertEquals(queue.get(timeout=5), mitogen.core._DEAD) class MessageSizeTest(testlib.BrokerMixin, unittest2.TestCase): klass = mitogen.master.Router def test_local_exceeded(self): router = self.klass(broker=self.broker, max_message_size=4096) recv = mitogen.core.Receiver(router) logs = testlib.LogCapturer() logs.start() sem = mitogen.core.Latch() router.route(mitogen.core.Message.pickled(' '*8192)) router.broker.defer(sem.put, ' ') # wlil always run after _async_route sem.get() expect = 'message too large (max 4096 bytes)' self.assertTrue(expect in logs.stop()) def test_remote_configured(self): router = self.klass(broker=self.broker, max_message_size=4096) remote = router.fork() size = remote.call(return_router_max_message_size) self.assertEquals(size, 4096) def test_remote_exceeded(self): # Ensure new contexts receive a router with the same value. router = self.klass(broker=self.broker, max_message_size=4096) recv = mitogen.core.Receiver(router) logs = testlib.LogCapturer() logs.start() remote = router.fork() remote.call(send_n_sized_reply, recv.to_sender(), 8192) expect = 'message too large (max 4096 bytes)' self.assertTrue(expect in logs.stop()) if __name__ == '__main__': unittest2.main()