test/test_utils.py

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