integrate_database.py migration with test robot sync

fcladera · fcladera · commit cd4350eeafa7 · 2025-08-13T00:14:38.000-04:00
diff --git a/mocha_core/CMakeLists.txt b/mocha_core/CMakeLists.txt
@@ -63,6 +63,8 @@ if(BUILD_TESTING)
   ament_add_pytest_test(test_zmq_comm_node test/test_zmq_comm_node.py)
   ament_add_pytest_test(test_synchronize_channel test/test_synchronize_channel.py)
   ament_add_pytest_test(test_database_server test/test_database_server.py)
+  ament_add_pytest_test(test_delay_2robots test/test_delay_2robots.py)
+  ament_add_pytest_test(test_multiple_robots_sync test/test_multiple_robots_sync.py)
 endif()
 
 ament_package()
diff --git a/mocha_core/mocha_core/integrate_database.py b/mocha_core/mocha_core/integrate_database.py
@@ -7,9 +7,9 @@
 import traceback
 from functools import partial
 
-import roslaunch
-import rospkg
-import rospy
+import rclpy
+import rclpy.logging
+import rclpy.time
 import yaml
 import std_msgs.msg
 import subprocess
@@ -31,35 +31,53 @@ def ping(host):
 
 
 class IntegrateDatabase:
-    def __init__(self):
-        rospy.init_node("integrate_database", anonymous=False)
-
-        self.this_robot = rospy.get_param("~robot_name")
-        self.rssi_threshold = rospy.get_param("~rssi_threshold", 20)
+    def __init__(self, node=None):
+        # Initialize ROS2 if not already done
+        if not rclpy.ok():
+            rclpy.init()
+        
+        # Create or use provided node
+        if node is None:
+            self.node = rclpy.create_node("integrate_database")
+        else:
+            self.node = node
+        
+        self.logger = self.node.get_logger()
+
+        # In ROS2, parameters need to be declared before use
+        self.node.declare_parameter("robot_name", "")
+        self.node.declare_parameter("rssi_threshold", 20)
+        self.node.declare_parameter("client_timeout", 6.0)
+        self.node.declare_parameter("robot_configs", "")
+        self.node.declare_parameter("radio_configs", "")
+        self.node.declare_parameter("topic_configs", "")
+
+        self.this_robot = self.node.get_parameter("robot_name").get_parameter_value().string_value
+        self.rssi_threshold = self.node.get_parameter("rssi_threshold").get_parameter_value().integer_value
         self.all_channels = []
-        rospy.loginfo(f"{self.this_robot} - Integrate - " +
-                      f"RSSI threshold: {self.rssi_threshold}")
-        self.client_timeout = rospy.get_param("~client_timeout", 6.)
-        rospy.loginfo(f"{self.this_robot} - Integrate - " +
-                      f"Client timeout: {self.client_timeout}")
+        self.logger.info(f"{self.this_robot} - Integrate - " +
+                        f"RSSI threshold: {self.rssi_threshold}")
+        self.client_timeout = self.node.get_parameter("client_timeout").get_parameter_value().double_value
+        self.logger.info(f"{self.this_robot} - Integrate - " +
+                        f"Client timeout: {self.client_timeout}")
 
         # Load and check robot configs
-        self.robot_configs_file = rospy.get_param("~robot_configs")
+        self.robot_configs_file = self.node.get_parameter("robot_configs").get_parameter_value().string_value
         with open(self.robot_configs_file, "r") as f:
             self.robot_configs = yaml.load(f, Loader=yaml.FullLoader)
         if self.this_robot not in self.robot_configs.keys():
             self.shutdown("Robot not in config file")
 
         # Load and check radio configs
-        self.radio_configs_file = rospy.get_param("~radio_configs")
+        self.radio_configs_file = self.node.get_parameter("radio_configs").get_parameter_value().string_value
         with open(self.radio_configs_file, "r") as f:
             self.radio_configs = yaml.load(f, Loader=yaml.FullLoader)
         self.radio = self.robot_configs[self.this_robot]["using-radio"]
         if self.radio not in self.radio_configs.keys():
             self.shutdown("Radio not in config file")
 
         # Load and check topic configs
