2 import unittest2 as unittest
6 from random import randint
8 sys.path.extend(['.', '..'])
9 from elftools.common.py3compat import int2byte, BytesIO
10 from elftools.common.utils import (parse_cstring_from_stream,
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))
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'')
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')
30 def test_large1(self):
31 text = b'i' * 400 + b'\x00' + b'bb'
33 self.assertEqual(parse_cstring_from_stream(sio), b'i' * 400)
34 self.assertEqual(parse_cstring_from_stream(sio, 150), b'i' * 250)
36 def test_large2(self):
37 text = self._make_random_bytes(5000) + b'\x00' + b'jujajaja'
39 self.assertEqual(parse_cstring_from_stream(sio), text[:5000])
40 self.assertEqual(parse_cstring_from_stream(sio, 2348), text[2348:5000])
43 class Test_preserve_stream_pos(object):
45 sio = BytesIO('abcdef')
46 with preserve_stream_pos(sio):
48 self.assertEqual(stream.tell(), 0)
51 with preserve_stream_pos(sio):
53 self.assertEqual(stream.tell(), 5)
56 if __name__ == '__main__':