391 lines
13 KiB
Python
391 lines
13 KiB
Python
import time
|
|
import zlib
|
|
|
|
import unittest2
|
|
|
|
import testlib
|
|
import mitogen.master
|
|
import mitogen.parent
|
|
import mitogen.utils
|
|
|
|
try:
|
|
import Queue
|
|
except ImportError:
|
|
import queue as Queue
|
|
|
|
|
|
def ping():
|
|
return True
|
|
|
|
|
|
@mitogen.core.takes_router
|
|
def ping_context(other, router):
|
|
other = mitogen.parent.Context(router, other.context_id)
|
|
other.call(ping)
|
|
|
|
|
|
@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, testlib.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.local()
|
|
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.local()
|
|
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,
|
|
in_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 %r 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.
|
|
self.assertTrue(self.router.refused_msg in log.stop())
|
|
|
|
# Verify message was not delivered.
|
|
self.assertTrue(recv.empty())
|
|
|
|
# Verify CallError received by reply_to target.
|
|
e = self.assertRaises(mitogen.core.ChannelError,
|
|
lambda: reply_target.get().unpickle())
|
|
self.assertEquals(e.args[0], self.router.refused_msg)
|
|
|
|
|
|
class CrashTest(testlib.BrokerMixin, testlib.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 message.
|
|
self.assertTrue(sem.get().is_dead)
|
|
|
|
# Ensure it was logged.
|
|
expect = '_broker_main() crashed'
|
|
self.assertTrue(expect in log.stop())
|
|
|
|
self.broker.join()
|
|
|
|
|
|
class AddHandlerTest(testlib.TestCase):
|
|
klass = mitogen.master.Router
|
|
|
|
def test_dead_message_sent_at_shutdown(self):
|
|
router = self.klass()
|
|
queue = Queue.Queue()
|
|
handle = router.add_handler(queue.put)
|
|
router.broker.shutdown()
|
|
self.assertTrue(queue.get(timeout=5).is_dead)
|
|
router.broker.join()
|
|
|
|
def test_cannot_double_register(self):
|
|
router = self.klass()
|
|
try:
|
|
router.add_handler((lambda: None), handle=1234)
|
|
e = self.assertRaises(mitogen.core.Error,
|
|
lambda: router.add_handler((lambda: None), handle=1234))
|
|
self.assertEquals(router.duplicate_handle_msg, e.args[0])
|
|
router.del_handler(1234)
|
|
finally:
|
|
router.broker.shutdown()
|
|
router.broker.join()
|
|
|
|
def test_can_reregister(self):
|
|
router = self.klass()
|
|
try:
|
|
router.add_handler((lambda: None), handle=1234)
|
|
router.del_handler(1234)
|
|
router.add_handler((lambda: None), handle=1234)
|
|
router.del_handler(1234)
|
|
finally:
|
|
router.broker.shutdown()
|
|
router.broker.join()
|
|
|
|
|
|
class MyselfTest(testlib.RouterMixin, testlib.TestCase):
|
|
def test_myself(self):
|
|
myself = self.router.myself()
|
|
self.assertEquals(myself.context_id, mitogen.context_id)
|
|
# TODO: context should know its own name too.
|
|
self.assertEquals(myself.name, 'self')
|
|
|
|
|
|
class MessageSizeTest(testlib.BrokerMixin, testlib.TestCase):
|
|
klass = mitogen.master.Router
|
|
|
|
def test_local_exceeded(self):
|
|
router = self.klass(broker=self.broker, max_message_size=4096)
|
|
|
|
logs = testlib.LogCapturer()
|
|
logs.start()
|
|
|
|
# Send message and block for one IO loop, so _async_route can run.
|
|
router.route(mitogen.core.Message.pickled(' '*8192))
|
|
router.broker.defer_sync(lambda: None)
|
|
|
|
expect = 'message too large (max 4096 bytes)'
|
|
self.assertTrue(expect in logs.stop())
|
|
|
|
def test_local_dead_message(self):
|
|
# Local router should generate dead message when reply_to is set.
|
|
router = self.klass(broker=self.broker, max_message_size=4096)
|
|
|
|
logs = testlib.LogCapturer()
|
|
logs.start()
|
|
|
|
expect = router.too_large_msg % (4096,)
|
|
|
|
# Try function call. Receiver should be woken by a dead message sent by
|
|
# router due to message size exceeded.
|
|
child = router.local()
|
|
e = self.assertRaises(mitogen.core.ChannelError,
|
|
lambda: child.call(zlib.crc32, ' '*8192))
|
|
self.assertEquals(e.args[0], expect)
|
|
|
|
self.assertTrue(expect in logs.stop())
|
|
|
|
def test_remote_configured(self):
|
|
router = self.klass(broker=self.broker, max_message_size=64*1024)
|
|
remote = router.local()
|
|
size = remote.call(return_router_max_message_size)
|
|
self.assertEquals(size, 64*1024)
|
|
|
|
def test_remote_exceeded(self):
|
|
# Ensure new contexts receive a router with the same value.
|
|
router = self.klass(broker=self.broker, max_message_size=64*1024)
|
|
recv = mitogen.core.Receiver(router)
|
|
|
|
logs = testlib.LogCapturer()
|
|
logs.start()
|
|
remote = router.local()
|
|
remote.call(send_n_sized_reply, recv.to_sender(), 128*1024)
|
|
|
|
expect = 'message too large (max %d bytes)' % (64*1024,)
|
|
self.assertTrue(expect in logs.stop())
|
|
|
|
|
|
class NoRouteTest(testlib.RouterMixin, testlib.TestCase):
|
|
def test_invalid_handle_returns_dead(self):
|
|
# Verify sending a message to an invalid handle yields a dead message
|
|
# from the target context.
|
|
l1 = self.router.local()
|
|
recv = l1.send_async(mitogen.core.Message(handle=999))
|
|
msg = recv.get(throw_dead=False)
|
|
self.assertEquals(msg.is_dead, True)
|
|
self.assertEquals(msg.src_id, l1.context_id)
|
|
self.assertEquals(msg.data, self.router.invalid_handle_msg.encode())
|
|
|
|
recv = l1.send_async(mitogen.core.Message(handle=999))
|
|
e = self.assertRaises(mitogen.core.ChannelError,
|
|
lambda: recv.get())
|
|
self.assertEquals(e.args[0], self.router.invalid_handle_msg)
|
|
|
|
def test_totally_invalid_context_returns_dead(self):
|
|
recv = mitogen.core.Receiver(self.router)
|
|
msg = mitogen.core.Message(
|
|
dst_id=1234,
|
|
handle=1234,
|
|
reply_to=recv.handle,
|
|
)
|
|
self.router.route(msg)
|
|
rmsg = recv.get(throw_dead=False)
|
|
self.assertEquals(rmsg.is_dead, True)
|
|
self.assertEquals(rmsg.src_id, mitogen.context_id)
|
|
self.assertEquals(rmsg.data, (self.router.no_route_msg % (
|
|
1234,
|
|
mitogen.context_id,
|
|
)).encode())
|
|
|
|
self.router.route(msg)
|
|
e = self.assertRaises(mitogen.core.ChannelError,
|
|
lambda: recv.get())
|
|
self.assertEquals(e.args[0], (self.router.no_route_msg % (
|
|
1234,
|
|
mitogen.context_id,
|
|
)))
|
|
|
|
def test_previously_alive_context_returns_dead(self):
|
|
l1 = self.router.local()
|
|
l1.shutdown(wait=True)
|
|
recv = mitogen.core.Receiver(self.router)
|
|
msg = mitogen.core.Message(
|
|
dst_id=l1.context_id,
|
|
handle=mitogen.core.CALL_FUNCTION,
|
|
reply_to=recv.handle,
|
|
)
|
|
self.router.route(msg)
|
|
rmsg = recv.get(throw_dead=False)
|
|
self.assertEquals(rmsg.is_dead, True)
|
|
self.assertEquals(rmsg.src_id, mitogen.context_id)
|
|
self.assertEquals(rmsg.data, (self.router.no_route_msg % (
|
|
l1.context_id,
|
|
mitogen.context_id,
|
|
)).encode())
|
|
|
|
self.router.route(msg)
|
|
e = self.assertRaises(mitogen.core.ChannelError,
|
|
lambda: recv.get())
|
|
self.assertEquals(e.args[0], self.router.no_route_msg % (
|
|
l1.context_id,
|
|
mitogen.context_id,
|
|
))
|
|
|
|
|
|
class UnidirectionalTest(testlib.RouterMixin, testlib.TestCase):
|
|
def test_siblings_cant_talk(self):
|
|
self.router.unidirectional = True
|
|
l1 = self.router.local()
|
|
l2 = self.router.local()
|
|
logs = testlib.LogCapturer()
|
|
logs.start()
|
|
e = self.assertRaises(mitogen.core.CallError,
|
|
lambda: l2.call(ping_context, l1))
|
|
|
|
msg = self.router.unidirectional_msg % (
|
|
l2.context_id,
|
|
l1.context_id,
|
|
)
|
|
self.assertTrue(msg in str(e))
|
|
self.assertTrue('routing mode prevents forward of ' in logs.stop())
|
|
|
|
def test_auth_id_can_talk(self):
|
|
self.router.unidirectional = True
|
|
# One stream has auth_id stamped to that of the master, so it should be
|
|
# treated like a parent.
|
|
l1 = self.router.local()
|
|
l1s = self.router.stream_by_id(l1.context_id)
|
|
l1s.auth_id = mitogen.context_id
|
|
l1s.is_privileged = True
|
|
|
|
l2 = self.router.local()
|
|
e = self.assertRaises(mitogen.core.CallError,
|
|
lambda: l2.call(ping_context, l1))
|
|
|
|
msg = 'mitogen.core.ChannelError: %s' % (self.router.refused_msg,)
|
|
self.assertTrue(str(e).startswith(msg))
|
|
|
|
|
|
class EgressIdsTest(testlib.RouterMixin, testlib.TestCase):
|
|
def test_egress_ids_populated(self):
|
|
# Ensure Stream.egress_ids is populated on message reception.
|
|
c1 = self.router.local()
|
|
stream = self.router.stream_by_id(c1.context_id)
|
|
self.assertEquals(set(), stream.egress_ids)
|
|
|
|
c1.call(time.sleep, 0)
|
|
self.assertEquals(set([mitogen.context_id]), stream.egress_ids)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest2.main()
|