IPython/kernel/zmq/tests/test_session.py

"""test building messages with streamsession"""

#-------------------------------------------------------------------------------
#  Copyright (C) 2011  The IPython Development Team
#
#  Distributed under the terms of the BSD License.  The full license is in
#  the file COPYING, distributed as part of this software.
#-------------------------------------------------------------------------------

#-------------------------------------------------------------------------------
# Imports
#-------------------------------------------------------------------------------

import os
import uuid
from datetime import datetime

import zmq

from zmq.tests import BaseZMQTestCase
from zmq.eventloop.zmqstream import ZMQStream

from IPython.kernel.zmq import session as ss

from IPython.testing.decorators import skipif, module_not_available
from IPython.utils.py3compat import string_types
from IPython.utils import jsonutil

def _bad_packer(obj):
    raise TypeError("I don't work")

def _bad_unpacker(bytes):
    raise TypeError("I don't work either")

class SessionTestCase(BaseZMQTestCase):

    def setUp(self):
        BaseZMQTestCase.setUp(self)
        self.session = ss.Session()


class TestSession(SessionTestCase):

    def test_msg(self):
        """message format"""
        msg = self.session.msg('execute')
        thekeys = set('header parent_header metadata content msg_type msg_id'.split())
        s = set(msg.keys())
        self.assertEqual(s, thekeys)
        self.assertTrue(isinstance(msg['content'],dict))
        self.assertTrue(isinstance(msg['metadata'],dict))
        self.assertTrue(isinstance(msg['header'],dict))
        self.assertTrue(isinstance(msg['parent_header'],dict))
        self.assertTrue(isinstance(msg['msg_id'],str))
        self.assertTrue(isinstance(msg['msg_type'],str))
        self.assertEqual(msg['header']['msg_type'], 'execute')
        self.assertEqual(msg['msg_type'], 'execute')

    def test_serialize(self):
        msg = self.session.msg('execute', content=dict(a=10, b=1.1))
        msg_list = self.session.serialize(msg, ident=b'foo')
        ident, msg_list = self.session.feed_identities(msg_list)
        new_msg = self.session.unserialize(msg_list)
        self.assertEqual(ident[0], b'foo')
        self.assertEqual(new_msg['msg_id'],msg['msg_id'])
        self.assertEqual(new_msg['msg_type'],msg['msg_type'])
        self.assertEqual(new_msg['header'],msg['header'])
        self.assertEqual(new_msg['content'],msg['content'])
        self.assertEqual(new_msg['parent_header'],msg['parent_header'])
        self.assertEqual(new_msg['metadata'],msg['metadata'])
        # ensure floats don't come out as Decimal:
        self.assertEqual(type(new_msg['content']['b']),type(new_msg['content']['b']))

    def test_send(self):
        ctx = zmq.Context.instance()
        A = ctx.socket(zmq.PAIR)
        B = ctx.socket(zmq.PAIR)
        A.bind("inproc://test")
        B.connect("inproc://test")

        msg = self.session.msg('execute', content=dict(a=10))
        self.session.send(A, msg, ident=b'foo', buffers=[b'bar'])
        
        ident, msg_list = self.session.feed_identities(B.recv_multipart())
        new_msg = self.session.unserialize(msg_list)
        self.assertEqual(ident[0], b'foo')
        self.assertEqual(new_msg['msg_id'],msg['msg_id'])
        self.assertEqual(new_msg['msg_type'],msg['msg_type'])
        self.assertEqual(new_msg['header'],msg['header'])
        self.assertEqual(new_msg['content'],msg['content'])
        self.assertEqual(new_msg['parent_header'],msg['parent_header'])
        self.assertEqual(new_msg['metadata'],msg['metadata'])
        self.assertEqual(new_msg['buffers'],[b'bar'])

        content = msg['content']
        header = msg['header']
        parent = msg['parent_header']
        metadata = msg['metadata']
        msg_type = header['msg_type']
        self.session.send(A, None, content=content, parent=parent,
            header=header, metadata=metadata, ident=b'foo', buffers=[b'bar'])
        ident, msg_list = self.session.feed_identities(B.recv_multipart())
        new_msg = self.session.unserialize(msg_list)
        self.assertEqual(ident[0], b'foo')
        self.assertEqual(new_msg['msg_id'],msg['msg_id'])
        self.assertEqual(new_msg['msg_type'],msg['msg_type'])
        self.assertEqual(new_msg['header'],msg['header'])
        self.assertEqual(new_msg['content'],msg['content'])
        self.assertEqual(new_msg['metadata'],msg['metadata'])
        self.assertEqual(new_msg['parent_header'],msg['parent_header'])
        self.assertEqual(new_msg['buffers'],[b'bar'])

        self.session.send(A, msg, ident=b'foo', buffers=[b'bar'])
        ident, new_msg = self.session.recv(B)
        self.assertEqual(ident[0], b'foo')
        self.assertEqual(new_msg['msg_id'],msg['msg_id'])
        self.assertEqual(new_msg['msg_type'],msg['msg_type'])
        self.assertEqual(new_msg['header'],msg['header'])
        self.assertEqual(new_msg['content'],msg['content'])
        self.assertEqual(new_msg['metadata'],msg['metadata'])
        self.assertEqual(new_msg['parent_header'],msg['parent_header'])
        self.assertEqual(new_msg['buffers'],[b'bar'])

        A.close()
        B.close()
        ctx.term()

    def test_args(self):
        """initialization arguments for Session"""
        s = self.session
        self.assertTrue(s.pack is ss.default_packer)
        self.assertTrue(s.unpack is ss.default_unpacker)
        self.assertEqual(s.username, os.environ.get('USER', u'username'))

        s = ss.Session()
        self.assertEqual(s.username, os.environ.get('USER', u'username'))

        self.assertRaises(TypeError, ss.Session, pack='hi')
        self.assertRaises(TypeError, ss.Session, unpack='hi')
        u = str(uuid.uuid4())
        s = ss.Session(username=u'carrot', session=u)
        self.assertEqual(s.session, u)
        self.assertEqual(s.username, u'carrot')

    def test_tracking(self):
        """test tracking messages"""
        a,b = self.create_bound_pair(zmq.PAIR, zmq.PAIR)
        s = self.session
        s.copy_threshold = 1
        stream = ZMQStream(a)
        msg = s.send(a, 'hello', track=False)
        self.assertTrue(msg['tracker'] is ss.DONE)
        msg = s.send(a, 'hello', track=True)
        self.assertTrue(isinstance(msg['tracker'], zmq.MessageTracker))
        M = zmq.Message(b'hi there', track=True)
        msg = s.send(a, 'hello', buffers=[M], track=True)
        t = msg['tracker']
        self.assertTrue(isinstance(t, zmq.MessageTracker))
        self.assertRaises(zmq.NotDone, t.wait, .1)
        del M
        t.wait(1) # this will raise


    def test_unique_msg_ids(self):
        """test that messages receive unique ids"""
        ids = set()
        for i in range(2**12):
            h = self.session.msg_header('test')
            msg_id = h['msg_id']
            self.assertTrue(msg_id not in ids)
            ids.add(msg_id)

    def test_feed_identities(self):
        """scrub the front for zmq IDENTITIES"""
        theids = "engine client other".split()
        content = dict(code='whoda',stuff=object())
        themsg = self.session.msg('execute',content=content)
        pmsg = theids

    def test_session_id(self):
        session = ss.Session()
        # get bs before us
        bs = session.bsession
        us = session.session
        self.assertEqual(us.encode('ascii'), bs)
        session = ss.Session()
        # get us before bs
        us = session.session
        bs = session.bsession
        self.assertEqual(us.encode('ascii'), bs)
        # change propagates:
        session.session = 'something else'
        bs = session.bsession
        us = session.session
        self.assertEqual(us.encode('ascii'), bs)
        session = ss.Session(session='stuff')
        # get us before bs
        self.assertEqual(session.bsession, session.session.encode('ascii'))
        self.assertEqual(b'stuff', session.bsession)

    def test_zero_digest_history(self):
        session = ss.Session(digest_history_size=0)
        for i in range(11):
            session._add_digest(uuid.uuid4().bytes)
        self.assertEqual(len(session.digest_history), 0)

    def test_cull_digest_history(self):
        session = ss.Session(digest_history_size=100)
        for i in range(100):
            session._add_digest(uuid.uuid4().bytes)
        self.assertTrue(len(session.digest_history) == 100)
        session._add_digest(uuid.uuid4().bytes)
        self.assertTrue(len(session.digest_history) == 91)
        for i in range(9):
            session._add_digest(uuid.uuid4().bytes)
        self.assertTrue(len(session.digest_history) == 100)
        session._add_digest(uuid.uuid4().bytes)
        self.assertTrue(len(session.digest_history) == 91)
    
    def test_bad_pack(self):
        try:
            session = ss.Session(pack=_bad_packer)
        except ValueError as e:
            self.assertIn("could not serialize", str(e))
            self.assertIn("don't work", str(e))
        else:
            self.fail("Should have raised ValueError")
    
    def test_bad_unpack(self):
        try:
            session = ss.Session(unpack=_bad_unpacker)
        except ValueError as e:
            self.assertIn("could not handle output", str(e))
            self.assertIn("don't work either", str(e))
        else:
            self.fail("Should have raised ValueError")
    
    def test_bad_packer(self):
        try:
            session = ss.Session(packer=__name__ + '._bad_packer')
        except ValueError as e:
            self.assertIn("could not serialize", str(e))
            self.assertIn("don't work", str(e))
        else:
            self.fail("Should have raised ValueError")
    
    def test_bad_unpacker(self):
        try:
            session = ss.Session(unpacker=__name__ + '._bad_unpacker')
        except ValueError as e:
            self.assertIn("could not handle output", str(e))
            self.assertIn("don't work either", str(e))
        else:
            self.fail("Should have raised ValueError")
    
    def test_bad_roundtrip(self):
        with self.assertRaises(ValueError):
            session = ss.Session(unpack=lambda b: 5)
    
    def _datetime_test(self, session):
        content = dict(t=datetime.now())
        metadata = dict(t=datetime.now())
        p = session.msg('msg')
        msg = session.msg('msg', content=content, metadata=metadata, parent=p['header'])
        smsg = session.serialize(msg)
        msg2 = session.unserialize(session.feed_identities(smsg)[1])
        assert isinstance(msg2['header']['date'], datetime)
        self.assertEqual(msg['header'], msg2['header'])
        self.assertEqual(msg['parent_header'], msg2['parent_header'])
        self.assertEqual(msg['parent_header'], msg2['parent_header'])
        assert isinstance(msg['content']['t'], datetime)
        assert isinstance(msg['metadata']['t'], datetime)
        assert isinstance(msg2['content']['t'], string_types)
        assert isinstance(msg2['metadata']['t'], string_types)
        self.assertEqual(msg['content'], jsonutil.extract_dates(msg2['content']))
        self.assertEqual(msg['content'], jsonutil.extract_dates(msg2['content']))
    
    def test_datetimes(self):
        self._datetime_test(self.session)
    
    def test_datetimes_pickle(self):
        session = ss.Session(packer='pickle')
        self._datetime_test(session)
    
    @skipif(module_not_available('msgpack'))
    def test_datetimes_msgpack(self):
        import msgpack
        
        session = ss.Session(
            pack=msgpack.packb,
            unpack=lambda buf: msgpack.unpackb(buf, encoding='utf8'),
        )
        self._datetime_test(session)
    
    def test_send_raw(self):
        ctx = zmq.Context.instance()
        A = ctx.socket(zmq.PAIR)
        B = ctx.socket(zmq.PAIR)
        A.bind("inproc://test")
        B.connect("inproc://test")

        msg = self.session.msg('execute', content=dict(a=10))
        msg_list = [self.session.pack(msg[part]) for part in 
                    ['header', 'parent_header', 'metadata', 'content']]
        self.session.send_raw(A, msg_list, ident=b'foo')
        
        ident, new_msg_list = self.session.feed_identities(B.recv_multipart())
        new_msg = self.session.unserialize(new_msg_list)
        self.assertEqual(ident[0], b'foo')
        self.assertEqual(new_msg['msg_type'],msg['msg_type'])
        self.assertEqual(new_msg['header'],msg['header'])
        self.assertEqual(new_msg['parent_header'],msg['parent_header'])
        self.assertEqual(new_msg['content'],msg['content'])
        self.assertEqual(new_msg['metadata'],msg['metadata'])

        A.close()
        B.close()
        ctx.term()