import bz2
import os
import unittest

import six

import bmemcached

if six.PY3:
    from unittest import mock
else:
    import mock


class MemcachedTests(unittest.TestCase):
    def setUp(self):
        self.server = '{}:11211'.format(os.environ['MEMCACHED_HOST'])
        self.client = bmemcached.Client(self.server, 'user', 'password')
        self.bzclient = bmemcached.Client(self.server, 'user', 'password',
                                          compression=bz2)
        self.data = 'this is test data. ' * 32

    def tearDown(self):
        self.client.delete('test_key')
        self.client.delete('test_key2')
        self.client.disconnect_all()
        self.bzclient.disconnect_all()

    def testCompressedData(self):
        self.client.set('test_key', self.data)
        self.assertEqual(self.data, self.client.get('test_key'))

    def testBZ2CompressedData(self):
        self.bzclient.set('test_key', self.data)
        self.assertEqual(self.data, self.bzclient.get('test_key'))

    def testCompressionMissmatch(self):
        self.client.set('test_key', self.data)
        self.bzclient.set('test_key2', self.data)
        self.assertEqual(self.client.get('test_key'),
                         self.bzclient.get('test_key2'))
        self.assertRaises(IOError, self.bzclient.get, 'test_key')

    def testCompressionEnabled(self):
        import zlib
        compression = mock.Mock()
        compression.compress.side_effect = zlib.compress
        compression.decompress.side_effect = zlib.decompress
        for proto in self.client._servers:
            proto.compression = compression
        self.client.set('test_key', self.data)
        self.assertEqual(self.data, self.client.get('test_key'))
        compression.compress.assert_called_with(self.data.encode('ascii'))
        self.assertEqual(1, compression.decompress.call_count)

    def testCompressionDisabled(self):
        compression = mock.Mock()
        for proto in self.client._servers:
            proto.compression = compression
        self.client.set('test_key', self.data, compress_level=0)
        self.assertEqual(self.data, self.client.get('test_key'))
        compression.compress.assert_not_called()
        compression.decompress.assert_not_called()