-        self.topic_configs_file = rospy.get_param("~topic_configs")
+        self.topic_configs_file = self.node.get_parameter("topic_configs").get_parameter_value().string_value
         with open(self.topic_configs_file, "r") as f:
             self.topic_configs = yaml.load(f, Loader=yaml.FullLoader)
         self.node_type = self.robot_configs[self.this_robot]["node-type"]
@@ -69,36 +87,36 @@ def __init__(self):
         # Check that we can ping the radios
         ip = self.robot_configs[self.this_robot]["IP-address"]
         if not ping(ip):
-            rospy.logerr(f"{self.this_robot} - Integrate - " +
-                         f"Cannot ping self {ip}. Is the radio on?")
-            rospy.signal_shutdown("Cannot ping self")
-            rospy.spin()
+            self.logger.error(f"{self.this_robot} - Integrate - " +
+                             f"Cannot ping self {ip}. Is the radio on?")
+            self.shutdown("Cannot ping self")
+            return
 
-        # Create a database server object
+        # Create a database server object with ROS2 node
         self.DBServer = ds.DatabaseServer(self.robot_configs,
-                                          self.topic_configs)
+                                          self.topic_configs, node=self.node)
 
         self.num_robot_in_comm = 0
 
         # Handle possible interruptions
         self.interrupted = False
 
         def signal_handler(sig, frame):
-            rospy.logwarn(f"{self.this_robot} - Integrate - " +
-                          f"Got signal. Killing comm nodes.")
+            self.logger.warning(f"{self.this_robot} - Integrate - " +
+                               f"Got signal. Killing comm nodes.")
             self.interrupted = True
-            rospy.signal_shutdown("Killed by user")
+            self.shutdown("Killed by user")
         signal.signal(signal.SIGINT, signal_handler)
 
-        rospy.loginfo(f"{self.this_robot} - Integrate - " +
-                      "Created all communication channels!")
+        self.logger.info(f"{self.this_robot} - Integrate - " +
+                        "Created all communication channels!")
 
         # Start comm channels with other robots
         self.other_robots = [i for i in list(self.robot_configs.keys()) if i !=
                              self.this_robot]
         for other_robot in self.other_robots:
             if other_robot not in self.robot_configs[self.this_robot]["clients"]:
