Improved zmq_comm_node and test to avoid creating node

Author: fcladera

mocha_core/mocha_core/zmq_comm_node.py

@@ -24,7 +24,7 @@ class SyncStatus(enum.Enum):
 
 class Comm_node:
     def __init__(self, this_node, client_node, robot_configs,
-                 client_callback, server_callback, client_timeout, node=None):
+                 client_callback, server_callback, client_timeout, ros_node):
         # Check input arguments
         assert isinstance(this_node, str)
         assert isinstance(client_node, str)
@@ -33,13 +33,11 @@ def __init__(self, this_node, client_node, robot_configs,
         assert callable(server_callback)
         assert isinstance(client_timeout, (int, float))
 
-        # Store or create the ROS2 node
-        self.node = node
-        if self.node is None:
-            # Create a minimal logger if no node provided
-            self.logger = rclpy.logging.get_logger('zmq_comm_node')
-        else:
-            self.logger = self.node.get_logger()
+        # Confirm that we are providing a ROS node
+        assert ros_node is not None
+        self.ros_node = ros_node
+        self.logger = self.ros_node.get_logger()
+        self.ros_node_name = self.ros_node.get_fully_qualified_name()
 
         # Check that this_node and client_node exist in the config file
         if this_node not in robot_configs:
@@ -50,7 +48,6 @@ def __init__(self, this_node, client_node, robot_configs,
         if client_node not in robot_configs[self.this_node]["clients"]:
             self.logger.error(f"{this_node} - Node: client_node not in config file")
             raise ValueError("client_node not in config file")
-
         self.client_node = client_node
 
         self.robot_configs = robot_configs
@@ -72,14 +69,11 @@ def __init__(self, this_node, client_node, robot_configs,
         self.client_timeout = client_timeout
 
         # Create a publisher for the client bandwidth
-        if self.node is not None:
-            self.pub_client_stats = self.node.create_publisher(
-                mocha_core.msg.ClientStats,
-                f"ddb/client_stats/{self.client_node}",
-                10
-            )
-        else:
-            self.pub_client_stats = None
+        self.pub_client_stats = self.ros_node.create_publisher(
+            mocha_core.msg.ClientStats,
+            f"{self.ros_node_name}/client_stats/{self.client_node}",
+            10
+        )
         self.pub_client_count = 0
 
         self.syncStatus = SyncStatus.IDLE
@@ -134,11 +128,8 @@ def connect_send_message(self, msg):
         self.logger.debug(f"{self.this_node} - Node - SENDMSG: " +
                        f"Sending ({full_msg})")
         client.send(full_msg)
-        if self.node is not None:
-            start_ts = self.node.get_clock().now()
-        else:
-            import time
-            start_ts = time.time()
+        start_ts = self.ros_node.get_clock().now()
+
         # Wait at most self.client_timeout * 1000 ms
         socks = dict(poll.poll(self.client_timeout*1000))
         if socks.get(client) == zmq.POLLIN:
@@ -157,21 +148,15 @@ def connect_send_message(self, msg):
                         f"{self.this_node} - Node - SENDMSG: Server replied " +
                         f"({len(reply)} bytes)"
                     )
-                    if self.node is not None:
-                        stop_ts = self.node.get_clock().now()
-                        time_d = (stop_ts - start_ts).nanoseconds
-                        time_s = float(time_d / 1e9)
-                    else:
-                        import time
-                        stop_ts = time.time()
-                        time_s = stop_ts - start_ts
-                    
+                    stop_ts = self.ros_node.get_clock().now()
+                    time_d = (stop_ts - start_ts).nanoseconds
+                    time_s = float(time_d / 1e9)
+
                     bw = len(reply)/time_s/1024/1024
-                    
+
                     if self.pub_client_stats is not None:
                         stats = mocha_core.msg.ClientStats()
-                        if self.node is not None:
-                            stats.header.stamp = self.node.get_clock().now().to_msg()
+                        stats.header.stamp = self.ros_node.get_clock().now().to_msg()
                         stats.header.frame_id = self.this_node
                         # Note: ROS2 doesn't have seq field in header
                         stats.msg = msg[:5].decode("utf-8")

mocha_core/test/test_zmq_comm_node.py

@@ -19,7 +19,16 @@
 import pdb
 import rclpy
 import rclpy.logging
+from rclpy.logging import LoggingSeverity
+from rclpy.executors import SingleThreadedExecutor
+from rclpy.node import Node
+import threading
 from colorama import Fore, Style
+import mocha_core.zmq_comm_node as zmq_comm_node
+
+class Comm_node_test(Node):
+    def __init__(self):
+        super().__init__("comm_node_test")
 
 
 class Test(unittest.TestCase):
@@ -28,44 +37,71 @@ def setUpClass(cls):
         # Load configurations at the class level
         current_dir = os.path.dirname(os.path.abspath(__file__))
         ddb_path = os.path.join(current_dir, "..")
-        
+
         # Load robot configs
         robot_configs_path = os.path.join(ddb_path, "config/testConfigs/robot_configs.yaml")
         with open(robot_configs_path, "r") as f:
             cls.robot_configs = yaml.load(f, Loader=yaml.FullLoader)
 
+
     def setUp(self):
         test_name = self._testMethodName
         print("\n", Fore.RED, 20*"=", test_name, 20*"=", Style.RESET_ALL)
+        # Create a mock node to run the tests
+        rclpy.init()
+        self.comm_node_test = Comm_node_test()
+        self.executor = SingleThreadedExecutor()
+        self.executor.add_node(self.comm_node_test)
+        executor_thread = threading.Thread(target=self.executor.spin, daemon=True)
+        executor_thread.start()
         super().setUp()
 
     def tearDown(self):
         time.sleep(1)
+        self.executor.shutdown()
+        self.comm_node_test.destroy_node()
+        rclpy.shutdown()
         super().tearDown()
 
     def test_simple_connection(self):
-        self.answer = None
-        logger = rclpy.logging.get_logger('test_zmq_comm_node')
-
-        def cb_groundstation(value):
-            logger.debug("cb_groundstation")
-            self.answer = value
-
-        def cb_charon(value):
+        self.answer_cb_gs_client = None
+        self.answer_cb_ch_client = None
+        logger = self.comm_node_test.get_logger()
+        logger.set_level(LoggingSeverity.DEBUG)
+
+        def cb_groundstation_client(value):
+            logger.debug(f"cb groundstation client")
+            self.answer_cb_gs_client = value
+
+        def cb_groundstation_server(value):
+            # This function is called upon reception of a message by the base
+            # station
+            logger.debug("cb groundstation server")
+            to_append = b"GSSERVER"
+            return value + to_append
+
+        def cb_charon_client(value):
+            logger.debug(f"cb charon client")
+            self.answer_cb_ch_client = value
+
+        def cb_charon_server(value):
             # This function is called upon reception of a message by charon.
             # The return value is transmitted as answer to the original
-            # message.
-            logger.debug(f"cb_charon: {value}")
-            return value
+            # message + b"CHARON"
+            logger.debug(f"cb charon server")
+            to_append = b"CHARONSERVER"
+            return value + to_append
 
         # Create the two robots
         node_groundstation = zmq_comm_node.Comm_node(
             "basestation", "charon", self.robot_configs,
-            cb_groundstation, cb_groundstation, 2
+            cb_groundstation_client, cb_groundstation_server,
+            2, self.comm_node_test
         )
         node_charon = zmq_comm_node.Comm_node(
             "charon", "basestation", self.robot_configs,
-            cb_charon, cb_charon, 2
+            cb_charon_client, cb_charon_server,
+            2, self.comm_node_test
         )
 
         # Generate random message
@@ -76,21 +112,15 @@ def cb_charon(value):
 
         # Send message from node_groundstation to robot 2
         node_groundstation.connect_send_message(random_msg)
-
-        # node_charon.connect_send_message(random_msg)
+        node_charon.connect_send_message(random_msg)
 
         # Terminate robots and test assertion
         node_groundstation.terminate()
         node_charon.terminate()
-        self.assertEqual(random_msg, self.answer, "Sent %s" % random_msg)
-
-
-# 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 mocha_core.zmq_comm_node as zmq_comm_node
+        self.assertEqual(random_msg + b"CHARONSERVER", self.answer_cb_gs_client,
+                         "Sent %s" % random_msg)
+        self.assertEqual(random_msg + b"GSSERVER", self.answer_cb_ch_client,
+                         "Sent %s" % random_msg)
 
 if __name__ == "__main__":
     # Run test cases!