diff --git a/mitogen/core.py b/mitogen/core.py index f2dece48..faa2d516 100644 --- a/mitogen/core.py +++ b/mitogen/core.py @@ -553,7 +553,7 @@ class Message(object): assert isinstance(self.data, BytesType) def _unpickle_context(self, context_id, name): - return _unpickle_context(self.router, context_id, name) + return _unpickle_context(context_id, name, router=self.router) def _unpickle_sender(self, context_id, dst_handle): return _unpickle_sender(self.router, context_id, dst_handle) @@ -1498,14 +1498,16 @@ class Context(object): return 'Context(%s, %r)' % (self.context_id, self.name) -def _unpickle_context(router, context_id, name): - if not (isinstance(router, Router) and - isinstance(context_id, (int, long)) and context_id >= 0 and ( - (name is None) or - (isinstance(name, UnicodeType) and len(name) < 100)) - ): +def _unpickle_context(context_id, name, router=None): + if not (isinstance(context_id, (int, long)) and context_id >= 0 and ( + (name is None) or + (isinstance(name, UnicodeType) and len(name) < 100)) + ): raise TypeError('cannot unpickle Context: bad input') - return router.context_by_id(context_id, name=name) + + if isinstance(router, Router): + return router.context_by_id(context_id, name=name) + return Context(None, context_id, name) # For plain Jane pickle. class Poller(object): diff --git a/tests/serialization_test.py b/tests/serialization_test.py index 3b56e10a..f108ff37 100644 --- a/tests/serialization_test.py +++ b/tests/serialization_test.py @@ -6,11 +6,14 @@ except ImportError: from StringIO import StringIO as StringIO from StringIO import StringIO as BytesIO +import pickle import unittest2 import mitogen.core from mitogen.core import b +import testlib + def roundtrip(v): msg = mitogen.core.Message.pickled(v) @@ -33,5 +36,30 @@ class BlobTest(unittest2.TestCase): self.assertEquals(b(''), roundtrip(v)) +class ContextTest(testlib.RouterMixin, unittest2.TestCase): + klass = mitogen.core.Context + + # Ensure Context can be round-tripped by regular pickle in addition to + # Mitogen's hacked pickle. Users may try to call pickle on a Context in + # strange circumstances, and it's often used to glue pieces of an app + # together (e.g. Ansible). + + def test_mitogen_roundtrip(self): + c = self.router.fork() + r = mitogen.core.Receiver(self.router) + r.to_sender().send(c) + c2 = r.get().unpickle() + self.assertEquals(None, c2.router) + self.assertEquals(c.context_id, c2.context_id) + self.assertEquals(c.name, c2.name) + + def test_vanilla_roundtrip(self): + c = self.router.fork() + c2 = pickle.loads(pickle.dumps(c)) + self.assertEquals(None, c2.router) + self.assertEquals(c.context_id, c2.context_id) + self.assertEquals(c.name, c2.name) + + if __name__ == '__main__': unittest2.main()