diff --git a/src/imetro_behavior/imetro_behavior/geometry_behaviors.py b/src/imetro_behavior/imetro_behavior/geometry_behaviors.py index 49b08a6..cc8f29f 100644 --- a/src/imetro_behavior/imetro_behavior/geometry_behaviors.py +++ b/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) diff --git a/src/imetro_behavior/tests/test_geometry_behaviors.py b/src/imetro_behavior/tests/test_geometry_behaviors.py index a34c49b..55030d1 100644 --- a/src/imetro_behavior/tests/test_geometry_behaviors.py +++ b/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)