topic_publisher.py migrated with new test cases.

Author: fcladera

mocha_core/CMakeLists.txt

@@ -71,6 +71,7 @@ if(BUILD_TESTING)
   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)
   ament_add_pytest_test(test_translator test/test_translator.py)
+  ament_add_pytest_test(test_topic_publisher test/test_topic_publisher.py)
 endif()
 
 ament_package()

mocha_core/mocha_core/topic_publisher.py

@@ -3,8 +3,10 @@
 import pdb
 import sys
 
-import rospkg
-import rospy
+import rclpy
+import rclpy.time
+import rclpy.clock
+import ament_index_python
 import yaml
 import re
 import mocha_core.msg
@@ -13,61 +15,86 @@
 
 
 class TopicPublisher():
-    def __init__(self, this_robot, target, msg_history="WHOLE_HISTORY"):
+    def __init__(self, this_robot, target, msg_history="WHOLE_HISTORY", node=None):
 
         self.this_robot = this_robot
 
+        # Store or create the ROS2 node
+        if node is None:
+            # Initialize ROS2 if not already done
+            if not rclpy.ok():
+                rclpy.init()
+            self.__node = rclpy.create_node(f'topic_publisher_{this_robot}')
+            self.__node_owned = True
+        else:
+            self.__node = node
+            self.__node_owned = False
+        
+        self.__logger = self.__node.get_logger()
 
         # Service configuration
-        self.__select_service = "integrate_database/SelectDB"
-        self.__get_header_service = "integrate_database/GetDataHeaderDB"
-
-        try:
-            rospy.wait_for_service(self.__select_service)
-            rospy.wait_for_service(self.__get_header_service)
-        except rospy.ROSInterruptException as exc:
-            rospy.logdebug("Service did not process request: " +
-                           str(exc))
-            rospy.signal_shutdown("Service did not process request")
-            rospy.spin()
-        self.__select_db = rospy.ServiceProxy(
-            self.__select_service, mocha_core.srv.SelectDB
-        )
+        self.__select_service = "select_db"
+        self.__get_header_service = "get_data_header_db"
 
