Parse length+message in the proto module

da/executor/executor.py

@@ -44,4 +44,3 @@ class Executor:
     async def run(self):
         await asyncio.gather(*(self.connect() for _ in range(self.col_num)))
         await self.execute()
-

da/executor/mock_network.py

@@ -28,4 +28,3 @@ def main():
 
 if __name__ == '__main__':
     main()
-

da/executor/proto.py

@@ -3,6 +3,12 @@ from itertools import count
 
 MAX_MSG_LEN_BYTES = 2
 
+async def parse_from_reader(reader):
+    length_prefix = await reader.readexactly(MAX_MSG_LEN_BYTES)
+    data_length = int.from_bytes(length_prefix, byteorder='big')
+    data = await reader.readexactly(data_length)
+    return unpack_message(data)
+
 def pack_message(message):
     # SerializeToString method returns an instance of bytes.
     data = message.SerializeToString()
@@ -56,4 +62,3 @@ def new_sample_res_not_found_error_msg(description):
     sample_res = dispersal_pb2.SampleRes(err=sample_err)
     dispersal_message = dispersal_pb2.DispersalMessage(sample_res=sample_res)
     return pack_message(dispersal_message)
-

da/executor/transport.py

@@ -11,12 +11,7 @@ class Transport:
     async def read_and_process(self):
         try:
             while True:
-                length_prefix = await self.reader.readexactly(proto.MAX_MSG_LEN_BYTES)
-                data_length = int.from_bytes(length_prefix, byteorder='big')
-                
-                data = await self.reader.readexactly(data_length)
-                message = proto.unpack_message(data)
-                
+                message = await proto.parse_from_reader(self.reader)
                 await self.handler(self.conn_id, self.writer, message)
         except asyncio.IncompleteReadError:
             print("Transport: Connection closed by the peer.")
@@ -29,5 +24,3 @@ class Transport:
     async def write(self, message):
         self.writer.write(message)
         await self.writer.drain()
-
-