python-kombu 3.0.33-1ubuntu2 / usr / lib / python2.7 / dist-packages / kombu / tests / test_serialization.py

#!/usr/bin/python
# -*- coding: utf-8 -*-
from __future__ import absolute_import
from __future__ import unicode_literals

import sys

from base64 import b64decode

from kombu.exceptions import ContentDisallowed, EncodeError, DecodeError
from kombu.five import text_t, bytes_t
from kombu.serialization import (
    registry, register, SerializerNotInstalled,
    raw_encode, register_yaml, register_msgpack,
    dumps, loads, pickle, pickle_protocol,
    unregister, register_pickle, enable_insecure_serializers,
    disable_insecure_serializers,
)
from kombu.utils.encoding import str_to_bytes

from .case import Case, call, mask_modules, patch, skip_if_not_module

# For content_encoding tests
unicode_string = 'abcdé\u8463'
unicode_string_as_utf8 = unicode_string.encode('utf-8')
latin_string = 'abcdé'
latin_string_as_latin1 = latin_string.encode('latin-1')
latin_string_as_utf8 = latin_string.encode('utf-8')


# For serialization tests
py_data = {
    'string': 'The quick brown fox jumps over the lazy dog',
    'int': 10,
    'float': 3.14159265,
    'unicode': 'Thé quick brown fox jumps over thé lazy dog',
    'list': ['george', 'jerry', 'elaine', 'cosmo'],
}

# JSON serialization tests
json_data = """\
{"int": 10, "float": 3.1415926500000002, \
"list": ["george", "jerry", "elaine", "cosmo"], \
"string": "The quick brown fox jumps over the lazy \
dog", "unicode": "Th\\u00e9 quick brown fox jumps over \
th\\u00e9 lazy dog"}\
"""

# Pickle serialization tests
pickle_data = pickle.dumps(py_data, protocol=pickle_protocol)

# YAML serialization tests
yaml_data = """\
float: 3.1415926500000002
int: 10
list: [george, jerry, elaine, cosmo]
string: The quick brown fox jumps over the lazy dog
unicode: "Th\\xE9 quick brown fox jumps over th\\xE9 lazy dog"
"""


msgpack_py_data = dict(py_data)
# Unicode chars are lost in transmit :(
msgpack_py_data['unicode'] = 'Th quick brown fox jumps over th lazy dog'
msgpack_data = b64decode(str_to_bytes("""\
haNpbnQKpWZsb2F0y0AJIftTyNTxpGxpc3SUpmdlb3JnZaVqZXJyeaZlbGFpbmWlY29zbW+mc3Rya\
W5n2gArVGhlIHF1aWNrIGJyb3duIGZveCBqdW1wcyBvdmVyIHRoZSBsYXp5IGRvZ6d1bmljb2Rl2g\
ApVGggcXVpY2sgYnJvd24gZm94IGp1bXBzIG92ZXIgdGggbGF6eSBkb2c=\
"""))


def say(m):
    sys.stderr.write('%s\n' % (m, ))


registry.register('testS', lambda s: s, lambda s: 'decoded',
                  'application/testS', 'utf-8')


