-
Notifications
You must be signed in to change notification settings - Fork 246
Expand file tree
/
Copy pathtest_unpack.py
More file actions
125 lines (97 loc) · 3.23 KB
/
test_unpack.py
File metadata and controls
125 lines (97 loc) · 3.23 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
import gc
import sys
import weakref
from io import BytesIO
from pytest import mark, raises
from msgpack import ExtType, OutOfData, Unpacker, packb
def test_unpack_array_header_from_file():
f = BytesIO(packb([1, 2, 3, 4]))
unpacker = Unpacker(f)
assert unpacker.read_array_header() == 4
assert unpacker.unpack() == 1
assert unpacker.unpack() == 2
assert unpacker.unpack() == 3
assert unpacker.unpack() == 4
with raises(OutOfData):
unpacker.unpack()
@mark.skipif(
"not hasattr(sys, 'getrefcount') == True",
reason="sys.getrefcount() is needed to pass this test",
)
def test_unpacker_hook_refcnt():
result = []
def hook(x):
result.append(x)
return x
basecnt = sys.getrefcount(hook)
up = Unpacker(object_hook=hook, list_hook=hook)
assert sys.getrefcount(hook) >= basecnt + 2
up.feed(packb([{}]))
up.feed(packb([{}]))
assert up.unpack() == [{}]
assert up.unpack() == [{}]
assert result == [{}, [{}], {}, [{}]]
del up
assert sys.getrefcount(hook) == basecnt
def test_unpacker_ext_hook():
class MyUnpacker(Unpacker):
def __init__(self):
super().__init__(ext_hook=self._hook, raw=False)
def _hook(self, code, data):
if code == 1:
return int(data)
else:
return ExtType(code, data)
unpacker = MyUnpacker()
unpacker.feed(packb({"a": 1}))
assert unpacker.unpack() == {"a": 1}
unpacker.feed(packb({"a": ExtType(1, b"123")}))
assert unpacker.unpack() == {"a": 123}
unpacker.feed(packb({"a": ExtType(2, b"321")}))
assert unpacker.unpack() == {"a": ExtType(2, b"321")}
def test_unpacker_tell():
objects = 1, 2, "abc", "def", "ghi"
packed = b"\x01\x02\xa3abc\xa3def\xa3ghi"
positions = 1, 2, 6, 10, 14
unpacker = Unpacker(BytesIO(packed))
for obj, unp, pos in zip(objects, unpacker, positions):
assert obj == unp
assert pos == unpacker.tell()
def test_unpacker_tell_read_bytes():
objects = 1, "abc", "ghi"
packed = b"\x01\x02\xa3abc\xa3def\xa3ghi"
raw_data = b"\x02", b"\xa3def", b""
lenghts = 1, 4, 999
positions = 1, 6, 14
unpacker = Unpacker(BytesIO(packed))
for obj, unp, pos, n, raw in zip(objects, unpacker, positions, lenghts, raw_data):
assert obj == unp
assert pos == unpacker.tell()
assert unpacker.read_bytes(n) == raw
@mark.skipif(
Unpacker.__module__ == "msgpack.fallback",
reason="specific to C extension reinit leak",
)
def test_unpacker_reinit_clears_partial_state():
refs = []
class Marker:
pass
def hook(code, data):
obj = Marker()
refs.append(weakref.ref(obj))
return obj
unpacker = Unpacker(ext_hook=hook, strict_map_key=False)
# Keep parser state mid-map with a live key object from ext_hook.
# Encodes: [ {ExtType(1, b"a"): <missing value>} ].
unpacker.feed(b"\x91\x81\xd4\x01a")
with raises(OutOfData):
unpacker.unpack()
assert len(refs) == 1
assert refs[0]() is not None
unpacker.__init__()
gc.collect()
assert refs[0]() is None
with raises(OutOfData):
unpacker.unpack()
unpacker.feed(packb({"a": 1}))
assert unpacker.unpack() == {"a": 1}