1import io 2import os 3from unittest import TestCase 4 5from snappy import snappy_formats as formats 6from snappy.snappy import _CHUNK_MAX, UncompressError 7 8 9class TestFormatBase(TestCase): 10 compress_format = formats.FORMAT_AUTO 11 decompress_format = formats.FORMAT_AUTO 12 success = True 13 14 def runTest(self): 15 data = os.urandom(1024 * 256 * 2) + os.urandom(13245 * 2) 16 compress_func = formats.get_compress_function(self.compress_format) 17 instream = io.BytesIO(data) 18 compressed_stream = io.BytesIO() 19 compress_func(instream, compressed_stream) 20 compressed_stream.seek(0) 21 if not self.success: 22 with self.assertRaises(UncompressError) as err: 23 decompress_func, read_chunk = formats.get_decompress_function( 24 self.decompress_format, compressed_stream 25 ) 26 decompressed_stream = io.BytesIO() 27 decompress_func( 28 compressed_stream, 29 decompressed_stream, 30 start_chunk=read_chunk 31 ) 32 return 33 decompress_func, read_chunk = formats.get_decompress_function( 34 self.decompress_format, compressed_stream 35 ) 36 decompressed_stream = io.BytesIO() 37 decompress_func( 38 compressed_stream, 39 decompressed_stream, 40 start_chunk=read_chunk 41 ) 42 decompressed_stream.seek(0) 43 self.assertEqual(data, decompressed_stream.read()) 44 45 46class TestFormatFramingFraming(TestFormatBase): 47 compress_format = formats.FRAMING_FORMAT 48 decompress_format = formats.FRAMING_FORMAT 49 success = True 50 51 52class TestFormatFramingHadoop(TestFormatBase): 53 compress_format = formats.FRAMING_FORMAT 54 decompress_format = formats.HADOOP_FORMAT 55 success = False 56 57 58class TestFormatFramingAuto(TestFormatBase): 59 compress_format = formats.FRAMING_FORMAT 60 decompress_format = formats.FORMAT_AUTO 61 success = True 62 63 64class TestFormatHadoopHadoop(TestFormatBase): 65 compress_format = formats.HADOOP_FORMAT 66 decompress_format = formats.HADOOP_FORMAT 67 success = True 68 69 70class TestFormatHadoopFraming(TestFormatBase): 71 compress_format = formats.HADOOP_FORMAT 72 decompress_format = formats.FRAMING_FORMAT 73 success = False 74 75 76class TestFormatHadoopAuto(TestFormatBase): 77 compress_format = formats.HADOOP_FORMAT 78 decompress_format = formats.FORMAT_AUTO 79 success = True 80 81 82class TestFormatAutoFraming(TestFormatBase): 83 compress_format = formats.FORMAT_AUTO 84 decompress_format = formats.FRAMING_FORMAT 85 success = True 86 87 88class TestFormatAutoHadoop(TestFormatBase): 89 compress_format = formats.FORMAT_AUTO 90 decompress_format = formats.HADOOP_FORMAT 91 success = False 92 93 94if __name__ == "__main__": 95 import unittest 96 unittest.main() 97