-                rospy.logwarn(
+                self.logger.warning(
                     f"{self.this_robot} - Integrate - "+
                     f"Skipping channel {self.this_robot}->{other_robot} " +
                     "as it is not in the graph of this robot"
@@ -115,40 +133,48 @@ def signal_handler(sig, frame):
             # Attach a radio trigger to each channel. This will be triggered
             # when the RSSI is high enough. You can use another approach here
             # such as using a timer to periodically trigger the sync
-            rospy.Subscriber('ddb/rajant/rssi/' + other_robot,
-                             std_msgs.msg.Int32,
-                             self.rssi_cb,
-                             channel)
+            def make_callback(ch):
+                return lambda msg: self.rssi_cb(msg, ch)
+            
+            self.node.create_subscription(
+                std_msgs.msg.Int32,
+                'ddb/rajant/rssi/' + other_robot,
+                make_callback(channel),
+                10
+            )
 
         # Wait for all the robots to start
         for _ in range(100):
-            if self.interrupted or rospy.is_shutdown():
+            if self.interrupted or not rclpy.ok():
                 self.shutdown("Killed while waiting for other robots")
                 return
-            rospy.sleep(.1)
+            time.sleep(0.1)
 
         # Spin!
-        rospy.spin()
+        try:
+            rclpy.spin(self.node)
+        except KeyboardInterrupt:
+            pass
 
     def shutdown(self, reason):
         assert isinstance(reason, str)
-        rospy.logerr(f"{self.this_robot} - Integrate - " + reason)
+        self.logger.error(f"{self.this_robot} - Integrate - " + reason)
         for channel in self.all_channels:
             channel.comm_node.terminate()
             self.all_channels.remove(channel)
-        rospy.logwarn(f"{self.this_robot} - Integrate - " + "Killed Channels")
+        self.logger.warning(f"{self.this_robot} - Integrate - " + "Killed Channels")
         self.DBServer.shutdown()
-        rospy.logwarn(f"{self.this_robot} - Integrate - " + "Killed DB")
-        rospy.signal_shutdown(reason)
-        rospy.logwarn(f"{self.this_robot} - Integrate - " + "Shutting down")
-        rospy.spin()
+        self.logger.warning(f"{self.this_robot} - Integrate - " + "Killed DB")
+        if rclpy.ok():
+            rclpy.shutdown()
+        self.logger.warning(f"{self.this_robot} - Integrate - " + "Shutting down")
 
     def rssi_cb(self, data, comm_node):
         rssi = data.data
         if rssi > self.rssi_threshold:
             self.num_robot_in_comm += 1
             try:
-                rospy.loginfo(f"{self.this_robot} <- {comm_node.target_robot}: Triggering comms")
+                self.logger.info(f"{self.this_robot} <- {comm_node.target_robot}: Triggering comms")
                 comm_node.trigger_sync()
             except:
                 traceback.print_exception(*sys.exc_info())
diff --git a/mocha_core/test/sample_db.py b/mocha_core/test/sample_db.py
@@ -2,7 +2,13 @@
 import copy
 from pprint import pprint
 import sys
-sys.path.append('..')
+import os
+
+# Add the mocha_core module path for imports
+current_dir = os.path.dirname(os.path.abspath(__file__))
+mocha_core_path = os.path.join(current_dir, "..", "mocha_core")
+sys.path.append(mocha_core_path)
+
 import hash_comm
 import database as db
 import database_utils as du
diff --git a/mocha_core/test/test_delay_2robots.py b/mocha_core/test/test_delay_2robots.py
@@ -6,34 +6,66 @@
 import time
 from pprint import pprint
 import multiprocessing
-sys.path.append('..')
-import get_sample_db as sdb
+import os
+import rclpy
+import rclpy.time
+import yaml
+
+# Add the mocha_core module path for imports
+current_dir = os.path.dirname(os.path.abspath(__file__))
+mocha_core_path = os.path.join(current_dir, "..", "mocha_core")
+sys.path.append(mocha_core_path)
+
+import sample_db as sdb
 import synchronize_channel as sync
-import synchronize_utils as su
-import database_server_utils as du
+import database as db
+import database_utils as du
 from colorama import Fore, Style
 
-CONFIG_FILE = "testConfigs/robotConfigs.yml"
+# Load robot configs for ROS2
+current_dir = os.path.dirname(os.path.abspath(__file__))
+robot_configs_path = os.path.join(
+    current_dir, "..", "config", "testConfigs", "robot_configs.yaml"
+)
+with open(robot_configs_path, "r") as f:
+    robot_configs = yaml.load(f, Loader=yaml.FullLoader)
+
+client_timeout = 6.0
 
 class test(unittest.TestCase):
     def setUp(self):
         test_name = self._testMethodName
         print("\n", Fore.RED, 20*"=", test_name, 20*"=", Style.RESET_ALL)
+        
+        # Initialize ROS2 if not already done
+        if not rclpy.ok():
+            rclpy.init()
+        
+        # Create a ROS2 node for testing
+        self.node = rclpy.create_node('test_delay_2robots_node')
+    
+    def tearDown(self):
+        # Shutdown ROS2 node
+        if hasattr(self, 'node'):
+            self.node.destroy_node()
+        
+        time.sleep(1)
+        super().tearDown()
 
     def test_delay_run(self):
         self.maxDiff = None
         db_gs = sdb.get_sample_dbl()
         db_r1 = sdb.get_sample_dbl()
         db_r2 = sdb.get_sample_dbl()
 
-        node_gs_r1 = sync.Channel(db_gs, 'groundstation', 'charon', CONFIG_FILE)
-        node_gs_r2 = sync.Channel(db_gs, 'groundstation', 'styx', CONFIG_FILE)
+        node_gs_r1 = sync.Channel(db_gs, 'basestation', 'charon', robot_configs, client_timeout, self.node)
+        node_gs_r2 = sync.Channel(db_gs, 'basestation', 'styx', robot_configs, client_timeout, self.node)
 
-        node_r1_gs = sync.Channel(db_r1, 'charon', 'groundstation', CONFIG_FILE)
-        node_r2_gs = sync.Channel(db_r2, 'styx', 'groundstation', CONFIG_FILE)
+        node_r1_gs = sync.Channel(db_r1, 'charon', 'basestation', robot_configs, client_timeout, self.node)
+        node_r2_gs = sync.Channel(db_r2, 'styx', 'basestation', robot_configs, client_timeout, self.node)
 
-        node_r1_r2 = sync.Channel(db_r1, 'charon', 'styx', CONFIG_FILE)
-        node_r2_r1 = sync.Channel(db_r2, 'styx', 'charon', CONFIG_FILE)
+        node_r1_r2 = sync.Channel(db_r1, 'charon', 'styx', robot_configs, client_timeout, self.node)
+        node_r2_r1 = sync.Channel(db_r2, 'styx', 'charon', robot_configs, client_timeout, self.node)
 
         pprint(db_gs.db)
         pprint(db_r1.db)
@@ -42,17 +74,17 @@ def test_delay_run(self):
         node_gs_r1.run()
         node_gs_r2.run()
         
-        dbm = su.DBMessage(1, 'Node_1', 1, 2, time.time(), bytes('r1_data', 'utf-8'), False)
-        du.add_modify_data_dbl(db_r1, dbm)
+        dbm = db.DBMessage(1, 1, 2, 1, rclpy.time.Time(seconds=int(time.time())), bytes('r1_data', 'utf-8'))
+        db_r1.add_modify_data(dbm)
 
         node_gs_r1.trigger_sync()
         time.sleep(4)
 
         node_r1_gs.run()
         node_r2_gs.run()
 
-        dbm = su.DBMessage(2, 'Node_2', 2, 1, time.time(), bytes('r2_data', 'utf-8'), False)
-        du.add_modify_data_dbl(db_r2, dbm)
+        dbm = db.DBMessage(2, 2, 1, 2, rclpy.time.Time(seconds=int(time.time())), bytes('r2_data', 'utf-8'))
+        db_r2.add_modify_data(dbm)
         
         node_r1_r2.run()
 
@@ -85,14 +117,14 @@ def test_multiple_trigger_sync(self):
         db_r1 = sdb.get_sample_dbl()
         db_r2 = sdb.get_sample_dbl()
 
-        node_gs_r1 = sync.Channel(db_gs, 'groundstation', 'charon', CONFIG_FILE)
-        node_gs_r2 = sync.Channel(db_gs, 'groundstation', 'styx', CONFIG_FILE)
+        node_gs_r1 = sync.Channel(db_gs, 'basestation', 'charon', robot_configs, client_timeout, self.node)
+        node_gs_r2 = sync.Channel(db_gs, 'basestation', 'styx', robot_configs, client_timeout, self.node)
 
-        node_r1_gs = sync.Channel(db_r1, 'charon', 'groundstation', CONFIG_FILE)
-        node_r2_gs = sync.Channel(db_r2, 'styx', 'groundstation', CONFIG_FILE)
+        node_r1_gs = sync.Channel(db_r1, 'charon', 'basestation', robot_configs, client_timeout, self.node)
+        node_r2_gs = sync.Channel(db_r2, 'styx', 'basestation', robot_configs, client_timeout, self.node)
 
-        node_r1_r2 = sync.Channel(db_r1, 'charon', 'styx', CONFIG_FILE)
-        node_r2_r1 = sync.Channel(db_r2, 'styx', 'charon', CONFIG_FILE)
+        node_r1_r2 = sync.Channel(db_r1, 'charon', 'styx', robot_configs, client_timeout, self.node)
+        node_r2_r1 = sync.Channel(db_r2, 'styx', 'charon', robot_configs, client_timeout, self.node)
 
 #        pprint(db_gs)
 #        pprint(db_r1)
@@ -101,8 +133,8 @@ def test_multiple_trigger_sync(self):
 #        node_r1_r2.run()
         node_r2_r1.run()
 
-        dbm = su.DBMessage(1, 'Node_1', 1, 2, time.time(), bytes('R1_data', 'utf-8'), False)
-        du.add_modify_data_dbl(db_r1, dbm)
+        dbm = db.DBMessage(1, 1, 2, 1, rclpy.time.Time(seconds=int(time.time())), bytes('R1_data', 'utf-8'))
+        db_r1.add_modify_data(dbm)
 
         node_r2_r1.trigger_sync()
         time.sleep(5)
diff --git a/mocha_core/test/test_multiple_robots_sync.py b/mocha_core/test/test_multiple_robots_sync.py