-        self.__get_header_db = rospy.ServiceProxy(
-            self.__get_header_service, mocha_core.srv.GetDataHeaderDB
+        # Create service clients
+        self.__select_db = self.__node.create_client(
+            mocha_core.srv.SelectDB, self.__select_service
+        )
+        self.__get_header_db = self.__node.create_client(
+            mocha_core.srv.GetDataHeaderDB, self.__get_header_service
         )
 
+        # Wait for services to be available
+        if not self.__select_db.wait_for_service(timeout_sec=10.0):
+            self.__logger.error(f"Service {self.__select_service} not available")
+            
+        if not self.__get_header_db.wait_for_service(timeout_sec=10.0):
+            self.__logger.error(f"Service {self.__get_header_service} not available")
+
         # List of robots to pull
         self.__robot_list = []
 
         # Publisher creation
         self.publishers = {}
-        self.header_pub = rospy.Publisher("ddb/topic_publisher/headers",
-                                          mocha_core.msg.Header_pub,
-                                          queue_size=10)
+        self.header_pub = self.__node.create_publisher(
+            mocha_core.msg.HeaderPub,
+            "ddb/topic_publisher/headers",
+            10
+        )
         for t in target:
             robot, robot_id, topic, topic_id, obj = t
             if robot not in self.__robot_list:
                 self.__robot_list.append(robot_id)
-            self.publishers[(robot_id, topic_id)] = {"pub": rospy.Publisher(f"/{robot}{topic}",
-                                                         obj,
-                                                         queue_size=10),
-                                                         "ts": rospy.Time(1, 0)}
+            self.publishers[(robot_id, topic_id)] = {
+                "pub": self.__node.create_publisher(obj, f"/{robot}{topic}", 10),
+                "ts": rclpy.time.Time(seconds=1, nanoseconds=0)
+            }
 
 
     def run(self):
-        rospy.loginfo(f"{self.this_robot} - Topic Publisher - Started")
-        rate = rospy.Rate(10)
+        self.__logger.info(f"{self.this_robot} - Topic Publisher - Started")
+        rate = self.__node.create_rate(10)
         headers = set()
-        while not rospy.is_shutdown():
+        while rclpy.ok():
             for robot_id in self.__robot_list:
                 headers_to_get = []
 
+                # Create service request
+                req = mocha_core.srv.SelectDB.Request()
+                req.robot_id = robot_id
+                req.topic_id = 0  # None equivalent for uint8
+                req.ts_limit = self.__node.get_clock().now().to_msg()
+                
                 try:
-                    answ = self.__select_db(robot_id, None, None)
-                except rospy.ServiceException as exc:
-                    rospy.logdebug(f"Service did not process request {exc}")
+                    future = self.__select_db.call_async(req)
+                    rclpy.spin_until_future_complete(self.__node, future, timeout_sec=1.0)
+                    if future.done():
+                        answ = future.result()
+                    else:
+                        self.__logger.debug("Service call timeout")
+                        continue
+                except Exception as exc:
+                    self.__logger.debug(f"Service did not process request {exc}")
                     continue
                 returned_headers = du.deserialize_headers(answ.headers)
                 if len(returned_headers) == 0:
@@ -79,75 +106,107 @@ def run(self):
                         headers_to_get.append(header_)
 
                 for get_header in headers_to_get:
-                    rospy.logdebug(f"Getting headers {get_header}")
+                    self.__logger.debug(f"Getting headers {get_header}")
+                    
+                    # Create service request for getting header data
+                    req = mocha_core.srv.GetDataHeaderDB.Request()
+                    req.msg_header = get_header
+                    
                     try:
-                        answ = self.__get_header_db(get_header)
-                    except rospy.ServiceException as exc:
-                        rospy.logerr("Service did not process request: " +
-                                     str(exc))
+                        future = self.__get_header_db.call_async(req)
+                        rclpy.spin_until_future_complete(self.__node, future, timeout_sec=1.0)
+                        if future.done():
+                            answ = future.result()
+                        else:
+                            self.__logger.debug("Get header service call timeout")
+                            continue
+                    except Exception as exc:
+                        self.__logger.error("Service did not process request: " +
+                                           str(exc))
                         continue
                     headers.add(get_header)
 
                     ans_robot_id, ans_topic_id, ans_ts, ans_data = du.parse_answer(answ,
                                                                                    msg_types)
                     for t in self.publishers.keys():
                         if t == (ans_robot_id, ans_topic_id):
-                            assert isinstance(ans_data,
-                                              self.publishers[t]['pub'].data_class)
+                            # Convert ROS1 timestamp to ROS2 Time for comparison
+                            ans_ts_ros2 = rclpy.time.Time.from_msg(ans_ts) if hasattr(ans_ts, 'sec') else ans_ts
+                            current_ts = self.publishers[t]["ts"]
+                            
                             # FIXME: remove this line once we have proper time
                             # filtering implemented
-                            if ans_ts > self.publishers[t]["ts"]:
-                                self.publishers[t]["ts"] = ans_ts
+                            if ans_ts_ros2 > current_ts:
+                                self.publishers[t]["ts"] = ans_ts_ros2
                                 self.publishers[t]["pub"].publish(ans_data)
                                 self.header_pub.publish(get_header)
-                                rospy.logdebug(f"Publishing robot_id: {ans_robot_id}" +
-                                               f" - topic: {ans_topic_id}")
+                                self.__logger.debug(f"Publishing robot_id: {ans_robot_id}" +
+                                                   f" - topic: {ans_topic_id}")
                             else:
-                                rospy.logdebug(f"Skipping robot_id: {ans_robot_id}" +
-                                               f" - topic: {ans_topic_id} as there is an old ts")
+                                self.__logger.debug(f"Skipping robot_id: {ans_robot_id}" +
+                                                   f" - topic: {ans_topic_id} as there is an old ts")
                 rate.sleep()
 
+    def shutdown(self):
+        """Cleanup method for shutting down the topic publisher"""
+        if self.__node_owned and self.__node is not None:
+            self.__node.destroy_node()
+            if rclpy.ok():
+                rclpy.shutdown()
 
-if __name__ == "__main__":
-    rospy.init_node("mocha_core_publisher", anonymous=False)
 
-    # Get the mocha_core path
-    rospack = rospkg.RosPack()
-    ddb_path = rospack.get_path('mocha_core')
-    scripts_path = os.path.join(ddb_path, "scripts/core")
-    sys.path.append(scripts_path)
-    import database_utils as du
+def create_topic_publisher_node(robot_name, robot_configs_path=None, topic_configs_path=None, node=None):
+    """
+    Create and setup topic publisher for a given robot.
+    Returns the node for cleanup.
+    """
+    # Use provided node or create new one
+    if node is None:
+        # Initialize ROS2 if not already done
+        if not rclpy.ok():
+            rclpy.init()
+        node = rclpy.create_node(f"topic_publisher_{robot_name}")
+        node_owned = True
+    else:
+        node_owned = False
+        
+    logger = node.get_logger()
 
-    # Get robot from the parameters
-    this_robot = rospy.get_param("~robot_name")
+    # Get the mocha_core path using ament
+    try:
+        from ament_index_python.packages import get_package_share_directory
+        ddb_path = get_package_share_directory('mocha_core')
+    except:
+        # Fallback for development
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        ddb_path = os.path.join(current_dir, "..")
+    
+    # Add path for database_utils
+    sys.path.append(os.path.join(ddb_path, "."))
+    import database_utils as du
 
     # Get the robot_config path and generate the robot_configs object
-    robot_configs_default = os.path.join(ddb_path,
-                                         "config/testConfigs/robot_configs.yaml")
-    robot_configs_path = rospy.get_param("~robot_configs", robot_configs_default)
+    if robot_configs_path is None:
+        robot_configs_path = os.path.join(ddb_path, "config", "testConfigs", "robot_configs.yaml")
+        
     with open(robot_configs_path, 'r') as f:
         robot_configs = yaml.load(f, Loader=yaml.FullLoader)
-    if this_robot not in robot_configs.keys():
-        rospy.logerr(f"Robot {this_robot} not in robot_configs")
-        rospy.signal_shutdown("Robot not in robot_configs")
-        rospy.spin()
-
+    if robot_name not in robot_configs.keys():
+        logger.error(f"Robot {robot_name} not in robot_configs")
+        if node_owned:
+            rclpy.shutdown()
+        raise ValueError(f"Robot {robot_name} not found in config")
 
     # Get the topic_configs path and generate the topic_configs object
-    topic_configs_default = os.path.join(ddb_path,
-                                         "config/testConfigs/topic_configs.yaml")
-    topic_configs_path = rospy.get_param("~topic_configs", topic_configs_default)
+    if topic_configs_path is None:
+        topic_configs_path = os.path.join(ddb_path, "config", "testConfigs", "topic_configs.yaml")
+        
     with open(topic_configs_path, 'r') as f:
         topic_configs = yaml.load(f, Loader=yaml.FullLoader)
 
     # Get msg_types dict used to publish the topics
     msg_types = du.msg_types(robot_configs, topic_configs)
 
-    # Flip the dict so that we have name: obj
-    msg_objects = {}
-    # for msg in msg_types:
-    #     msg_objects[msg_types[msg]["name"]] = msg_types[msg]["obj"]
-
     list_of_topics = set()
 
     # Iterate for each robot in robot_configs, and generate the topics
@@ -161,5 +220,50 @@ def run(self):
             robot_tuple = (robot, robot_id, topic["msg_topic"], topic_id, obj)
             list_of_topics.add(robot_tuple)
 
-    pub = TopicPublisher(this_robot, list_of_topics)
-    pub.run()
+    # Create TopicPublisher instance
+    pub = TopicPublisher(robot_name, list_of_topics, node=node)
+    
+    return node, pub
+
+
+if __name__ == "__main__":
+    # Initialize ROS2
+    rclpy.init()
+    
+    # Create main node
+    node = rclpy.create_node("mocha_core_publisher")
+    logger = node.get_logger()
+
+    # Declare parameters with defaults
+    node.declare_parameter("robot_name", "")
+    node.declare_parameter("robot_configs", "")
+    node.declare_parameter("topic_configs", "")
+
+    # Get robot from the parameters
+    this_robot = node.get_parameter("robot_name").get_parameter_value().string_value
+    if not this_robot:
+        logger.error("robot_name parameter is required")
+        rclpy.shutdown()
+        sys.exit(1)
+
+    # Get config file paths from parameters
+    robot_configs_path = node.get_parameter("robot_configs").get_parameter_value().string_value
+    topic_configs_path = node.get_parameter("topic_configs").get_parameter_value().string_value
+    
+    # Use empty strings as None for the function
+    robot_configs_path = robot_configs_path if robot_configs_path else None
+    topic_configs_path = topic_configs_path if topic_configs_path else None
+
+    try:
+        # Create topic publisher node and publisher instance
+        _, pub = create_topic_publisher_node(this_robot, robot_configs_path, topic_configs_path, node)
+        
+        # Run the publisher
+        pub.run()
+    except KeyboardInterrupt:
+        pass
+    finally:
+        # Cleanup
+        pub.shutdown()
+        node.destroy_node()
+        rclpy.shutdown()