aboutsummaryrefslogtreecommitdiff
path: root/python/tests/compressor_test.py
blob: 6983c31ce0b11a7c3eff5196354f7095d6e476d8 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
# Copyright 2016 The Brotli Authors. All rights reserved.
#
# Distributed under MIT license.
# See file LICENSE for detail or copy at https://opensource.org/licenses/MIT

import filecmp
import functools
import os
import sys
import types
import unittest

import test_utils

import brotli


TEST_DATA_FILES = [
    'empty',  # Empty file
    '10x10y',  # Small text
    'alice29.txt',  # Large text
    'random_org_10k.bin',  # Small data
    'mapsdatazrh',  # Large data
]
TEST_DATA_PATHS = [os.path.join(test_utils.TESTDATA_DIR, f)
                   for f in TEST_DATA_FILES]


def get_compressed_name(filename):
    return filename + '.compressed'


def get_temp_compressed_name(filename):
    return filename + '.bro'


def get_temp_uncompressed_name(filename):
    return filename + '.unbro'


def bind_method_args(method, *args):
    return lambda self: method(self, *args)


# Do not inherit from unittest.TestCase here to ensure that test methods
# are not run automatically and instead are run as part of a specific
# configuration below.
class _TestCompressor(object):

    def _check_decompression_matches(self, test_data):
        # Write decompression to temp file and verify it matches the original.
        with open(get_temp_uncompressed_name(test_data), 'wb') as out_file:
            with open(get_temp_compressed_name(test_data), 'rb') as in_file:
                out_file.write(brotli.decompress(in_file.read()))
        self.assertTrue(
            filecmp.cmp(get_temp_uncompressed_name(test_data),
                        test_data,
                        shallow=False))

    def _test_single_process(self, test_data):
        # Write single-shot compression to temp file.
        with open(get_temp_compressed_name(test_data), 'wb') as out_file:
            with open(test_data, 'rb') as in_file:
                out_file.write(self.compressor.process(in_file.read()))
            out_file.write(self.compressor.finish())
        self._check_decompression_matches(test_data)

    def _test_multiple_process(self, test_data):
        # Write chunked compression to temp file.
        chunk_size = 2048
        with open(get_temp_compressed_name(test_data), 'wb') as out_file:
            with open(test_data, 'rb') as in_file:
                read_chunk = functools.partial(in_file.read, chunk_size)
                for data in iter(read_chunk, b''):
                    out_file.write(self.compressor.process(data))
            out_file.write(self.compressor.finish())
        self._check_decompression_matches(test_data)

    def _test_multiple_process_and_flush(self, test_data):
        # Write chunked and flushed compression to temp file.
        chunk_size = 2048
        with open(get_temp_compressed_name(test_data), 'wb') as out_file:
            with open(test_data, 'rb') as in_file:
                read_chunk = functools.partial(in_file.read, chunk_size)
                for data in iter(read_chunk, b''):
                    out_file.write(self.compressor.process(data))
                    out_file.write(self.compressor.flush())
            out_file.write(self.compressor.finish())
        self._check_decompression_matches(test_data)


# Add test methods for each test data file.  This makes identifying problems
# with specific compression scenarios easier.
for methodname in [m for m in dir(_TestCompressor) if m.startswith('_test')]:
    for test_data in TEST_DATA_PATHS:
        filename = os.path.splitext(os.path.basename(test_data))[0]
        name = 'test_{method}_{file}'.format(method=methodname, file=filename)
        func = bind_method_args(getattr(_TestCompressor, methodname), test_data)
        setattr(_TestCompressor, name, func)


class _CompressionTestCase(unittest.TestCase):

    def tearDown(self):
        for f in TEST_DATA_PATHS:
            try:
                os.unlink(get_temp_compressed_name(f))
            except OSError:
                pass
            try:
                os.unlink(get_temp_uncompressed_name(f))
            except OSError:
                pass


class TestCompressorQuality1(_TestCompressor, _CompressionTestCase):

    def setUp(self):
        self.compressor = brotli.Compressor(quality=1)


class TestCompressorQuality6(_TestCompressor, _CompressionTestCase):

    def setUp(self):
        self.compressor = brotli.Compressor(quality=6)


class TestCompressorQuality9(_TestCompressor, _CompressionTestCase):

    def setUp(self):
        self.compressor = brotli.Compressor(quality=9)


class TestCompressorQuality11(_TestCompressor, _CompressionTestCase):

    def setUp(self):
        self.compressor = brotli.Compressor(quality=11)


if __name__ == '__main__':
    unittest.main()