try:
    from io import StringIO
    from io import BytesIO
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)
    return mitogen.core.Message(data=msg.data).unpickle()


class BlobTest(testlib.TestCase):
    klass = mitogen.core.Blob

    # 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_nonempty_bytes(self):
        v = mitogen.core.Blob(b('dave'))
        self.assertEquals(b('dave'), roundtrip(v))

    def test_empty_bytes(self):
        v = mitogen.core.Blob(b(''))
        self.assertEquals(b(''), roundtrip(v))


class ContextTest(testlib.RouterMixin, testlib.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()