From 9b6647c29e5740d8ace3e9b9b0c4a5d443bc30c4 Mon Sep 17 00:00:00 2001
From: Pedro <ptsmonteiro@gmail.com>
Date: Sun, 17 Dec 2023 01:43:23 +0100
Subject: [PATCH] Move decode mode setting to demodulator. Fix ARQ tests.

---
 modem/arq_session_irs.py  | 26 +++++---------------------
 modem/demodulator.py      | 12 ++++++++++++
 tests/test_arq_session.py |  2 ++
 3 files changed, 19 insertions(+), 21 deletions(-)

diff --git a/modem/arq_session_irs.py b/modem/arq_session_irs.py
index 28235d0f..8a170e3e 100644
--- a/modem/arq_session_irs.py
+++ b/modem/arq_session_irs.py
@@ -1,6 +1,4 @@
 import threading
-import data_frame_factory
-import queue
 import arq_session
 import helpers
 from modem_frametypes import FRAME_TYPE
@@ -56,6 +54,9 @@ class ARQSessionIRS(arq_session.ARQSession):
 
         self.transmitted_acks = 0
 
+    def set_decode_mode(self):
+        self.modem.demodulator.set_decode_mode(self.get_mode_by_speed_level(self.speed_level))
+
     def all_data_received(self):
         return self.total_length == self.received_bytes
 
@@ -100,7 +101,7 @@ class ARQSessionIRS(arq_session.ARQSession):
         self.dx_snr.append(info_frame['snr'])
 
         self.calibrate_speed_settings()
-        self.set_modem_decode_modes(self.speed_level)
+        self.set_decode_mode()
         info_ack = self.frame_factory.build_arq_session_info_ack(
             self.id, self.total_crc, self.snr[0],
             self.speed_level, self.frames_per_burst)
@@ -109,27 +110,10 @@ class ARQSessionIRS(arq_session.ARQSession):
 
     def send_burst_nack(self):
         self.calibrate_speed_settings()
-        self.set_modem_decode_modes(self.speed_level)
+        self.set_decode_mode()
         nack = self.frame_factory.build_arq_burst_ack(self.id, self.received_bytes, self.speed_level, self.frames_per_burst, self.snr[0])
         self.launch_transmit_and_wait(nack, None, mode=FREEDV_MODE.signalling)
 
-
-    def set_modem_decode_modes(self, speed_level):
-
-        for mode in self.modem.demodulator.MODE_DICT:
-            self.modem.demodulator.MODE_DICT[mode]["decode"] = False
-
-        # signalling is always true
-        self.modem.demodulator.MODE_DICT[FREEDV_MODE.signalling.value]["decode"] = True
-
-        mode = self.get_mode_by_speed_level(self.speed_level)
-        # Enable mode based on speed_level
-        self.modem.demodulator.MODE_DICT[mode.value]["decode"] = True
-        self.log(f"Modem set to speed level {speed_level}")
-
-        return
-
-
     def process_incoming_data(self, frame):
         if frame['offset'] != self.received_bytes:
             self.logger.info(f"Discarding data frame due to wrong offset", frame=self.frame_received)
diff --git a/modem/demodulator.py b/modem/demodulator.py
index 4b87d951..af27ddbe 100644
--- a/modem/demodulator.py
+++ b/modem/demodulator.py
@@ -394,3 +394,15 @@ class Demodulator():
         for mode in self.MODE_DICT:
             codec2.api.freedv_set_sync(self.MODE_DICT[mode]["instance"], 0)
 
+    def set_decode_mode(self, mode):
+
+        for m in self.MODE_DICT: self.MODE_DICT[m]["decode"] = False
+
+        # signalling is always true
+        self.MODE_DICT[codec2.FREEDV_MODE.signalling.value]["decode"] = True
+
+        # Enable mode based on speed_level
+        self.MODE_DICT[mode.value]["decode"] = True
+        self.log(f"Demodulator data mode {mode.name}")
+
+        return
diff --git a/tests/test_arq_session.py b/tests/test_arq_session.py
index 8717b1d5..1c4d1d8d 100644
--- a/tests/test_arq_session.py
+++ b/tests/test_arq_session.py
@@ -2,6 +2,7 @@ import sys
 sys.path.append('modem')
 
 import unittest
+import unittest.mock
 from config import CONFIG
 import helpers
 import queue
@@ -16,6 +17,7 @@ import structlog
 class TestModem:
     def __init__(self):
         self.data_queue_received = queue.Queue()
+        self.demodulator = unittest.mock.Mock()
 
     def transmit(self, mode, repeats: int, repeat_delay: int, frames: bytearray) -> bool:
         self.data_queue_received.put(frames)