class test_Serialization(Case):

    def test_disable(self):
        disabled = registry._disabled_content_types
        try:
            registry.disable('testS')
            self.assertIn('application/testS', disabled)
            disabled.clear()

            registry.disable('application/testS')
            self.assertIn('application/testS', disabled)
        finally:
            disabled.clear()

    def test_enable(self):
        registry._disabled_content_types.add('application/json')
        registry.enable('json')
        self.assertNotIn('application/json', registry._disabled_content_types)
        registry._disabled_content_types.add('application/json')
        registry.enable('application/json')
        self.assertNotIn('application/json', registry._disabled_content_types)

    def test_loads_when_disabled(self):
        disabled = registry._disabled_content_types
        try:
            registry.disable('testS')

            with self.assertRaises(SerializerNotInstalled):
                loads('xxd', 'application/testS', 'utf-8', force=False)

            ret = loads('xxd', 'application/testS', 'utf-8', force=True)
            self.assertEqual(ret, 'decoded')
        finally:
            disabled.clear()

    def test_loads_when_data_is_None(self):
        loads(None, 'application/testS', 'utf-8')

    def test_content_type_decoding(self):
        self.assertEqual(
            unicode_string,
            loads(unicode_string_as_utf8,
                  content_type='plain/text', content_encoding='utf-8'),
        )
        self.assertEqual(
            latin_string,
            loads(latin_string_as_latin1,
                  content_type='application/data', content_encoding='latin-1'),
        )

    def test_content_type_binary(self):
        self.assertIsInstance(
            loads(unicode_string_as_utf8,
                  content_type='application/data', content_encoding='binary'),
            bytes_t,
        )

        self.assertEqual(
            unicode_string_as_utf8,
            loads(unicode_string_as_utf8,
                  content_type='application/data', content_encoding='binary'),
        )

    def test_content_type_encoding(self):
        # Using the 'raw' serializer
        self.assertEqual(
            unicode_string_as_utf8,
            dumps(unicode_string, serializer='raw')[-1],
        )
        self.assertEqual(
            latin_string_as_utf8,
            dumps(latin_string, serializer='raw')[-1],
        )
        # And again w/o a specific serializer to check the
        # code where we force unicode objects into a string.
        self.assertEqual(
            unicode_string_as_utf8,
            dumps(unicode_string)[-1],
        )
        self.assertEqual(
            latin_string_as_utf8,
            dumps(latin_string)[-1],
        )

    def test_enable_insecure_serializers(self):
        with patch('kombu.serialization.registry') as registry:
            enable_insecure_serializers()
            registry.assert_has_calls([
                call.enable('pickle'), call.enable('yaml'),
                call.enable('msgpack'),
            ])
            registry.enable.side_effect = KeyError()
            enable_insecure_serializers()

        with patch('kombu.serialization.registry') as registry:
            enable_insecure_serializers(['msgpack'])
            registry.assert_has_calls([call.enable('msgpack')])

    def test_disable_insecure_serializers(self):
        with patch('kombu.serialization.registry') as registry:
            registry._decoders = ['pickle', 'yaml', 'doomsday']
            disable_insecure_serializers(allowed=['doomsday'])
            registry.disable.assert_has_calls([call('pickle'), call('yaml')])
            registry.enable.assert_has_calls([call('doomsday')])
            disable_insecure_serializers(allowed=None)
            registry.disable.assert_has_calls([
                call('pickle'), call('yaml'), call('doomsday')
            ])

    def test_reraises_EncodeError(self):
        with self.assertRaises(EncodeError):
            dumps([object()], serializer='json')

    def test_reraises_DecodeError(self):
        with self.assertRaises(DecodeError):
            loads(object(), content_type='application/json',
                  content_encoding='utf-8')

    def test_json_loads(self):
        self.assertEqual(
            py_data,
            loads(json_data,
                  content_type='application/json', content_encoding='utf-8'),
        )

    def test_json_dumps(self):
        self.assertEqual(
            loads(
                dumps(py_data, serializer='json')[-1],
                content_type='application/json',
                content_encoding='utf-8',
            ),
            loads(
                json_data,
                content_type='application/json',
                content_encoding='utf-8',
            ),
        )

    @skip_if_not_module('msgpack', (ImportError, ValueError))
    def test_msgpack_loads(self):
        register_msgpack()
        res = loads(msgpack_data,
                    content_type='application/x-msgpack',
                    content_encoding='binary')
        if sys.version_info[0] < 3:
            for k, v in res.items():
                if isinstance(v, text_t):
                    res[k] = v.encode()
                if isinstance(v, (list, tuple)):
                    res[k] = [i.encode() for i in v]
        self.assertEqual(
            msgpack_py_data,
            res,
        )

    @skip_if_not_module('msgpack', (ImportError, ValueError))
    def test_msgpack_dumps(self):
        register_msgpack()
        self.assertEqual(
            loads(
                dumps(msgpack_py_data, serializer='msgpack')[-1],
                content_type='application/x-msgpack',
                content_encoding='binary',
            ),
            loads(
                msgpack_data,
                content_type='application/x-msgpack',
                content_encoding='binary',
            ),
        )

    @skip_if_not_module('yaml')
    def test_yaml_loads(self):
        register_yaml()
        self.assertEqual(
            py_data,
            loads(yaml_data,
                  content_type='application/x-yaml',
                  content_encoding='utf-8'),
        )

    @skip_if_not_module('yaml')
    def test_yaml_dumps(self):
        register_yaml()
        self.assertEqual(
            loads(
                dumps(py_data, serializer='yaml')[-1],
                content_type='application/x-yaml',
                content_encoding='utf-8',
            ),
            loads(
                yaml_data,
                content_type='application/x-yaml',
                content_encoding='utf-8',
            ),
        )

    def test_pickle_loads(self):
        self.assertEqual(
            py_data,
            loads(pickle_data,
                  content_type='application/x-python-serialize',
                  content_encoding='binary'),
        )

    def test_pickle_dumps(self):
        self.assertEqual(
            pickle.loads(pickle_data),
            pickle.loads(dumps(py_data, serializer='pickle')[-1]),
        )

    def test_register(self):
        register(None, None, None, None)

    def test_unregister(self):
        with self.assertRaises(SerializerNotInstalled):
            unregister('nonexisting')
        dumps('foo', serializer='pickle')
        unregister('pickle')
        with self.assertRaises(SerializerNotInstalled):
            dumps('foo', serializer='pickle')
        register_pickle()

    def test_set_default_serializer_missing(self):
        with self.assertRaises(SerializerNotInstalled):
            registry._set_default_serializer('nonexisting')

    def test_dumps_missing(self):
        with self.assertRaises(SerializerNotInstalled):
            dumps('foo', serializer='nonexisting')

    def test_dumps__no_serializer(self):
        ctyp, cenc, data = dumps(str_to_bytes('foo'))
        self.assertEqual(ctyp, 'application/data')
        self.assertEqual(cenc, 'binary')

    def test_loads__trusted_content(self):
        loads('tainted', 'application/data', 'binary', accept=[])
        loads('tainted', 'application/text', 'utf-8', accept=[])

    def test_loads__not_accepted(self):
        with self.assertRaises(ContentDisallowed):
            loads('tainted', 'application/x-evil', 'binary', accept=[])
        with self.assertRaises(ContentDisallowed):
            loads('tainted', 'application/x-evil', 'binary',
                  accept=['application/x-json'])
        self.assertTrue(
            loads('tainted', 'application/x-doomsday', 'binary',
                  accept=['application/x-doomsday'])
        )

    def test_raw_encode(self):
        self.assertTupleEqual(
            raw_encode('foo'.encode('utf-8')),
            ('application/data', 'binary', 'foo'.encode('utf-8')),
        )

    @mask_modules('yaml')
    def test_register_yaml__no_yaml(self):
        register_yaml()
        with self.assertRaises(SerializerNotInstalled):
            loads('foo', 'application/x-yaml', 'utf-8')

    @mask_modules('msgpack')
    def test_register_msgpack__no_msgpack(self):
        register_msgpack()
        with self.assertRaises(SerializerNotInstalled):
            loads('foo', 'application/x-msgpack', 'utf-8')