SetWorld behavior

src/imetro_behavior/imetro_behavior/geometry_behaviors.py

@@ -40,6 +40,7 @@
 from tf2_ros import TransformBroadcaster
 
 from imetro_behavior.helpers import set_ros_node
+from imetro_behavior.ros_behaviors.service_client import RosServiceClientBase
 
 
 class CreatePoseStamped(BehaviourWithPorts):
@@ -752,3 +753,38 @@ def terminate(self, new_status: Status) -> None:
         """Cleanup the publisher"""
         if self.status == Status.RUNNING and new_status == Status.INVALID and self.publisher is not None:
             self.publisher.destroy()
+
+
+try:
+    from phoebe_interfaces.srv import SetWorld
+
+    class SetWorldBehavior(RosServiceClientBase):
+        """
+        Update transform for the world_publisher node.
+        """
+
+        def __init__(self, name: str, **kwargs: Any):
+            super().__init__(name, service_type=SetWorld, **kwargs)
+
+        INPUT_PORTS = {"transform_stamped": PortInformation(data_type=TransformStamped)}
+
+        OUTPUT_PORTS = {}
+
+        def create_request(self) -> SetWorld.Request:
+            """Create a SetWorld service request."""
+            request = SetWorld.Request()
+            request.transform = self.get_input("transform_stamped").transform
+            return request
+
+        def process_response(self, response: SetWorld.Response) -> Status:
+            """Process the SetWorld service response."""
+            if response.success:
+                self.logger.info("Successfully set the world frame!")
+                return Status.SUCCESS
+            else:
+                self.logger.error("Error: failed to set the world frame.")
+                return Status.FAILURE
+
+
+except ModuleNotFoundError as e:
+    print(e)

src/imetro_behavior/tests/test_geometry_behaviors.py

@@ -620,3 +620,35 @@ def test_decompose_pose_stamped() -> None:
     assert frame_id_msg == "odom"
     assert translation_msg == [1.0, 2.0, 3.0]
     assert orientation_msg == [0.0, 0.0, 0.0, 1.0]
+
+
+try:
+    from phoebe_interfaces.srv import SetWorld
+
+    from imetro_behavior.geometry_behaviors import SetWorldBehavior
+
+    def test_set_world(ros_node: Node) -> None:
+        behavior = SetWorldBehavior(name="set_world_behavior", service_name="/set_world")
+        behavior.setup(node=ros_node)
+
+        behavior.setup_ports()
+        transform_stamped = TransformStamped()
+        transform_stamped.transform.translation.x = 0.0
+        transform_stamped.transform.translation.y = 0.1
+        transform_stamped.transform.translation.z = 0.2
+
+        transform_stamped.transform.rotation.x = 0.3
+        transform_stamped.transform.rotation.y = 0.4
+        transform_stamped.transform.rotation.z = 0.5
+        transform_stamped.transform.rotation.w = 0.6
+
+        Blackboard.set(behavior._get_blackboard_key("transform_stamped"), transform_stamped)
+
+        assert isinstance(behavior.create_request(), SetWorld.Request)
+
+        response = SetWorld.Response()
+        response.success = True
+        assert behavior.process_response(response) == Status.SUCCESS
+
+except ModuleNotFoundError as e:
+    print(e)