Python: Add tests for streamed compression (#458)

Progress on #191

2 files changed, 159 insertions, 15 deletions

diff --git a/python/tests/compressor_test.py b/python/tests/compressor_test.py
new file mode 100644
index 0000000..6983c31
--- /dev/null
+++ b/python/tests/compressor_test.py
@@ -0,0 +1,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()
diff --git a/python/tests/test_utils.py b/python/tests/test_utils.py
index 733f7b5..2a62a2f 100644
--- a/python/tests/test_utils.py
+++ b/python/tests/test_utils.py
@@ -8,29 +8,32 @@ import filecmp
 def diff_q(first_file, second_file):
     """Simulate call to POSIX diff with -q argument"""
     if not filecmp.cmp(first_file, second_file, shallow=False):
-        print("Files %s and %s differ" % (first_file, second_file),
-              file=sys.stderr)
+        print(
+            'Files %s and %s differ' % (first_file, second_file),
+            file=sys.stderr)
         return 1
     return 0
 
 
-PYTHON = sys.executable or "python"
+project_dir = os.path.abspath(os.path.join(__file__, '..', '..', '..'))
 
-# 'bro.py' script should be in parent directory
-BRO = os.path.abspath("../bro.py")
+PYTHON = sys.executable or 'python'
 
-# get platform- and version-specific build/lib folder
-platform_lib_name = "lib.{platform}-{version[0]}.{version[1]}".format(
-    platform=sysconfig.get_platform(),
-    version=sys.version_info)
+BRO = os.path.join(project_dir, 'python', 'bro.py')
 
-# by default, distutils' build base is in the same location as setup.py
-build_base = os.path.abspath(os.path.join("..", "..", "bin"))
-build_lib = os.path.join(build_base, platform_lib_name)
+TESTDATA_DIR = os.path.join(project_dir, 'tests', 'testdata')
 
-# prepend build/lib to PYTHONPATH environment variable
+# Get the platform/version-specific build folder.
+# By default, the distutils build base is in the same location as setup.py.
+platform_lib_name = 'lib.{platform}-{version[0]}.{version[1]}'.format(
+    platform=sysconfig.get_platform(), version=sys.version_info)
+build_dir = os.path.join(project_dir, 'bin', platform_lib_name)
+
+# Prepend the build folder to sys.path and the PYTHONPATH environment variable.
+if build_dir not in sys.path:
+    sys.path.insert(0, build_dir)
 TEST_ENV = os.environ.copy()
 if 'PYTHONPATH' not in TEST_ENV:
-    TEST_ENV['PYTHONPATH'] = build_lib
+    TEST_ENV['PYTHONPATH'] = build_dir
 else:
-    TEST_ENV['PYTHONPATH'] = build_lib + os.pathsep + TEST_ENV['PYTHONPATH']
+    TEST_ENV['PYTHONPATH'] = build_dir + os.pathsep + TEST_ENV['PYTHONPATH']