arq session type: from string to enum

modem/arq_data_type_handler.py

@@ -4,53 +4,66 @@ import structlog
 import lzma
 import gzip
 from message_p2p import message_received, message_failed
+from enum import Enum
+
+class ARQ_SESSION_TYPES(Enum):
+    raw = 0
+    raw_lzma = 10
+    raw_gzip = 11
+    p2pmsg_lzma = 20
 
 class ARQDataTypeHandler:
     def __init__(self, event_manager):
         self.logger = structlog.get_logger(type(self).__name__)
         self.event_manager = event_manager
         self.handlers = {
-            "raw": {
+            ARQ_SESSION_TYPES.raw: {
                 'prepare': self.prepare_raw,
                 'handle': self.handle_raw,
                 'failed': self.failed_raw,
             },
-            "raw_lzma": {
+            ARQ_SESSION_TYPES.raw_lzma: {
                 'prepare': self.prepare_raw_lzma,
                 'handle': self.handle_raw_lzma,
                 'failed': self.failed_raw_lzma,
             },
-            "raw_gzip": {
+            ARQ_SESSION_TYPES.raw_gzip: {
                 'prepare': self.prepare_raw_gzip,
                 'handle': self.handle_raw_gzip,
                 'failed': self.failed_raw_gzip,
             },
-            "p2pmsg_lzma": {
+            ARQ_SESSION_TYPES.p2pmsg_lzma: {
                 'prepare': self.prepare_p2pmsg_lzma,
                 'handle': self.handle_p2pmsg_lzma,
                 'failed' : self.failed_p2pmsg_lzma,
             },
         }
 
+    @staticmethod
+    def get_session_type_from_value(value):
+        for session_type in ARQ_SESSION_TYPES:
+            if session_type.value == value:
+                return session_type
+        return None
+
     def dispatch(self, type_byte: int, data: bytearray):
-        endpoint_name = list(self.handlers.keys())[type_byte]
-        if endpoint_name in self.handlers and 'handle' in self.handlers[endpoint_name]:
-            return self.handlers[endpoint_name]['handle'](data)
+        session_type = self.get_session_type_from_value(type_byte)
+        if session_type and session_type in self.handlers and 'handle' in self.handlers[session_type]:
+            return self.handlers[session_type]['handle'](data)
         else:
-            self.log(f"Unknown handling endpoint: {endpoint_name}", isWarning=True)
+            self.log(f"Unknown handling endpoint for type: {type_byte}", isWarning=True)
 
-    def failed(self, type_byte: int, data: bytearray):
-        endpoint_name = list(self.handlers.keys())[type_byte]
-        if endpoint_name in self.handlers and 'failed' in self.handlers[endpoint_name]:
-            return self.handlers[endpoint_name]['failed'](data)
+    def failed(self, session_type: ARQ_SESSION_TYPES, data: bytearray):
+        if session_type in self.handlers and 'failed' in self.handlers[session_type]:
+            return self.handlers[session_type]['failed'](data)
         else:
-            self.log(f"Unknown handling endpoint: {endpoint_name}", isWarning=True)
+            self.log(f"Unknown handling endpoint: {session_type.name}", isWarning=True)
 
-    def prepare(self, data: bytearray, endpoint_name="raw" ):
-        if endpoint_name in self.handlers and 'prepare' in self.handlers[endpoint_name]:
-            return self.handlers[endpoint_name]['prepare'](data), list(self.handlers.keys()).index(endpoint_name)
+    def prepare(self, data: bytearray, session_type=ARQ_SESSION_TYPES.raw):
+        if session_type in self.handlers and 'prepare' in self.handlers[session_type]:
+            return self.handlers[session_type]['prepare'](data), session_type.value
         else:
-            self.log(f"Unknown preparation endpoint: {endpoint_name}", isWarning=True)
+            self.log(f"Unknown preparation endpoint: {session_type.name}", isWarning=True)
 
     def log(self, message, isWarning=False):
         msg = f"[{type(self).__name__}]: {message}"

modem/command_arq_raw.py

@@ -4,6 +4,7 @@ import api_validations
 import base64
 from queue import Queue
 from arq_session_iss import ARQSessionISS
+from arq_data_type_handler import ARQ_SESSION_TYPES
 
 
 class ARQRawCommand(TxCommand):
@@ -14,9 +15,9 @@ class ARQRawCommand(TxCommand):
             self.dxcall = f"{self.dxcall}-0"
 
         try:
-            self.type = apiParams['type']
+            self.type = ARQ_SESSION_TYPES[apiParams['type']]
         except KeyError:
-            self.type = "raw"
+            self.type = ARQ_SESSION_TYPES.raw
 
         self.data = base64.b64decode(apiParams['data'])
 

modem/command_message_send.py

@@ -4,7 +4,7 @@ import base64
 from queue import Queue
 from arq_session_iss import ARQSessionISS
 from message_p2p import MessageP2P
-from arq_data_type_handler import ARQDataTypeHandler
+from arq_data_type_handler import ARQ_SESSION_TYPES
 from message_system_db_manager import DatabaseManager
 
 class SendMessageCommand(TxCommand):
@@ -22,7 +22,7 @@ class SendMessageCommand(TxCommand):
 
         payload = self.message.to_payload().encode('utf-8')
         json_bytearray = bytearray(payload)
-        data, data_type = self.arq_data_type_handler.prepare(json_bytearray, 'p2pmsg_lzma')
+        data, data_type = self.arq_data_type_handler.prepare(json_bytearray, ARQ_SESSION_TYPES.p2pmsg_lzma)
 
 
 

tests/test_data_type_handler.py

@@ -3,7 +3,7 @@ sys.path.append('modem')
 
 import unittest
 import queue
-from arq_data_type_handler import ARQDataTypeHandler
+from arq_data_type_handler import ARQDataTypeHandler, ARQ_SESSION_TYPES
 from event_manager import EventManager
 
 
@@ -19,21 +19,21 @@ class TestDispatcher(unittest.TestCase):
     def testDataTypeHevent_managerandlerRaw(self):
         # Example usage
         example_data = b"Hello FreeDATA!"
-        formatted_data, type_byte = self.arq_data_type_handler.prepare(example_data, "raw")
+        formatted_data, type_byte = self.arq_data_type_handler.prepare(example_data, ARQ_SESSION_TYPES.raw)
         dispatched_data = self.arq_data_type_handler.dispatch(type_byte, formatted_data)
         self.assertEqual(example_data, dispatched_data)
 
     def testDataTypeHandlerLZMA(self):
         # Example usage
         example_data = b"Hello FreeDATA!"
-        formatted_data, type_byte = self.arq_data_type_handler.prepare(example_data, "raw_lzma")
+        formatted_data, type_byte = self.arq_data_type_handler.prepare(example_data, ARQ_SESSION_TYPES.raw_lzma)
         dispatched_data = self.arq_data_type_handler.dispatch(type_byte, formatted_data)
         self.assertEqual(example_data, dispatched_data)
 
     def testDataTypeHandlerGZIP(self):
         # Example usage
         example_data = b"Hello FreeDATA!"
-        formatted_data, type_byte = self.arq_data_type_handler.prepare(example_data, "raw_gzip")
+        formatted_data, type_byte = self.arq_data_type_handler.prepare(example_data, ARQ_SESSION_TYPES.raw_gzip)
         dispatched_data = self.arq_data_type_handler.dispatch(type_byte, formatted_data)
         self.assertEqual(example_data, dispatched_data)