1*f4ee7fbaSAndroid Build Coastguard Worker# Copyright 2016 The Brotli Authors. All rights reserved. 2*f4ee7fbaSAndroid Build Coastguard Worker# 3*f4ee7fbaSAndroid Build Coastguard Worker# Distributed under MIT license. 4*f4ee7fbaSAndroid Build Coastguard Worker# See file LICENSE for detail or copy at https://opensource.org/licenses/MIT 5*f4ee7fbaSAndroid Build Coastguard Worker 6*f4ee7fbaSAndroid Build Coastguard Workerimport functools 7*f4ee7fbaSAndroid Build Coastguard Workerimport unittest 8*f4ee7fbaSAndroid Build Coastguard Worker 9*f4ee7fbaSAndroid Build Coastguard Workerfrom . import _test_utils 10*f4ee7fbaSAndroid Build Coastguard Workerimport brotli 11*f4ee7fbaSAndroid Build Coastguard Worker 12*f4ee7fbaSAndroid Build Coastguard Worker 13*f4ee7fbaSAndroid Build Coastguard Workerdef _get_original_name(test_data): 14*f4ee7fbaSAndroid Build Coastguard Worker return test_data.split('.compressed')[0] 15*f4ee7fbaSAndroid Build Coastguard Worker 16*f4ee7fbaSAndroid Build Coastguard Worker 17*f4ee7fbaSAndroid Build Coastguard Workerclass TestDecompressor(_test_utils.TestCase): 18*f4ee7fbaSAndroid Build Coastguard Worker 19*f4ee7fbaSAndroid Build Coastguard Worker CHUNK_SIZE = 1 20*f4ee7fbaSAndroid Build Coastguard Worker 21*f4ee7fbaSAndroid Build Coastguard Worker def setUp(self): 22*f4ee7fbaSAndroid Build Coastguard Worker self.decompressor = brotli.Decompressor() 23*f4ee7fbaSAndroid Build Coastguard Worker 24*f4ee7fbaSAndroid Build Coastguard Worker def tearDown(self): 25*f4ee7fbaSAndroid Build Coastguard Worker self.decompressor = None 26*f4ee7fbaSAndroid Build Coastguard Worker 27*f4ee7fbaSAndroid Build Coastguard Worker def _check_decompression(self, test_data): 28*f4ee7fbaSAndroid Build Coastguard Worker # Verify decompression matches the original. 29*f4ee7fbaSAndroid Build Coastguard Worker temp_uncompressed = _test_utils.get_temp_uncompressed_name(test_data) 30*f4ee7fbaSAndroid Build Coastguard Worker original = _get_original_name(test_data) 31*f4ee7fbaSAndroid Build Coastguard Worker self.assertFilesMatch(temp_uncompressed, original) 32*f4ee7fbaSAndroid Build Coastguard Worker 33*f4ee7fbaSAndroid Build Coastguard Worker def _decompress(self, test_data): 34*f4ee7fbaSAndroid Build Coastguard Worker temp_uncompressed = _test_utils.get_temp_uncompressed_name(test_data) 35*f4ee7fbaSAndroid Build Coastguard Worker with open(temp_uncompressed, 'wb') as out_file: 36*f4ee7fbaSAndroid Build Coastguard Worker with open(test_data, 'rb') as in_file: 37*f4ee7fbaSAndroid Build Coastguard Worker read_chunk = functools.partial(in_file.read, self.CHUNK_SIZE) 38*f4ee7fbaSAndroid Build Coastguard Worker for data in iter(read_chunk, b''): 39*f4ee7fbaSAndroid Build Coastguard Worker out_file.write(self.decompressor.process(data)) 40*f4ee7fbaSAndroid Build Coastguard Worker self.assertTrue(self.decompressor.is_finished()) 41*f4ee7fbaSAndroid Build Coastguard Worker 42*f4ee7fbaSAndroid Build Coastguard Worker def _test_decompress(self, test_data): 43*f4ee7fbaSAndroid Build Coastguard Worker self._decompress(test_data) 44*f4ee7fbaSAndroid Build Coastguard Worker self._check_decompression(test_data) 45*f4ee7fbaSAndroid Build Coastguard Worker 46*f4ee7fbaSAndroid Build Coastguard Worker def test_garbage_appended(self): 47*f4ee7fbaSAndroid Build Coastguard Worker with self.assertRaises(brotli.error): 48*f4ee7fbaSAndroid Build Coastguard Worker self.decompressor.process(brotli.compress(b'a') + b'a') 49*f4ee7fbaSAndroid Build Coastguard Worker 50*f4ee7fbaSAndroid Build Coastguard Worker def test_already_finished(self): 51*f4ee7fbaSAndroid Build Coastguard Worker self.decompressor.process(brotli.compress(b'a')) 52*f4ee7fbaSAndroid Build Coastguard Worker with self.assertRaises(brotli.error): 53*f4ee7fbaSAndroid Build Coastguard Worker self.decompressor.process(b'a') 54*f4ee7fbaSAndroid Build Coastguard Worker 55*f4ee7fbaSAndroid Build Coastguard Worker 56*f4ee7fbaSAndroid Build Coastguard Worker_test_utils.generate_test_methods(TestDecompressor, for_decompression=True) 57*f4ee7fbaSAndroid Build Coastguard Worker 58*f4ee7fbaSAndroid Build Coastguard Workerif __name__ == '__main__': 59*f4ee7fbaSAndroid Build Coastguard Worker unittest.main() 60