1 #-------------------------------------------------------------------------------
4 # Eli Bendersky (eliben@gmail.com)
5 # This code is in the public domain
6 #-------------------------------------------------------------------------------
8 import unittest2
as unittest
11 from random
import randint
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
,
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
))
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
'')
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')
35 def test_large1(self
):
36 text
= b
'i' * 400 + b
'\x00' + b
'bb'
38 self
.assertEqual(parse_cstring_from_stream(sio
), b
'i' * 400)
39 self
.assertEqual(parse_cstring_from_stream(sio
, 150), b
'i' * 250)
41 def test_large2(self
):
42 text
= self
._make
_random
_bytes
(5000) + b
'\x00' + b
'jujajaja'
44 self
.assertEqual(parse_cstring_from_stream(sio
), text
[:5000])
45 self
.assertEqual(parse_cstring_from_stream(sio
, 2348), text
[2348:5000])
48 class Test_preserve_stream_pos(unittest
.TestCase
):
50 sio
= BytesIO(b
'abcdef')
51 with
preserve_stream_pos(sio
):
53 self
.assertEqual(sio
.tell(), 0)
56 with
preserve_stream_pos(sio
):
58 self
.assertEqual(sio
.tell(), 5)
61 if __name__
== '__main__':