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