test_snappy.py 9.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220
  1. #!/usr/bin/env python
  2. #
  3. # Copyright (c) 2011, Andres Moreira <andres@andresmoreira.com>
  4. # All rights reserved.
  5. #
  6. # Redistribution and use in source and binary forms, with or without
  7. # modification, are permitted provided that the following conditions are met:
  8. # * Redistributions of source code must retain the above copyright
  9. # notice, this list of conditions and the following disclaimer.
  10. # * Redistributions in binary form must reproduce the above copyright
  11. # notice, this list of conditions and the following disclaimer in the
  12. # documentation and/or other materials provided with the distribution.
  13. # * Neither the name of the authors nor the
  14. # names of its contributors may be used to endorse or promote products
  15. # derived from this software without specific prior written permission.
  16. #
  17. # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  18. # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  19. # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
  20. # ARE DISCLAIMED. IN NO EVENT SHALL ANDRES MOREIRA BE LIABLE FOR ANY DIRECT,
  21. # INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
  22. # (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
  23. # LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
  24. # ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  25. # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
  26. # SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  27. #
  28. import os
  29. import sys
  30. import random
  31. import snappy
  32. import struct
  33. from unittest import TestCase
  34. class SnappyCompressionTest(TestCase):
  35. def test_simple_compress(self):
  36. text = "hello world!".encode('utf-8')
  37. compressed = snappy.compress(text)
  38. self.assertEqual(text, snappy.uncompress(compressed))
  39. def test_moredata_compress(self):
  40. text = "snappy +" * 1000 + " " + "by " * 1000 + " google"
  41. text = text.encode('utf-8')
  42. compressed = snappy.compress(text)
  43. self.assertEqual(text, snappy.uncompress(compressed))
  44. def test_randombytes_compress(self):
  45. _bytes = repr(os.urandom(1000)).encode('utf-8')
  46. compressed = snappy.compress(_bytes)
  47. self.assertEqual(_bytes, snappy.uncompress(compressed))
  48. def test_randombytes2_compress(self):
  49. _bytes = bytes(os.urandom(10000))
  50. compressed = snappy.compress(_bytes)
  51. self.assertEqual(_bytes, snappy.uncompress(compressed))
  52. def test_uncompress_error(self):
  53. self.assertRaises(snappy.UncompressError, snappy.uncompress,
  54. "hoa".encode('utf-8'))
  55. if sys.version_info[0] == 2:
  56. def test_unicode_compress(self):
  57. text = "hello unicode world!".decode('utf-8')
  58. compressed = snappy.compress(text)
  59. self.assertEqual(text, snappy.uncompress(compressed))
  60. def test_decompress(self):
  61. # decompress == uncompress, just to support compatibility with zlib
  62. text = "hello world!".encode('utf-8')
  63. compressed = snappy.compress(text)
  64. self.assertEqual(text, snappy.decompress(compressed))
  65. def test_big_string(self):
  66. text = ('a'*10000000).encode('utf-8')
  67. compressed = snappy.compress(text)
  68. self.assertEqual(text, snappy.decompress(compressed))
  69. class SnappyValidBufferTest(TestCase):
  70. def test_valid_compressed_buffer(self):
  71. text = "hello world!".encode('utf-8')
  72. compressed = snappy.compress(text)
  73. uncompressed = snappy.uncompress(compressed)
  74. self.assertEqual(text == uncompressed,
  75. snappy.isValidCompressed(compressed))
  76. def test_invalid_compressed_buffer(self):
  77. self.assertFalse(snappy.isValidCompressed(
  78. "not compressed".encode('utf-8')))
  79. class SnappyStreaming(TestCase):
  80. def test_random(self):
  81. for _ in range(100):
  82. compressor = snappy.StreamCompressor()
  83. decompressor = snappy.StreamDecompressor()
  84. data = b""
  85. compressed = b""
  86. for _ in range(random.randint(0, 3)):
  87. chunk = os.urandom(random.randint(0, snappy._CHUNK_MAX * 2))
  88. data += chunk
  89. compressed += compressor.add_chunk(
  90. chunk, compress=random.choice([True, False, None]))
  91. upper_bound = random.choice([256, snappy._CHUNK_MAX * 2])
  92. while compressed:
  93. size = random.randint(0, upper_bound)
  94. chunk, compressed = compressed[:size], compressed[size:]
  95. chunk = decompressor.decompress(chunk)
  96. self.assertEqual(data[:len(chunk)], chunk)
  97. data = data[len(chunk):]
  98. decompressor.flush()
  99. self.assertEqual(len(data), 0)
  100. def test_compression(self):
  101. # test that we can add compressed chunks
  102. compressor = snappy.StreamCompressor()
  103. data = b"\0" * 50
  104. compressed_data = snappy.compress(data)
  105. crc = struct.pack("<L", snappy._masked_crc32c(data))
  106. self.assertEqual(crc, b"\x8f)H\xbd")
  107. self.assertEqual(len(compressed_data), 6)
  108. self.assertEqual(compressor.add_chunk(data, compress=True),
  109. b"\xff\x06\x00\x00sNaPpY"
  110. b"\x00\x0a\x00\x00" + crc + compressed_data)
  111. # test that we can add uncompressed chunks
  112. data = b"\x01" * 50
  113. crc = struct.pack("<L", snappy._masked_crc32c(data))
  114. self.assertEqual(crc, b"\xb2\x14)\x8a")
  115. self.assertEqual(compressor.add_chunk(data, compress=False),
  116. b"\x01\x36\x00\x00" + crc + data)
  117. # test that we can add more data than will fit in one chunk
  118. data = b"\x01" * (snappy._CHUNK_MAX * 2 - 5)
  119. crc1 = struct.pack("<L",
  120. snappy._masked_crc32c(data[:snappy._CHUNK_MAX]))
  121. self.assertEqual(crc1, b"h#6\x8e")
  122. crc2 = struct.pack("<L",
  123. snappy._masked_crc32c(data[snappy._CHUNK_MAX:]))
  124. self.assertEqual(crc2, b"q\x8foE")
  125. self.assertEqual(compressor.add_chunk(data, compress=False),
  126. b"\x01\x04\x00\x01" + crc1 + data[:snappy._CHUNK_MAX] +
  127. b"\x01\xff\xff\x00" + crc2 + data[snappy._CHUNK_MAX:])
  128. def test_decompression(self):
  129. # test that we check for the initial stream identifier
  130. data = b"\x01" * 50
  131. self.assertRaises(snappy.UncompressError,
  132. snappy.StreamDecompressor().decompress,
  133. b"\x01\x36\x00\00" +
  134. struct.pack("<L", snappy._masked_crc32c(data)) + data)
  135. self.assertEqual(
  136. snappy.StreamDecompressor().decompress(
  137. b"\xff\x06\x00\x00sNaPpY"
  138. b"\x01\x36\x00\x00" +
  139. struct.pack("<L", snappy._masked_crc32c(data)) + data),
  140. data)
  141. decompressor = snappy.StreamDecompressor()
  142. decompressor.decompress(b"\xff\x06\x00\x00sNaPpY")
  143. self.assertEqual(
  144. decompressor.copy().decompress(
  145. b"\x01\x36\x00\x00" +
  146. struct.pack("<L", snappy._masked_crc32c(data)) + data),
  147. data)
  148. # test that we throw errors for unknown unskippable chunks
  149. self.assertRaises(snappy.UncompressError,
  150. decompressor.copy().decompress, b"\x03\x01\x00\x00")
  151. # test that we skip unknown skippable chunks
  152. self.assertEqual(b"",
  153. decompressor.copy().decompress(b"\xfe\x01\x00\x00"))
  154. # test that we check CRCs
  155. compressed_data = snappy.compress(data)
  156. real_crc = struct.pack("<L", snappy._masked_crc32c(data))
  157. fake_crc = os.urandom(4)
  158. self.assertRaises(snappy.UncompressError,
  159. decompressor.copy().decompress,
  160. b"\x00\x0a\x00\x00" + fake_crc + compressed_data)
  161. self.assertEqual(
  162. decompressor.copy().decompress(
  163. b"\x00\x0a\x00\x00" + real_crc + compressed_data),
  164. data)
  165. # test that we buffer when we don't have enough
  166. uncompressed_data = os.urandom(100)
  167. compressor = snappy.StreamCompressor()
  168. compressed_data = (compressor.compress(uncompressed_data[:50]) +
  169. compressor.compress(uncompressed_data[50:]))
  170. for split1 in range(len(compressed_data) - 1):
  171. for split2 in range(split1, len(compressed_data)):
  172. decompressor = snappy.StreamDecompressor()
  173. self.assertEqual(
  174. (decompressor.decompress(compressed_data[:split1]) +
  175. decompressor.decompress(compressed_data[split1:split2]) +
  176. decompressor.decompress(compressed_data[split2:])),
  177. uncompressed_data)
  178. def test_concatenation(self):
  179. data1 = os.urandom(snappy._CHUNK_MAX * 2)
  180. data2 = os.urandom(4096)
  181. decompressor = snappy.StreamDecompressor()
  182. self.assertEqual(
  183. decompressor.decompress(
  184. snappy.StreamCompressor().compress(data1) +
  185. snappy.StreamCompressor().compress(data2)),
  186. data1 + data2)
  187. if __name__ == "__main__":
  188. import unittest
  189. unittest.main()