PageRenderTime 55ms CodeModel.GetById 23ms RepoModel.GetById 1ms app.codeStats 0ms

/zfec-1.4.22/zfec/test/test_zfec.py

#
Python | 360 lines | 352 code | 7 blank | 1 comment | 4 complexity | 4621646ba207b9fb2d3ea58b09d9f1ae MD5 | raw file
Possible License(s): GPL-2.0, CC-BY-SA-3.0
  1. #!/usr/bin/env python
  2. import cStringIO, os, random, re
  3. import unittest
  4. global VERBOSE
  5. VERBOSE=False
  6. import zfec
  7. from pyutil import fileutil
  8. from base64 import b32encode
  9. def ab(x): # debuggery
  10. if len(x) >= 3:
  11. return "%s:%s" % (len(x), b32encode(x[-3:]),)
  12. elif len(x) == 2:
  13. return "%s:%s" % (len(x), b32encode(x[-2:]),)
  14. elif len(x) == 1:
  15. return "%s:%s" % (len(x), b32encode(x[-1:]),)
  16. elif len(x) == 0:
  17. return "%s:%s" % (len(x), "--empty--",)
  18. def randstr(n):
  19. return ''.join(map(chr, map(random.randrange, [0]*n, [256]*n)))
  20. def _h(k, m, ss):
  21. encer = zfec.Encoder(k, m)
  22. nums_and_blocks = list(enumerate(encer.encode(ss)))
  23. assert isinstance(nums_and_blocks, list), nums_and_blocks
  24. assert len(nums_and_blocks) == m, (len(nums_and_blocks), m,)
  25. nums_and_blocks = random.sample(nums_and_blocks, k)
  26. blocks = [ x[1] for x in nums_and_blocks ]
  27. nums = [ x[0] for x in nums_and_blocks ]
  28. decer = zfec.Decoder(k, m)
  29. decoded = decer.decode(blocks, nums)
  30. assert len(decoded) == len(ss), (len(decoded), len(ss),)
  31. assert tuple([str(s) for s in decoded]) == tuple([str(s) for s in ss]), (tuple([ab(str(s)) for s in decoded]), tuple([ab(str(s)) for s in ss]),)
  32. def _help_test_random():
  33. m = random.randrange(1, 257)
  34. k = random.randrange(1, m+1)
  35. l = random.randrange(0, 2**9)
  36. ss = [ randstr(l/k) for x in range(k) ]
  37. _h(k, m, ss)
  38. def _help_test_random_with_l(l):
  39. m = random.randrange(1, 257)
  40. k = random.randrange(1, m+1)
  41. ss = [ randstr(l/k) for x in range(k) ]
  42. _h(k, m, ss)
  43. def _h_easy(k, m, s):
  44. encer = zfec.easyfec.Encoder(k, m)
  45. nums_and_blocks = list(enumerate(encer.encode(s)))
  46. assert isinstance(nums_and_blocks, list), nums_and_blocks
  47. assert len(nums_and_blocks) == m, (len(nums_and_blocks), m,)
  48. nums_and_blocks = random.sample(nums_and_blocks, k)
  49. blocks = [ x[1] for x in nums_and_blocks ]
  50. nums = [ x[0] for x in nums_and_blocks ]
  51. decer = zfec.easyfec.Decoder(k, m)
  52. decodeds = decer.decode(blocks, nums, padlen=k*len(blocks[0]) - len(s))
  53. assert len(decodeds) == len(s), (ab(decodeds), ab(s), k, m)
  54. assert decodeds == s, (ab(decodeds), ab(s),)
  55. def _help_test_random_easy():
  56. m = random.randrange(1, 257)
  57. k = random.randrange(1, m+1)
  58. l = random.randrange(0, 2**9)
  59. s = randstr(l)
  60. _h_easy(k, m, s)
  61. def _help_test_random_with_l_easy(l):
  62. m = random.randrange(1, 257)
  63. k = random.randrange(1, m+1)
  64. s = randstr(l)
  65. _h_easy(k, m, s)
  66. class ZFecTest(unittest.TestCase):
  67. def test_instantiate_encoder_no_args(self):
  68. try:
  69. e = zfec.Encoder()
  70. except TypeError:
  71. # Okay, so that's because we're required to pass constructor args.
  72. pass
  73. else:
  74. # Oops, it should have raised an exception.
  75. self.fail("Should have raised exception from incorrect arguments to constructor.")
  76. def test_instantiate_decoder_no_args(self):
  77. try:
  78. e = zfec.Decoder()
  79. except TypeError:
  80. # Okay, so that's because we're required to pass constructor args.
  81. pass
  82. else:
  83. # Oops, it should have raised an exception.
  84. self.fail("Should have raised exception from incorrect arguments to constructor.")
  85. def test_from_agl_c(self):
  86. self.failUnless(zfec._fec.test_from_agl())
  87. def test_from_agl_py(self):
  88. e = zfec.Encoder(3, 5)
  89. b0 = '\x01'*8 ; b1 = '\x02'*8 ; b2 = '\x03'*8
  90. # print "_from_py before encoding:"
  91. # print "b0: %s, b1: %s, b2: %s" % tuple(base64.b16encode(x) for x in [b0, b1, b2])
  92. b3, b4 = e.encode([b0, b1, b2], (3, 4))
  93. # print "after encoding:"
  94. # print "b3: %s, b4: %s" % tuple(base64.b16encode(x) for x in [b3, b4])
  95. d = zfec.Decoder(3, 5)
  96. r0, r1, r2 = d.decode((b2, b3, b4), (1, 2, 3))
  97. # print "after decoding:"
  98. # print "b0: %s, b1: %s" % tuple(base64.b16encode(x) for x in [b0, b1])
  99. def test_small(self):
  100. for i in range(16):
  101. _help_test_random_with_l(i)
  102. if VERBOSE:
  103. print "%d randomized tests pass." % (i+1)
  104. def test_random(self):
  105. for i in range(3):
  106. _help_test_random()
  107. if VERBOSE:
  108. print "%d randomized tests pass." % (i+1)
  109. def test_bad_args_construct_decoder(self):
  110. try:
  111. zfec.Decoder(-1, -1)
  112. except zfec.Error, e:
  113. assert "argument is required to be greater than or equal to 1" in str(e), e
  114. else:
  115. self.fail("Should have gotten an exception from out-of-range arguments.")
  116. try:
  117. zfec.Decoder(1, 257)
  118. except zfec.Error, e:
  119. assert "argument is required to be less than or equal to 256" in str(e), e
  120. else:
  121. self.fail("Should have gotten an exception from out-of-range arguments.")
  122. try:
  123. zfec.Decoder(3, 2)
  124. except zfec.Error, e:
  125. assert "first argument is required to be less than or equal to the second argument" in str(e), e
  126. else:
  127. self.fail("Should have gotten an exception from out-of-range arguments.")
  128. def test_bad_args_construct_encoder(self):
  129. try:
  130. zfec.Encoder(-1, -1)
  131. except zfec.Error, e:
  132. assert "argument is required to be greater than or equal to 1" in str(e), e
  133. else:
  134. self.fail("Should have gotten an exception from out-of-range arguments.")
  135. try:
  136. zfec.Encoder(1, 257)
  137. except zfec.Error, e:
  138. assert "argument is required to be less than or equal to 256" in str(e), e
  139. else:
  140. self.fail("Should have gotten an exception from out-of-range arguments.")
  141. def test_bad_args_dec(self):
  142. decer = zfec.Decoder(2, 4)
  143. try:
  144. decer.decode(98, []) # first argument is not a sequence
  145. except TypeError, e:
  146. assert "First argument was not a sequence" in str(e), e
  147. else:
  148. self.fail("Should have gotten TypeError for wrong type of second argument.")
  149. try:
  150. decer.decode(["a", "b", ], ["c", "d",])
  151. except zfec.Error, e:
  152. assert "Precondition violation: second argument is required to contain int" in str(e), e
  153. else:
  154. self.fail("Should have gotten zfec.Error for wrong type of second argument.")
  155. try:
  156. decer.decode(["a", "b", ], 98) # not a sequence at all
  157. except TypeError, e:
  158. assert "Second argument was not a sequence" in str(e), e
  159. else:
  160. self.fail("Should have gotten TypeError for wrong type of second argument.")
  161. class EasyFecTest(unittest.TestCase):
  162. def test_small(self):
  163. for i in range(16):
  164. _help_test_random_with_l_easy(i)
  165. if VERBOSE:
  166. print "%d randomized tests pass." % (i+1)
  167. def test_random(self):
  168. for i in range(3):
  169. _help_test_random_easy()
  170. if VERBOSE:
  171. print "%d randomized tests pass." % (i+1)
  172. def test_bad_args_dec(self):
  173. decer = zfec.easyfec.Decoder(2, 4)
  174. try:
  175. decer.decode(98, [0, 1], 0) # first argument is not a sequence
  176. except TypeError, e:
  177. assert "First argument was not a sequence" in str(e), e
  178. else:
  179. self.fail("Should have gotten TypeError for wrong type of second argument.")
  180. try:
  181. decer.decode("ab", ["c", "d",], 0)
  182. except zfec.Error, e:
  183. assert "Precondition violation: second argument is required to contain int" in str(e), e
  184. else:
  185. self.fail("Should have gotten zfec.Error for wrong type of second argument.")
  186. try:
  187. decer.decode("ab", 98, 0) # not a sequence at all
  188. except TypeError, e:
  189. assert "Second argument was not a sequence" in str(e), e
  190. else:
  191. self.fail("Should have gotten TypeError for wrong type of second argument.")
  192. class FileFec(unittest.TestCase):
  193. def test_filefec_header(self):
  194. for m in [1, 2, 3, 5, 7, 9, 11, 17, 19, 33, 35, 65, 66, 67, 129, 130, 131, 254, 255, 256,]:
  195. for k in [1, 2, 3, 5, 9, 17, 33, 65, 129, 255, 256,]:
  196. if k >= m:
  197. continue
  198. for pad in [0, 1, k-1,]:
  199. if pad >= k:
  200. continue
  201. for sh in [0, 1, m-1,]:
  202. if sh >= m:
  203. continue
  204. h = zfec.filefec._build_header(m, k, pad, sh)
  205. hio = cStringIO.StringIO(h)
  206. (rm, rk, rpad, rsh,) = zfec.filefec._parse_header(hio)
  207. assert (rm, rk, rpad, rsh,) == (m, k, pad, sh,), h
  208. def _help_test_filefec(self, teststr, k, m, numshs=None):
  209. if numshs == None:
  210. numshs = m
  211. TESTFNAME = "testfile.txt"
  212. PREFIX = "test"
  213. SUFFIX = ".fec"
  214. fsize = len(teststr)
  215. tempdir = fileutil.NamedTemporaryDirectory(cleanup=True)
  216. try:
  217. tempf = tempdir.file(TESTFNAME, 'w+b')
  218. tempf.write(teststr)
  219. tempf.flush()
  220. tempf.seek(0)
  221. # encode the file
  222. zfec.filefec.encode_to_files(tempf, fsize, tempdir.name, PREFIX, k, m, SUFFIX, verbose=VERBOSE)
  223. # select some share files
  224. RE=re.compile(zfec.filefec.RE_FORMAT % (PREFIX, SUFFIX,))
  225. fns = os.listdir(tempdir.name)
  226. assert len(fns) >= m, (fns, tempdir, tempdir.name,)
  227. sharefs = [ open(os.path.join(tempdir.name, fn), "rb") for fn in fns if RE.match(fn) ]
  228. for sharef in sharefs:
  229. tempdir.register_file(sharef)
  230. random.shuffle(sharefs)
  231. del sharefs[numshs:]
  232. # decode from the share files
  233. outf = tempdir.file('recovered-testfile.txt', 'w+b')
  234. zfec.filefec.decode_from_files(outf, sharefs, verbose=VERBOSE)
  235. outf.flush()
  236. outf.seek(0)
  237. recovereddata = outf.read()
  238. assert recovereddata == teststr, (ab(recovereddata), ab(teststr),)
  239. finally:
  240. tempdir.shutdown()
  241. def test_filefec_all_shares(self):
  242. return self._help_test_filefec("Yellow Whirled!", 3, 8)
  243. def test_filefec_all_shares_1_b(self):
  244. return self._help_test_filefec("Yellow Whirled!", 4, 16)
  245. def test_filefec_all_shares_2(self):
  246. return self._help_test_filefec("Yellow Whirled", 3, 8)
  247. def test_filefec_all_shares_2_b(self):
  248. return self._help_test_filefec("Yellow Whirled", 4, 16)
  249. def test_filefec_all_shares_3(self):
  250. return self._help_test_filefec("Yellow Whirle", 3, 8)
  251. def test_filefec_all_shares_3_b(self):
  252. return self._help_test_filefec("Yellow Whirle", 4, 16)
  253. def test_filefec_all_shares_with_padding(self, noisy=VERBOSE):
  254. return self._help_test_filefec("Yellow Whirled!A", 3, 8)
  255. def test_filefec_min_shares_with_padding(self, noisy=VERBOSE):
  256. return self._help_test_filefec("Yellow Whirled!A", 3, 8, numshs=3)
  257. def test_filefec_min_shares_with_crlf(self, noisy=VERBOSE):
  258. return self._help_test_filefec("llow Whirled!A\r\n", 3, 8, numshs=3)
  259. def test_filefec_min_shares_with_lf(self, noisy=VERBOSE):
  260. return self._help_test_filefec("Yellow Whirled!A\n", 3, 8, numshs=3)
  261. def test_filefec_min_shares_with_lflf(self, noisy=VERBOSE):
  262. return self._help_test_filefec("Yellow Whirled!A\n\n", 3, 8, numshs=3)
  263. def test_filefec_min_shares_with_crcrlflf(self, noisy=VERBOSE):
  264. return self._help_test_filefec("Yellow Whirled!A\r\r\n\n", 3, 8, numshs=3)
  265. class Cmdline(unittest.TestCase):
  266. def test_basic(self, noisy=VERBOSE):
  267. tempdir = fileutil.NamedTemporaryDirectory(cleanup=True)
  268. fo = tempdir.file("test.data", "w+b")
  269. fo.write("WHEHWHJEKWAHDLJAWDHWALKDHA")
  270. import sys
  271. realargv = sys.argv
  272. try:
  273. DEFAULT_M=8
  274. DEFAULT_K=3
  275. sys.argv = ["zfec", os.path.join(tempdir.name, "test.data"),]
  276. retcode = zfec.cmdline_zfec.main()
  277. assert retcode == 0, retcode
  278. RE=re.compile(zfec.filefec.RE_FORMAT % ('test.data', ".fec",))
  279. fns = os.listdir(tempdir.name)
  280. assert len(fns) >= DEFAULT_M, (fns, DEFAULT_M, tempdir, tempdir.name,)
  281. sharefns = [ os.path.join(tempdir.name, fn) for fn in fns if RE.match(fn) ]
  282. random.shuffle(sharefns)
  283. del sharefns[DEFAULT_K:]
  284. sys.argv = ["zunfec",]
  285. sys.argv.extend(sharefns)
  286. sys.argv.extend(['-o', os.path.join(tempdir.name, 'test.data-recovered'),])
  287. retcode = zfec.cmdline_zunfec.main()
  288. assert retcode == 0, retcode
  289. import filecmp
  290. assert filecmp.cmp(os.path.join(tempdir.name, 'test.data'), os.path.join(tempdir.name, 'test.data-recovered'))
  291. finally:
  292. sys.argv = realargv
  293. if __name__ == "__main__":
  294. unittest.main()