Merge pull request #105 from msgpack/max-xxx-size

Add max_<type>_len option to unpacker. (fixes #97).

1 files changed, 72 insertions, 1 deletions

diff --git a/test/test_limits.py b/test/test_limits.py
index 1cfa2d6..3c1cf2a 100644
--- a/test/test_limits.py
+++ b/test/test_limits.py
@@ -3,7 +3,7 @@
 from __future__ import absolute_import, division, print_function, unicode_literals
 import pytest
 
-from msgpack import packb, unpackb, Packer
+from msgpack import packb, unpackb, Packer, Unpacker, ExtType
 
 
 def test_integer():
@@ -32,6 +32,77 @@ def test_map_header():
         packer.pack_array_header(2**32)
 
 
+def test_max_str_len():
+    d = 'x' * 3
+    packed = packb(d)
+
+    unpacker = Unpacker(max_str_len=3, encoding='utf-8')
+    unpacker.feed(packed)
+    assert unpacker.unpack() == d
+
+    unpacker = Unpacker(max_str_len=2, encoding='utf-8')
+    with pytest.raises(ValueError):
+        unpacker.feed(packed)
+        unpacker.unpack()
+
+
+def test_max_bin_len():
+    d = b'x' * 3
+    packed = packb(d, use_bin_type=True)
+
+    unpacker = Unpacker(max_bin_len=3)
+    unpacker.feed(packed)
+    assert unpacker.unpack() == d
+
+    unpacker = Unpacker(max_bin_len=2)
+    with pytest.raises(ValueError):
+        unpacker.feed(packed)
+        unpacker.unpack()
+
+
+def test_max_array_len():
+    d = [1,2,3]
+    packed = packb(d)
+
+    unpacker = Unpacker(max_array_len=3)
+    unpacker.feed(packed)
+    assert unpacker.unpack() == d
+
+    unpacker = Unpacker(max_array_len=2)
+    with pytest.raises(ValueError):
+        unpacker.feed(packed)
+        unpacker.unpack()
+
+
+def test_max_map_len():
+    d = {1: 2, 3: 4, 5: 6}
+    packed = packb(d)
+
+    unpacker = Unpacker(max_map_len=3)
+    unpacker.feed(packed)
+    assert unpacker.unpack() == d
+
+    unpacker = Unpacker(max_map_len=2)
+    with pytest.raises(ValueError):
+        unpacker.feed(packed)
+        unpacker.unpack()
+
+
+def test_max_ext_len():
+    d = ExtType(42, b"abc")
+    packed = packb(d)
+
+    unpacker = Unpacker(max_ext_len=3)
+    unpacker.feed(packed)
+    assert unpacker.unpack() == d
+
+    unpacker = Unpacker(max_ext_len=2)
+    with pytest.raises(ValueError):
+        unpacker.feed(packed)
+        unpacker.unpack()
+
+
+
 # PyPy fails following tests because of constant folding?
 # https://bugs.pypy.org/issue1721
 #@pytest.mark.skipif(True, reason="Requires very large memory.")