test/test_utils.py

   1 try:
   2     import unittest2 as unittest
   3 except ImportError:
   4     import unittest
   5 import sys
   6 from random import randint
   7
   8 sys.path.extend(['.', '..'])
   9 from elftools.common.py3compat import int2byte, BytesIO
  10 from elftools.common.utils import (parse_cstring_from_stream,
  11         preserve_stream_pos)
  12
  13
  14 class Test_parse_cstring_from_stream(unittest.TestCase):
  15     def _make_random_bytes(self, n):
  16         return b''.join(int2byte(randint(32, 127)) for i in range(n))
  17
  18     def test_small1(self):
  19         sio = BytesIO(b'abcdefgh\x0012345')
  20         self.assertEqual(parse_cstring_from_stream(sio), b'abcdefgh')
  21         self.assertEqual(parse_cstring_from_stream(sio, 2), b'cdefgh')
  22         self.assertEqual(parse_cstring_from_stream(sio, 8), b'')
  23
  24     def test_small2(self):
  25         sio = BytesIO(b'12345\x006789\x00abcdefg\x00iii')
  26         self.assertEqual(parse_cstring_from_stream(sio), b'12345')
  27         self.assertEqual(parse_cstring_from_stream(sio, 5), b'')
  28         self.assertEqual(parse_cstring_from_stream(sio, 6), b'6789')
  29
  30     def test_large1(self):
  31         text = b'i' * 400 + b'\x00' + b'bb'
  32         sio = BytesIO(text)
  33         self.assertEqual(parse_cstring_from_stream(sio), b'i' * 400)
  34         self.assertEqual(parse_cstring_from_stream(sio, 150), b'i' * 250)
  35
  36     def test_large2(self):
  37         text = self._make_random_bytes(5000) + b'\x00' + b'jujajaja'
  38         sio = BytesIO(text)
  39         self.assertEqual(parse_cstring_from_stream(sio), text[:5000])
  40         self.assertEqual(parse_cstring_from_stream(sio, 2348), text[2348:5000])
  41
  42
  43 class Test_preserve_stream_pos(object):
  44     def test_basic(self):
  45         sio = BytesIO('abcdef')
  46         with preserve_stream_pos(sio):
  47             sio.seek(4)
  48         self.assertEqual(stream.tell(), 0)
  49
  50         sio.seek(5)
  51         with preserve_stream_pos(sio):
  52             sio.seek(0)
  53         self.assertEqual(stream.tell(), 5)
  54
  55
  56 if __name__ == '__main__':
  57     unittest.main()
  58
  59
  60