546 lines
18 KiB
Python
546 lines
18 KiB
Python
|
|
import sys
|
|
import struct
|
|
|
|
import mock
|
|
import unittest2
|
|
|
|
import mitogen.core
|
|
import mitogen.master
|
|
import testlib
|
|
|
|
from mitogen.core import b
|
|
|
|
|
|
class ConstructorTest(testlib.TestCase):
|
|
klass = mitogen.core.Message
|
|
|
|
def test_dst_id_default(self):
|
|
self.assertEquals(self.klass().dst_id, None)
|
|
|
|
def test_dst_id_explicit(self):
|
|
self.assertEquals(self.klass(dst_id=1111).dst_id, 1111)
|
|
|
|
@mock.patch('mitogen.context_id', 1234)
|
|
def test_src_id_default(self):
|
|
self.assertEquals(self.klass().src_id, 1234)
|
|
|
|
def test_src_id_explicit(self):
|
|
self.assertEquals(self.klass(src_id=4321).src_id, 4321)
|
|
|
|
@mock.patch('mitogen.context_id', 5555)
|
|
def test_auth_id_default(self):
|
|
self.assertEquals(self.klass().auth_id, 5555)
|
|
|
|
def test_auth_id_explicit(self):
|
|
self.assertEquals(self.klass(auth_id=2222).auth_id, 2222)
|
|
|
|
def test_handle_default(self):
|
|
self.assertEquals(self.klass().handle, None)
|
|
|
|
def test_handle_explicit(self):
|
|
self.assertEquals(self.klass(handle=1234).handle, 1234)
|
|
|
|
def test_reply_to_default(self):
|
|
self.assertEquals(self.klass().reply_to, None)
|
|
|
|
def test_reply_to_explicit(self):
|
|
self.assertEquals(self.klass(reply_to=8888).reply_to, 8888)
|
|
|
|
def test_data_default(self):
|
|
m = self.klass()
|
|
self.assertEquals(m.data, b(''))
|
|
self.assertTrue(isinstance(m.data, mitogen.core.BytesType))
|
|
|
|
def test_data_explicit(self):
|
|
m = self.klass(data=b('asdf'))
|
|
self.assertEquals(m.data, b('asdf'))
|
|
self.assertTrue(isinstance(m.data, mitogen.core.BytesType))
|
|
|
|
def test_data_hates_unicode(self):
|
|
self.assertRaises(Exception,
|
|
lambda: self.klass(data=u'asdf'))
|
|
|
|
|
|
class PackTest(testlib.TestCase):
|
|
klass = mitogen.core.Message
|
|
|
|
def test_header_format_sanity(self):
|
|
self.assertEquals(self.klass.HEADER_LEN,
|
|
struct.calcsize(self.klass.HEADER_FMT))
|
|
|
|
def test_header_length_correct(self):
|
|
s = self.klass(dst_id=123, handle=123).pack()
|
|
self.assertEquals(len(s), self.klass.HEADER_LEN)
|
|
|
|
def test_magic(self):
|
|
s = self.klass(dst_id=123, handle=123).pack()
|
|
magic, = struct.unpack('>h', s[:2])
|
|
self.assertEquals(self.klass.HEADER_MAGIC, magic)
|
|
|
|
def test_dst_id(self):
|
|
s = self.klass(dst_id=123, handle=123).pack()
|
|
dst_id, = struct.unpack('>L', s[2:6])
|
|
self.assertEquals(123, dst_id)
|
|
|
|
def test_src_id(self):
|
|
s = self.klass(src_id=5432, dst_id=123, handle=123).pack()
|
|
src_id, = struct.unpack('>L', s[6:10])
|
|
self.assertEquals(5432, src_id)
|
|
|
|
def test_auth_id(self):
|
|
s = self.klass(auth_id=1919, src_id=5432, dst_id=123, handle=123).pack()
|
|
auth_id, = struct.unpack('>L', s[10:14])
|
|
self.assertEquals(1919, auth_id)
|
|
|
|
def test_handle(self):
|
|
s = self.klass(dst_id=123, handle=9999).pack()
|
|
handle, = struct.unpack('>L', s[14:18])
|
|
self.assertEquals(9999, handle)
|
|
|
|
def test_reply_to(self):
|
|
s = self.klass(dst_id=1231, handle=7777, reply_to=9132).pack()
|
|
reply_to, = struct.unpack('>L', s[18:22])
|
|
self.assertEquals(9132, reply_to)
|
|
|
|
def test_data_length_empty(self):
|
|
s = self.klass(dst_id=1231, handle=7777).pack()
|
|
data_length, = struct.unpack('>L', s[22:26])
|
|
self.assertEquals(0, data_length)
|
|
|
|
def test_data_length_present(self):
|
|
s = self.klass(dst_id=1231, handle=7777, data=b('hello')).pack()
|
|
data_length, = struct.unpack('>L', s[22:26])
|
|
self.assertEquals(5, data_length)
|
|
|
|
def test_data_empty(self):
|
|
s = self.klass(dst_id=1231, handle=7777).pack()
|
|
data = s[26:]
|
|
self.assertEquals(b(''), data)
|
|
|
|
def test_data_present(self):
|
|
s = self.klass(dst_id=11, handle=77, data=b('hello')).pack()
|
|
data = s[26:]
|
|
self.assertEquals(b('hello'), data)
|
|
|
|
|
|
class IsDeadTest(testlib.TestCase):
|
|
klass = mitogen.core.Message
|
|
|
|
def test_is_dead(self):
|
|
msg = self.klass(reply_to=mitogen.core.IS_DEAD)
|
|
self.assertTrue(msg.is_dead)
|
|
|
|
def test_is_not_dead(self):
|
|
msg = self.klass(reply_to=5555)
|
|
self.assertFalse(msg.is_dead)
|
|
|
|
|
|
class DeadTest(testlib.TestCase):
|
|
klass = mitogen.core.Message
|
|
|
|
def test_no_reason(self):
|
|
msg = self.klass.dead()
|
|
self.assertEquals(msg.reply_to, mitogen.core.IS_DEAD)
|
|
self.assertTrue(msg.is_dead)
|
|
self.assertEquals(msg.data, b(''))
|
|
|
|
def test_with_reason(self):
|
|
msg = self.klass.dead(reason=u'oh no')
|
|
self.assertEquals(msg.reply_to, mitogen.core.IS_DEAD)
|
|
self.assertTrue(msg.is_dead)
|
|
self.assertEquals(msg.data, b('oh no'))
|
|
|
|
|
|
class EvilObject(object):
|
|
pass
|
|
|
|
|
|
class PickledTest(testlib.TestCase):
|
|
# getting_started.html#rpc-serialization-rules
|
|
klass = mitogen.core.Message
|
|
|
|
def roundtrip(self, v, router=None):
|
|
msg = self.klass.pickled(v)
|
|
msg2 = self.klass(data=msg.data)
|
|
msg2.router = router
|
|
return msg2.unpickle()
|
|
|
|
def test_bool(self):
|
|
for b in True, False:
|
|
self.assertEquals(b, self.roundtrip(b))
|
|
|
|
@unittest2.skipIf(condition=sys.version_info < (2, 6),
|
|
reason='bytearray missing on <2.6')
|
|
def test_bytearray(self):
|
|
ba = bytearray(b('123'))
|
|
self.assertRaises(mitogen.core.StreamError,
|
|
lambda: self.roundtrip(ba)
|
|
)
|
|
|
|
def test_bytes(self):
|
|
by = b('123')
|
|
self.assertEquals(by, self.roundtrip(by))
|
|
|
|
def test_dict(self):
|
|
d = {1: 2, u'a': 3, b('b'): 4, 'c': {}}
|
|
roundtrip = self.roundtrip(d)
|
|
self.assertEquals(d, roundtrip)
|
|
self.assertTrue(isinstance(roundtrip, dict))
|
|
for k in d:
|
|
self.assertTrue(isinstance(roundtrip[k], type(d[k])))
|
|
|
|
def test_int(self):
|
|
self.assertEquals(123, self.klass.pickled(123).unpickle())
|
|
|
|
def test_list(self):
|
|
l = [1, u'b', b('c')]
|
|
roundtrip = self.roundtrip(l)
|
|
self.assertTrue(isinstance(roundtrip, list))
|
|
self.assertEquals(l, roundtrip)
|
|
for k in range(len(l)):
|
|
self.assertTrue(isinstance(roundtrip[k], type(l[k])))
|
|
|
|
@unittest2.skipIf(condition=sys.version_info > (3, 0),
|
|
reason='long missing in >3.x')
|
|
def test_long(self):
|
|
l = long(0xffffffffffff)
|
|
roundtrip = self.roundtrip(l)
|
|
self.assertEquals(l, roundtrip)
|
|
self.assertTrue(isinstance(roundtrip, long))
|
|
|
|
def test_tuple(self):
|
|
l = (1, u'b', b('c'))
|
|
roundtrip = self.roundtrip(l)
|
|
self.assertEquals(l, roundtrip)
|
|
self.assertTrue(isinstance(roundtrip, tuple))
|
|
for k in range(len(l)):
|
|
self.assertTrue(isinstance(roundtrip[k], type(l[k])))
|
|
|
|
def test_unicode(self):
|
|
u = u'abcd'
|
|
roundtrip = self.roundtrip(u)
|
|
self.assertEquals(u, roundtrip)
|
|
self.assertTrue(isinstance(roundtrip, mitogen.core.UnicodeType))
|
|
|
|
#### custom types. see also: types_test.py, call_error_test.py
|
|
|
|
# Python 3 pickle protocol 2 does weird stuff depending on whether an empty
|
|
# or nonempty bytes is being serialized. For non-empty, it yields a
|
|
# _codecs.encode() call. For empty, it yields a bytes() call.
|
|
|
|
def test_blob_nonempty(self):
|
|
v = mitogen.core.Blob(b('dave'))
|
|
roundtrip = self.roundtrip(v)
|
|
self.assertTrue(isinstance(roundtrip, mitogen.core.Blob))
|
|
self.assertEquals(b('dave'), roundtrip)
|
|
|
|
def test_blob_empty(self):
|
|
v = mitogen.core.Blob(b(''))
|
|
roundtrip = self.roundtrip(v)
|
|
self.assertTrue(isinstance(roundtrip, mitogen.core.Blob))
|
|
self.assertEquals(b(''), v)
|
|
|
|
def test_secret_nonempty(self):
|
|
s = mitogen.core.Secret(u'dave')
|
|
roundtrip = self.roundtrip(s)
|
|
self.assertTrue(isinstance(roundtrip, mitogen.core.Secret))
|
|
self.assertEquals(u'dave', roundtrip)
|
|
|
|
def test_secret_empty(self):
|
|
s = mitogen.core.Secret(u'')
|
|
roundtrip = self.roundtrip(s)
|
|
self.assertTrue(isinstance(roundtrip, mitogen.core.Secret))
|
|
self.assertEquals(u'', roundtrip)
|
|
|
|
def test_call_error(self):
|
|
ce = mitogen.core.CallError('nope')
|
|
ce2 = self.assertRaises(mitogen.core.CallError,
|
|
lambda: self.roundtrip(ce))
|
|
self.assertEquals(ce.args[0], ce2.args[0])
|
|
|
|
def test_context(self):
|
|
router = mitogen.master.Router()
|
|
try:
|
|
c = router.context_by_id(1234)
|
|
roundtrip = self.roundtrip(c)
|
|
self.assertTrue(isinstance(roundtrip, mitogen.core.Context))
|
|
self.assertEquals(c.context_id, 1234)
|
|
finally:
|
|
router.broker.shutdown()
|
|
router.broker.join()
|
|
|
|
def test_sender(self):
|
|
router = mitogen.master.Router()
|
|
try:
|
|
recv = mitogen.core.Receiver(router)
|
|
sender = recv.to_sender()
|
|
roundtrip = self.roundtrip(sender, router=router)
|
|
self.assertTrue(isinstance(roundtrip, mitogen.core.Sender))
|
|
self.assertEquals(roundtrip.context.context_id, mitogen.context_id)
|
|
self.assertEquals(roundtrip.dst_handle, sender.dst_handle)
|
|
finally:
|
|
router.broker.shutdown()
|
|
router.broker.join()
|
|
|
|
####
|
|
|
|
def test_custom_object_deserialization_fails(self):
|
|
self.assertRaises(mitogen.core.StreamError,
|
|
lambda: self.roundtrip(EvilObject())
|
|
)
|
|
|
|
|
|
class ReplyTest(testlib.TestCase):
|
|
# getting_started.html#rpc-serialization-rules
|
|
klass = mitogen.core.Message
|
|
|
|
def test_reply_calls_router_route(self):
|
|
msg = self.klass(src_id=1234, reply_to=9191)
|
|
router = mock.Mock()
|
|
msg.reply(123, router=router)
|
|
self.assertEquals(1, router.route.call_count)
|
|
|
|
def test_reply_pickles_object(self):
|
|
msg = self.klass(src_id=1234, reply_to=9191)
|
|
router = mock.Mock()
|
|
msg.reply(123, router=router)
|
|
_, (reply,), _ = router.route.mock_calls[0]
|
|
self.assertEquals(reply.dst_id, 1234)
|
|
self.assertEquals(reply.unpickle(), 123)
|
|
|
|
def test_reply_uses_preformatted_message(self):
|
|
msg = self.klass(src_id=1234, reply_to=9191)
|
|
router = mock.Mock()
|
|
my_reply = mitogen.core.Message.pickled(4444)
|
|
msg.reply(my_reply, router=router)
|
|
_, (reply,), _ = router.route.mock_calls[0]
|
|
self.assertTrue(my_reply is reply)
|
|
self.assertEquals(reply.dst_id, 1234)
|
|
self.assertEquals(reply.unpickle(), 4444)
|
|
|
|
def test_reply_sets_dst_id(self):
|
|
msg = self.klass(src_id=1234, reply_to=9191)
|
|
router = mock.Mock()
|
|
msg.reply(123, router=router)
|
|
_, (reply,), _ = router.route.mock_calls[0]
|
|
self.assertEquals(reply.dst_id, 1234)
|
|
|
|
def test_reply_sets_handle(self):
|
|
msg = self.klass(src_id=1234, reply_to=9191)
|
|
router = mock.Mock()
|
|
msg.reply(123, router=router)
|
|
_, (reply,), _ = router.route.mock_calls[0]
|
|
self.assertEquals(reply.handle, 9191)
|
|
|
|
|
|
class UnpickleTest(testlib.TestCase):
|
|
# mostly done by PickleTest, just check behaviour of parameters
|
|
klass = mitogen.core.Message
|
|
|
|
def test_throw(self):
|
|
ce = mitogen.core.CallError('nope')
|
|
m = self.klass.pickled(ce)
|
|
ce2 = self.assertRaises(mitogen.core.CallError,
|
|
lambda: m.unpickle())
|
|
self.assertEquals(ce.args[0], ce2.args[0])
|
|
|
|
def test_no_throw(self):
|
|
ce = mitogen.core.CallError('nope')
|
|
m = self.klass.pickled(ce)
|
|
ce2 = m.unpickle(throw=False)
|
|
self.assertEquals(ce.args[0], ce2.args[0])
|
|
|
|
def test_throw_dead(self):
|
|
m = self.klass.pickled('derp', reply_to=mitogen.core.IS_DEAD)
|
|
self.assertRaises(mitogen.core.ChannelError,
|
|
lambda: m.unpickle())
|
|
|
|
def test_no_throw_dead(self):
|
|
m = self.klass.pickled('derp', reply_to=mitogen.core.IS_DEAD)
|
|
self.assertEquals('derp', m.unpickle(throw_dead=False))
|
|
|
|
|
|
class UnpickleCompatTest(testlib.TestCase):
|
|
# try weird variations of pickles from different Python versions.
|
|
klass = mitogen.core.Message
|
|
|
|
def check(self, value, encoded, **kwargs):
|
|
if isinstance(encoded, mitogen.core.UnicodeType):
|
|
encoded = encoded.encode('latin1')
|
|
m = self.klass(data=encoded)
|
|
m.router = mitogen.master.Router()
|
|
try:
|
|
return m.unpickle(**kwargs)
|
|
finally:
|
|
m.router.broker.shutdown()
|
|
m.router.broker.join()
|
|
|
|
def test_py24_bytes(self):
|
|
self.check('test',
|
|
('\x80\x02U\x04testq\x00.'))
|
|
|
|
def test_py24_unicode(self):
|
|
self.check(u'test',
|
|
('\x80\x02X\x04\x00\x00\x00testq\x00.'))
|
|
|
|
def test_py24_int(self):
|
|
self.check(123,
|
|
('\x80\x02K{.'))
|
|
|
|
def test_py24_long(self):
|
|
self.check(17592186044415,
|
|
('\x80\x02\x8a\x06\xff\xff\xff\xff\xff\x0f.'))
|
|
|
|
def test_py24_dict(self):
|
|
self.check({},
|
|
('\x80\x02}q\x00.'))
|
|
|
|
def test_py24_tuple(self):
|
|
self.check((1, 2, u'b'),
|
|
('\x80\x02K\x01K\x02X\x01\x00\x00\x00bq\x00\x87q\x01.'))
|
|
|
|
def test_py24_bool(self):
|
|
self.check(True,
|
|
('\x80\x02\x88.'))
|
|
|
|
def test_py24_list(self):
|
|
self.check([1, 2, u'b'],
|
|
('\x80\x02]q\x00(K\x01K\x02X\x01\x00\x00\x00bq\x01e.'))
|
|
|
|
def test_py24_blob(self):
|
|
self.check(mitogen.core.mitogen.core.Blob(b('bigblob')),
|
|
('\x80\x02cmitogen.core\nBlob\nq\x00U\x07bigblobq\x01\x85q\x02Rq\x03.'))
|
|
|
|
def test_py24_secret(self):
|
|
self.check(mitogen.core.Secret(u'mypassword'),
|
|
('\x80\x02cmitogen.core\nSecret\nq\x00X\n\x00\x00\x00mypasswordq\x01\x85q\x02Rq\x03.'))
|
|
|
|
def test_py24_call_error(self):
|
|
self.check(mitogen.core.CallError('big error'),
|
|
('\x80\x02cmitogen.core\n_unpickle_call_error\nq\x00X\t\x00\x00\x00big errorq\x01\x85q\x02R.'), throw=False)
|
|
|
|
def test_py24_context(self):
|
|
self.check(mitogen.core.Context(1234, None),
|
|
('\x80\x02cmitogen.core\n_unpickle_context\nq\x00M\xd2\x04N\x86q\x01Rq\x02.'))
|
|
|
|
def test_py24_sender(self):
|
|
self.check(mitogen.core.Sender(mitogen.core.Context(55555, None), 4444),
|
|
('\x80\x02cmitogen.core\n_unpickle_sender\nq\x00M\x03\xd9M\\\x11\x86q\x01Rq\x02.'))
|
|
|
|
def test_py27_bytes(self):
|
|
self.check(b('test'),
|
|
('\x80\x02U\x04testq\x01.'))
|
|
|
|
def test_py27_unicode(self):
|
|
self.check(u'test',
|
|
('\x80\x02X\x04\x00\x00\x00testq\x01.'))
|
|
|
|
def test_py27_int(self):
|
|
self.check(123,
|
|
('\x80\x02K{.'))
|
|
|
|
def test_py27_long(self):
|
|
self.check(17592186044415,
|
|
('\x80\x02\x8a\x06\xff\xff\xff\xff\xff\x0f.'))
|
|
|
|
def test_py27_dict(self):
|
|
self.check({},
|
|
('\x80\x02}q\x01.'))
|
|
|
|
def test_py27_tuple(self):
|
|
self.check((1, 2, u'b'),
|
|
('\x80\x02K\x01K\x02X\x01\x00\x00\x00b\x87q\x01.'))
|
|
|
|
def test_py27_bool(self):
|
|
self.check(True,
|
|
('\x80\x02\x88.'))
|
|
|
|
def test_py27_list(self):
|
|
self.check([1, 2, u'b'],
|
|
('\x80\x02]q\x01(K\x01K\x02X\x01\x00\x00\x00be.'))
|
|
|
|
def test_py27_blob(self):
|
|
self.check(mitogen.core.mitogen.core.Blob(b('bigblob')),
|
|
('\x80\x02cmitogen.core\nBlob\nq\x01U\x07bigblob\x85Rq\x02.'))
|
|
|
|
def test_py27_secret(self):
|
|
self.check(mitogen.core.Secret(u'mypassword'),
|
|
('\x80\x02cmitogen.core\nSecret\nq\x01X\n\x00\x00\x00mypassword\x85Rq\x02.'))
|
|
|
|
def test_py27_call_error(self):
|
|
self.check(mitogen.core.CallError(u'big error',),
|
|
('\x80\x02cmitogen.core\n_unpickle_call_error\nq\x01X\t\x00\x00\x00big errorq\x02\x85Rq\x03.'), throw=False)
|
|
|
|
def test_py27_context(self):
|
|
self.check(mitogen.core.Context(1234, None),
|
|
('\x80\x02cmitogen.core\n_unpickle_context\nq\x01M\xd2\x04N\x86Rq\x02.'))
|
|
|
|
def test_py27_sender(self):
|
|
self.check(mitogen.core.Sender(mitogen.core.Context(55555, None), 4444),
|
|
('\x80\x02cmitogen.core\n_unpickle_sender\nq\x01M\x03\xd9M\\\x11\x86Rq\x02.'))
|
|
|
|
def test_py36_bytes(self):
|
|
self.check(b('test'),
|
|
('\x80\x02c_codecs\nencode\nq\x00X\x04\x00\x00\x00testq\x01X\x06\x00\x00\x00latin1q\x02\x86q\x03Rq\x04.'))
|
|
|
|
def test_py36_unicode(self):
|
|
self.check('test',
|
|
('\x80\x02X\x04\x00\x00\x00testq\x00.'))
|
|
|
|
def test_py36_int(self):
|
|
self.check(123,
|
|
('\x80\x02K{.'))
|
|
|
|
def test_py36_long(self):
|
|
self.check(17592186044415,
|
|
('\x80\x02\x8a\x06\xff\xff\xff\xff\xff\x0f.'))
|
|
|
|
def test_py36_dict(self):
|
|
self.check({},
|
|
('\x80\x02}q\x00.'))
|
|
|
|
def test_py36_tuple(self):
|
|
self.check((1, 2, u'b'),
|
|
('\x80\x02K\x01K\x02X\x01\x00\x00\x00bq\x00\x87q\x01.'))
|
|
|
|
def test_py36_bool(self):
|
|
self.check(True,
|
|
('\x80\x02\x88.'))
|
|
|
|
def test_py36_list(self):
|
|
self.check([1, 2, u'b'],
|
|
('\x80\x02]q\x00(K\x01K\x02X\x01\x00\x00\x00bq\x01e.'))
|
|
|
|
def test_py36_blob(self):
|
|
self.check(mitogen.core.mitogen.core.Blob(b('bigblob')),
|
|
('\x80\x02cmitogen.core\nBlob\nq\x00c_codecs\nencode\nq\x01X\x07\x00\x00\x00bigblobq\x02X\x06\x00\x00\x00latin1q\x03\x86q\x04Rq\x05\x85q\x06Rq\x07.'))
|
|
|
|
def test_py36_secret(self):
|
|
self.check(mitogen.core.Secret('mypassword'),
|
|
('\x80\x02cmitogen.core\nSecret\nq\x00X\n\x00\x00\x00mypasswordq\x01\x85q\x02Rq\x03.'))
|
|
|
|
def test_py36_call_error(self):
|
|
self.check(mitogen.core.CallError('big error'),
|
|
('\x80\x02cmitogen.core\n_unpickle_call_error\nq\x00X\t\x00\x00\x00big errorq\x01\x85q\x02Rq\x03.'), throw=False)
|
|
|
|
def test_py36_context(self):
|
|
self.check(mitogen.core.Context(1234, None),
|
|
('\x80\x02cmitogen.core\n_unpickle_context\nq\x00M\xd2\x04N\x86q\x01Rq\x02.'))
|
|
|
|
def test_py36_sender(self):
|
|
self.check(mitogen.core.Sender(mitogen.core.Context(55555, None), 4444),
|
|
('\x80\x02cmitogen.core\n_unpickle_sender\nq\x00M\x03\xd9M\\\x11\x86q\x01Rq\x02.'))
|
|
|
|
|
|
class ReprTest(testlib.TestCase):
|
|
klass = mitogen.core.Message
|
|
|
|
def test_repr(self):
|
|
# doesn't crash
|
|
repr(self.klass.pickled('test'))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest2.main()
|