From 6dc7985dc7c852f49ecae8d98d2cb273081f295e Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Thu, 22 Aug 2019 16:39:23 +0200
Subject: [PATCH] python: add more protobuf tests

also fix field numbers to start at 1
---
 python/tests/test_protobuf.py | 84 +++++++++++++++++++++++++++++++----
 1 file changed, 76 insertions(+), 8 deletions(-)

diff --git a/python/tests/test_protobuf.py b/python/tests/test_protobuf.py
index b8e51b855..cfe024912 100644
--- a/python/tests/test_protobuf.py
+++ b/python/tests/test_protobuf.py
@@ -26,25 +26,35 @@ class PrimitiveMessage(protobuf.MessageType):
     @classmethod
     def get_fields(cls):
         return {
-            0: ("uvarint", protobuf.UVarintType, 0),
-            1: ("svarint", protobuf.SVarintType, 0),
-            2: ("bool", protobuf.BoolType, 0),
-            3: ("bytes", protobuf.BytesType, 0),
-            4: ("unicode", protobuf.UnicodeType, 0),
-            5: ("enum", protobuf.EnumType("t", (0, 5, 25)), 0),
+            1: ("uvarint", protobuf.UVarintType, 0),
+            2: ("svarint", protobuf.SVarintType, 0),
+            3: ("bool", protobuf.BoolType, 0),
+            4: ("bytes", protobuf.BytesType, 0),
+            5: ("unicode", protobuf.UnicodeType, 0),
+            6: ("enum", protobuf.EnumType("t", (0, 5, 25)), 0),
         }
 
 
 class EnumMessageMoreValues(protobuf.MessageType):
     @classmethod
     def get_fields(cls):
-        return {0: ("enum", protobuf.EnumType("t", (0, 1, 2, 3, 4, 5)), 0)}
+        return {1: ("enum", protobuf.EnumType("t", (0, 1, 2, 3, 4, 5)), 0)}
 
 
 class EnumMessageLessValues(protobuf.MessageType):
     @classmethod
     def get_fields(cls):
-        return {0: ("enum", protobuf.EnumType("t", (0, 5)), 0)}
+        return {1: ("enum", protobuf.EnumType("t", (0, 5)), 0)}
+
+
+class RepeatedFields(protobuf.MessageType):
+    @classmethod
+    def get_fields(cls):
+        return {
+            1: ("uintlist", protobuf.UVarintType, protobuf.FLAG_REPEATED),
+            2: ("enumlist", protobuf.EnumType("t", (0, 1)), protobuf.FLAG_REPEATED),
+            3: ("strlist", protobuf.UnicodeType, protobuf.FLAG_REPEATED),
+        }
 
 
 def load_uvarint(buffer):
@@ -73,6 +83,12 @@ def test_load_uvarint():
     assert load_uvarint(b"\x01") == 1
     assert load_uvarint(b"\xff\x01") == 0xFF
     assert load_uvarint(b"\xc0\xc4\x07") == 123456
+    assert load_uvarint(b"\x80\x80\x80\x80\x00") == 0
+
+
+def test_broken_uvarint():
+    with pytest.raises(IOError):
+        load_uvarint(b"\x80\x80")
 
 
 def test_sint_uint():
@@ -170,3 +186,55 @@ def test_enum_to_str():
 
     with pytest.raises(TypeError):
         enum_type.to_str(999)
+
+
+def test_repeated():
+    msg = RepeatedFields(
+        uintlist=[1, 2, 3], enumlist=[0, 1, 0, 1], strlist=["hello", "world"]
+    )
+    buf = BytesIO()
+    protobuf.dump_message(buf, msg)
+    buf.seek(0)
+    retr = protobuf.load_message(buf, RepeatedFields)
+
+    assert retr == msg
+
+
+def test_enum_in_repeated(caplog):
+    msg = RepeatedFields(enumlist=[0, 1, 2, 3])
+    buf = BytesIO()
+    protobuf.dump_message(buf, msg)
+    assert len(caplog.records) == 2
+    for record in caplog.records:
+        assert record.levelname == "WARNING"
+        assert "unknown for type t" in record.getMessage()
+
+
+def test_packed():
+    values = [4, 44, 444]
+    packed_values = b"".join(dump_uvarint(v) for v in values)
+    field_id = 1 << 3 | 2  # field number 1, wire type 2
+    field_len = len(packed_values)
+    message_bytes = dump_uvarint(field_id) + dump_uvarint(field_len) + packed_values
+
+    buf = BytesIO(message_bytes)
+    msg = protobuf.load_message(buf, RepeatedFields)
+    assert msg
+    assert msg.uintlist == values
+    assert not msg.enumlist
+    assert not msg.strlist
+
+
+def test_packed_enum():
+    values = [0, 0, 0, 0]
+    packed_values = b"".join(dump_uvarint(v) for v in values)
+    field_id = 2 << 3 | 2  # field number 2, wire type 2
+    field_len = len(packed_values)
+    message_bytes = dump_uvarint(field_id) + dump_uvarint(field_len) + packed_values
+
+    buf = BytesIO(message_bytes)
+    msg = protobuf.load_message(buf, RepeatedFields)
+    assert msg
+    assert msg.enumlist == values
+    assert not msg.uintlist
+    assert not msg.strlist