From 55ff270f167d36cd67c637332d7db9ad1b5c68ce Mon Sep 17 00:00:00 2001
From: Jeppe Klitgaard <jeppe@dapj.dk>
Date: Thu, 16 Apr 2015 23:38:02 +0200
Subject: [PATCH] Tests for VarInt

Removed unused import (SerializationError)

We now check if Datatype.SIZE is a number using ABC.

Added 2 decorators for raising serialized and deserialized data
exceptions.

Datatype.SIZE can now be either a number or a sequence. If sequence,
first value is MIN_SIZE and second value is MAX_SIZE

VarInt.serialize now raises ValueError instead of SerializationError
when number is too big to serialize.
---
 minecraft/networking/datatypes.py | 69 ++++++++++++++++++++++++++-----
 tests/test_datatypes.py           |  6 +++
 2 files changed, 65 insertions(+), 10 deletions(-)

diff --git a/minecraft/networking/datatypes.py b/minecraft/networking/datatypes.py
index e36db06..9410ca6 100644
--- a/minecraft/networking/datatypes.py
+++ b/minecraft/networking/datatypes.py
@@ -18,19 +18,52 @@ __all__ = ["ENDIANNESS",
            "VarInt", "VarLong",
            "String"]
 
-from minecraft.exceptions import DeserializationError, SerializationError
+from minecraft.exceptions import DeserializationError
 from minecraft.compat import long
 from io import BytesIO
 import struct
 import collections
+import numbers
 
 ENDIANNESS = "!"  # Network, big-endian
 
 
+def raise_serialization_data(func):
+    """
+    A decorator to be used on a ``Datatype``.serialize definition.
+
+    Must be placed before a classmethod decorator.
+    """
+    def wrapped(cls, data):
+        cls.raise_serialization_data(data)
+
+        return func(cls, data)
+
+    return wrapped
+
+
+def raise_deserialization_data(func):
+    """
+    A decorator to be used on a ``Datatype``.serialize definition.
+
+    Must be placed before a classmethod decorator.
+    """
+    def wrapped(cls, data):
+        cls.raise_deserialization_data(data)
+
+        return func(cls, data)
+
+    return wrapped
+
+
 class Datatype(object):
     """
     Base object for all `pyminecraft` networking datatypes.
 
+    ``Datatype``.SIZE can be either a number, specifying an exact required size
+    of data to be deserialized, or it can be a tuple like this:
+    ``(MIN_SIZE, MAX_SIZE)``
+
 
     .. note::
         If ``ALLOWED_SERIALIZATION_TYPES`` is not empty, only the types found
@@ -65,9 +98,8 @@ class Datatype(object):
         return cls.deserialize(bin_data)
 
     @classmethod
+    @raise_deserialization_data
     def deserialize(cls, data):
-        cls.raise_deserialization_data(data)
-
         deserialized_data = struct.unpack(ENDIANNESS + cls.FORMAT, data)[0]
         return deserialized_data
 
@@ -76,9 +108,8 @@ class Datatype(object):
         return fileobject.write(cls.serialize(data))
 
     @classmethod
+    @raise_serialization_data
     def serialize(cls, data):
-        cls.raise_serialization_data(data)
-
         serialized_data = struct.pack(ENDIANNESS + cls.FORMAT, data)
         return serialized_data
 
@@ -137,11 +168,23 @@ class Datatype(object):
 
             raise TypeError(err)
 
-        if cls.SIZE != len(data):
-            err = "'data' must have a length of {}, not {}"
-            err = err.format(str(cls.SIZE), str(len(data)))
+        if isinstance(cls.SIZE, numbers.Number):
+            if cls.SIZE != len(data):
+                err = "'data' must have a length of {}, not {}"
+                err = err.format(str(cls.SIZE), str(len(data)))
 
-            raise ValueError(err)
+                raise ValueError(err)
+
+        elif isinstance(cls.SIZE, collections.Sequence):
+            if not cls.SIZE[0] <= len(data) <= cls.SIZE[1]:
+                err = "'data' must have a length between {} and {}, not {}"
+                err = err.format(str(cls.SIZE[0]), str(cls.SIZE[1]),
+                                 str(len(data)))
+
+                raise ValueError(err)
+
+        else:
+            raise TypeError("'cls.SIZE' must be a number or a sequence.")
 
         return None
 
@@ -304,6 +347,8 @@ class VarInt(NumberDatatype):
     # Largest element in SIZE_TABLE, assuming largest element is last.
     MAX_SIZE = list(SIZE_TABLE.items())[-1][-1]
 
+    SIZE = (1, MAX_SIZE)
+
     @classmethod
     def read(cls, fileobject):
         number = 0  # The decoded number
@@ -329,16 +374,18 @@ class VarInt(NumberDatatype):
         return number
 
     @classmethod
+    @raise_deserialization_data
     def deserialize(cls, data):
         data_fileobject = BytesIO(bytes(data))
         return cls.read(data_fileobject)
 
     @classmethod
+    @raise_serialization_data
     def serialize(cls, data):
         if data > cls.SIZE_TABLE[-1][0]:
             name_of_self = str(type(cls))
             e = "Number too big to serialize as {}".format(name_of_self)
-            raise SerializationError(e)
+            raise ValueError(e)
 
         result = bytes()  # Where we store the serialized number
 
@@ -370,6 +417,8 @@ class VarLong(VarInt):
 
     MAX_SIZE = list(SIZE_TABLE.items())[-1][-1]
 
+    SIZE = (1, MAX_SIZE)
+
 
 class String(Datatype):
     FORMAT = "utf-8"
diff --git a/tests/test_datatypes.py b/tests/test_datatypes.py
index 07dfce7..4a1f490 100644
--- a/tests/test_datatypes.py
+++ b/tests/test_datatypes.py
@@ -308,6 +308,12 @@ class DoubleTest(FloatTest):
         (5324342541.72123, b"A\xf3\xd5\xb0P\xdb\x8a(")
     ]
 
+
+class VarIntTest(BaseNumberDatatypeTester):
+    DATATYPE_CLS = VarInt
+
+    INVALID_DESERIALIZATION_VALUES = BASE_INVALID_DESERIALIZATION_VALUES
+
 # def _bin(binstr):
 #     """
 #     Accepts a pretty looking string of binary numbers and