#2 Add Protocol.unpack_many

pikhovkin · Oct 5, 2023 · 25e85d5 · 25e85d5
1 parent 8319aaa
commit 25e85d5
Show file tree

Hide file tree

Showing 2 changed files with 72 additions and 0 deletions.
diff --git a/omnicomm/protocol.py b/omnicomm/protocol.py
@@ -77,6 +77,23 @@ def unpack(cls, data: bytes) -> tuple[BaseCommand, bytes]:
         cmd = commands[cmd_num].unpack(original_data)
         return cmd, remain
 
+    @classmethod
+    def unpack_many(cls, data: bytes) -> tuple[[BaseCommand], bytes]:
+        cmds = []
+        cmd, remain = cls.unpack(data)
+        cmds.append(cmd)
+        data = remain
+        try:
+            while data:
+                cmd, remain = cls.unpack(data)
+                cmds.append(cmd)
+                data = remain
+        except Exception as err:
+            if not cmds:
+                raise err
+
+        return cmds, remain
+
     @classmethod
     def register_proto(cls, item: RegFwCmd, module: str) -> None:
         proto_class = import_string(module)

diff --git a/omnicomm/tests/test_protocol.py b/omnicomm/tests/test_protocol.py
@@ -125,3 +125,58 @@ def test_95(self):
         cmd, remain = protocol.ServerProtocol.unpack(data)
         self.assertTrue(cmd.id == commands.Cmd95.id)
         self.assertTrue(cmd.value == value)
+
+
+class ProtocolUnpackMany(TestCase):
+    def test_cmd1_no_remain(self):
+        value = {'reg_id': 1234567890, 'firmware': 1}
+        data = protocol.RegistrarProtocol.pack(commands.Cmd80(value))
+        cmds, remain = protocol.ServerProtocol.unpack_many(data)
+        self.assertTrue(len(cmds) == 1)
+        self.assertTrue(remain == b'')
+
+    def test_cmd1_broken_remain(self):
+        value = {'reg_id': 1234567890, 'firmware': 1}
+        data = protocol.RegistrarProtocol.pack(commands.Cmd80(value))
+        original_remain = data[:-1]
+        cmds, remain = protocol.ServerProtocol.unpack_many(data + original_remain)
+        self.assertTrue(len(cmds) == 1)
+        self.assertTrue(remain == original_remain)
+
+    def test_cmd2_no_remain(self):
+        value = {'reg_id': 1234567890, 'firmware': 1}
+        data = protocol.RegistrarProtocol.pack(commands.Cmd80(value))
+        cmds, remain = protocol.ServerProtocol.unpack_many(data + data)
+        self.assertTrue(len(cmds) == 2)
+        self.assertTrue(remain == b'')
+
+    def test_cmd2_broken_remain(self):
+        value = {'reg_id': 1234567890, 'firmware': 1}
+        data = protocol.RegistrarProtocol.pack(commands.Cmd80(value))
+        original_remain = data[:-1]
+        cmds, remain = protocol.ServerProtocol.unpack_many(data + data + original_remain)
+        self.assertTrue(len(cmds) == 2)
+        self.assertTrue(remain == original_remain)
+
+    def test_empty_data_error(self):
+        with self.assertRaises(exceptions.EmptyDataError):
+            protocol.ServerProtocol.unpack_many(b'')
+
+    def test_frame_marker_does_not_exist_error(self):
+        value = {'reg_id': 1234567890, 'firmware': 1}
+        data = protocol.RegistrarProtocol.pack(commands.Cmd80(value))
+        with self.assertRaises(exceptions.FrameMarkerDoesNotExistError):
+            protocol.ServerProtocol.unpack_many(data[1:])
+
+    def test_unpacking_crc_error(self):
+        value = {'reg_id': 1234567890, 'firmware': 1}
+        data = protocol.RegistrarProtocol.pack(commands.Cmd80(value))
+        with self.assertRaises(exceptions.UnpackingCRCError):
+            protocol.ServerProtocol.unpack_many(data[:-1])
+
+    def test_crc_does_not_match_error(self):
+        value = {'reg_id': 1234567890, 'firmware': 1}
+        data = protocol.RegistrarProtocol.pack(commands.Cmd80(value))
+        data = data[:-1] + bytes.fromhex(hex(data[-1] - 1)[2:])
+        with self.assertRaises(exceptions.CRCDoesNotMatchError):
+            protocol.ServerProtocol.unpack_many(data)