diff --git a/stretch4_body/core/client_server.py b/stretch4_body/core/client_server.py index 2d8535b..cba7eb3 100644 --- a/stretch4_body/core/client_server.py +++ b/stretch4_body/core/client_server.py @@ -42,6 +42,11 @@ def __init__(self): self.last_cmd_seq = None self.logger = logging.getLogger(name='stretch_body_server') + self.context=None + self.socket_cmd=None + self.socket_admin=None + self.socket_status=None + def startup(self): if not is_user_in_group('users'): self.logger.error("Cannot start Stretch Body Server: User is not a member of the 'users' group. The user should be a member of the 'users' group for locks to work properly.") @@ -88,14 +93,18 @@ def startup(self): return False def stop(self): - if hasattr(self, 'socket_status'): + time.sleep(0.1) # Required here for all freewheel, etc. commands to transmit + if self.socket_cmd is not None: + self.socket_cmd.close(linger=0) + if self.socket_status is not None: self.socket_status.close(linger=0) - if hasattr(self, 'socket_admin'): + if self.socket_admin is not None: self.socket_admin.close(linger=0) - if hasattr(self, 'socket_cmd'): - self.socket_cmd.close(linger=0) - if hasattr(self, 'context'): - self.context.term() + if self.context is not None: + try: + self.context.destroy(linger=0) + except Exception: + pass def dispatch_command_messages(self,cb_dispatch, is_routine_active): # Check if messages available @@ -178,8 +187,9 @@ def get_server_owning_user(): class StretchBodyClient: def __init__(self, name=None, ip_address=None): - self.server_connected=False + self._server_connected=False self.admin_poller=None + self.context=None self.socket_cmd=None self.socket_admin=None self.socket_status=None @@ -187,9 +197,31 @@ def __init__(self, name=None, ip_address=None): self.ip_address=ip_address self.cmd_seq = 0 + self.logger = logging.getLogger(name='stretch_body_client') + @property def connected(self): - return self.is_valid + """Warning: server_connected may become stale; it's not a computed property. Call `check_connection()` to refresh it.""" + return self._server_connected + + def check_connection(self): + """ + Warning: mutation of `self._server_connected` happens here. + + Updates `self._server_connected` with whether the server responds to a ping. + """ + if self.socket_admin is None: + self._server_connected = False + return self._server_connected + + ack = self._do_send_recv_admin_str(b"ping") + self._server_connected = ack is not None + + return self._server_connected + + def free_up_control(self): + ack = self._do_send_recv_admin_str(b"free_up_control", timeout=3.0) + self._server_connected = (ack == b"free_up_control") def client_id(self): pid = os.getpid() @@ -199,7 +231,7 @@ def client_id(self): def startup(self, *, verbose:bool = True, allow_different_user_connection:bool = False): if not is_user_in_group('users'): if verbose: - print("StretchBodyClient: Cannot connect to the server because the current user is not a member of the 'users' group. The user should be a member of the 'users' group for locks to work properly.") + self.logger.error("StretchBodyClient: Cannot connect to the server because the current user is not a member of the 'users' group. The user should be a member of the 'users' group for locks to work properly.") return False # Start admin REQ-REP connection @@ -210,14 +242,15 @@ def startup(self, *, verbose:bool = True, allow_different_user_connection:bool = self.socket_admin.connect(f"tcp://{self.ip_address}:23114") else: self.socket_admin.connect(f"ipc://{PORT_ADMIN}") + self.admin_poller = zmq.Poller() self.admin_poller.register(self.socket_admin, zmq.POLLIN) - self.is_valid=True - ack = self._do_send_recv_admin_str(b"ping") - self.server_connected = (ack is not None) - if ack is None: + + self.check_connection() + + if not self.connected: if verbose: - print(""" + self.logger.error(""" =============================================== StretchBodyClient: Not able to connect to Stretch Body Server. Check that server is running @@ -226,12 +259,12 @@ def startup(self, *, verbose:bool = True, allow_different_user_connection:bool = =============================================== """) - self.is_valid=False + self.stop() return False if not allow_different_user_connection and not StretchBodyServer.is_server_owned_by_current_user(): if verbose: - print(f""" + self.logger.error(f""" =============================================== StretchBodyClient: A server is already running, but it was started by a different user ({StretchBodyServer.get_server_owning_user()}). @@ -240,6 +273,7 @@ def startup(self, *, verbose:bool = True, allow_different_user_connection:bool = =============================================== """) + self.stop() return False # Start command PUB connection @@ -262,28 +296,36 @@ def startup(self, *, verbose:bool = True, allow_different_user_connection:bool = else: self.socket_status.connect(f"ipc://{PORT_STATUS}") - # self.status_poller = zmq.Poller() - # self.status_poller.register(self.socket_status, zmq.POLLIN) - - self.is_valid=True return True def stop(self): - if self.is_valid: - time.sleep(0.1) # Required here for all freewheel, etc. commands to transmit - self.socket_cmd.close() - self.socket_status.close() - self.socket_admin.close() - self.context.term() - self.is_valid=False + time.sleep(0.1) # Required here for all freewheel, etc. commands to transmit + if self.socket_cmd is not None: + self.socket_cmd.close(linger=0) + if self.socket_status is not None: + self.socket_status.close(linger=0) + if self.socket_admin is not None: + if self.admin_poller is not None: + try: + self.admin_poller.unregister(self.socket_admin) + except Exception: + pass + self.socket_admin.close(linger=0) + if self.context is not None: + try: + self.context.destroy(linger=0) + except Exception: + pass + self._server_connected = False + + def __del__(self): + try: + self.stop() + except Exception: + pass @require_connection - def _do_recv_status(self, timeout_ms=None): - # Check if messages available - if timeout_ms is not None: - if not self.status_poller.poll(int(timeout_ms)): - return None - + def _do_recv_status(self): try: message = self.socket_status.recv_pyobj(flags=zmq.NOBLOCK) return message @@ -301,10 +343,16 @@ def _do_send_cmd(self, cmd_dict, priority=0): self.socket_cmd.send_pyobj(message) @require_connection - def _do_send_recv_admin_str(self,send,timeout=1.0): - self.socket_admin.send(send) + def do_send_recv_admin_str(self, send, timeout=1.0): + return self._do_send_recv_admin_str(send, timeout) + + def _do_send_recv_admin_str(self, send, timeout=1.0): + try: + self.socket_admin.send(send) - # Poll to check if a status message is available within the timeout period - if self.admin_poller.poll(int(timeout * 1000)): - message = self.socket_admin.recv(flags=zmq.NOBLOCK) - return message + # Poll to check if a status message is available within the timeout period + if self.admin_poller.poll(int(timeout * 1000)): + message = self.socket_admin.recv(flags=zmq.NOBLOCK) + return message + except zmq.ZMQError as e: + return None \ No newline at end of file diff --git a/stretch4_body/core/subsystem_client.py b/stretch4_body/core/subsystem_client.py index 130931c..2a777a3 100644 --- a/stretch4_body/core/subsystem_client.py +++ b/stretch4_body/core/subsystem_client.py @@ -37,7 +37,7 @@ def __init__(self, name, client_id=None, parent=None, ip_address=None): @property def connected(self): - return self.is_valid and self.client.server_connected + return self.is_valid and self.client.connected def startup(self, *, verbose:bool = True, allow_different_user_connection:bool=False): """ @@ -64,10 +64,17 @@ def stop(self): for k in self.subsystems: self.subsystems[k].stop() if self.parent is None: - self.push_command(ignore_control_lock=True) #Push out clean shutdown commands to the subsystems + if self.connected: + self.push_command(ignore_control_lock=True) #Push out clean shutdown commands to the subsystems self.client.stop() #Close the sockets Device.stop(self) + def __del__(self): + try: + self.stop() + except Exception: + pass + # ########## Control loop #############3 @require_connection @@ -84,10 +91,6 @@ def pull_status(self, blocking=True): return self.parent.pull_status(blocking=blocking) status_server = self.client._do_recv_status() - # if blocking: - # t_sleep_ms = (1 / self.robot_params['robot']['server']['control_loop_rate_Hz']) * 1000 - # while status_server is None: - # status_server = self.client._do_recv_status(timeout_ms=t_sleep_ms/10) #poll 10x faster than control loop if blocking: while status_server is None: time.sleep(0.005) @@ -173,32 +176,29 @@ def is_moving(self): # ########## Server Admin ############# def is_server_active(self): - return self.client.server_connected + return self.client.connected def ping_server(self): - ack = self.client._do_send_recv_admin_str(b"ping") - self.client.server_connected = (ack == b"ping") - return self.client.server_connected + return self.client.check_connection() @require_connection def kill_server(self): - ack = self.client._do_send_recv_admin_str(b"kill") - self.client.server_connected = False + ack = self.client.do_send_recv_admin_str(b"kill") + self.client.stop() # @require_connection # def pause_control_loop(self): - # ack = self.client._do_send_recv_admin_str(b"pause") + # ack = self.client.do_send_recv_admin_str(b"pause") # self.client.server_connected = (ack ==b"pause") # # @require_connection # def unpause_control_loop(self): - # ack = self.client._do_send_recv_admin_str(b"unpause") + # ack = self.client.do_send_recv_admin_str(b"unpause") # self.client.server_connected = (ack ==b"unpause") @require_connection def free_up_control(self): - ack = self.client._do_send_recv_admin_str(b"free_up_control", timeout=3.0) - self.client.server_connected = (ack == b"free_up_control") + self.client.free_up_control() # ########## Utility #############3 diff --git a/stretch4_body/robot/robot_client.py b/stretch4_body/robot/robot_client.py index 7851107..a8cdf15 100644 --- a/stretch4_body/robot/robot_client.py +++ b/stretch4_body/robot/robot_client.py @@ -16,6 +16,19 @@ class RobotClient(SubsystemClient): This class provides access to the robot's subsystems (arm, lift, base, etc.) and high-level routines. It communicates with the robot server to execute commands and retrieve status. + + Usage: + ``` + robot = RobotClient() + success = robot.startup() + if not success: raise Exception("Could not start robot client") + ``` + + or + ``` + with RobotClient() as robot: + if robot is None: raise Exception("Could not start robot client") + ``` """ def __init__(self, client_id=None, ip_address=None): """ diff --git a/stretch4_body/robot/robot_server.py b/stretch4_body/robot/robot_server.py index 43369cf..5cb5162 100644 --- a/stretch4_body/robot/robot_server.py +++ b/stretch4_body/robot/robot_server.py @@ -323,7 +323,7 @@ def run_controller(self): except KeyboardInterrupt: self.logger.info("Received request to stop stretch body server via keyboard interrupt.") except Exception as e: - self.logger.error(f"Error in stretch body server: {e}") + self.logger.error(f"Error in Stretch Body Server's Control Loop. Control Loop has ended. Error message: {e}.") def cb_routine_update_controller(self): diff --git a/stretch4_body/tools/stretch_body_server.py b/stretch4_body/tools/stretch_body_server.py index cef3f53..e518159 100644 --- a/stretch4_body/tools/stretch_body_server.py +++ b/stretch4_body/tools/stretch_body_server.py @@ -37,6 +37,11 @@ def print_status(robot_client): print('Current rate (Hz): %.2f'%robot_client.status['server']['control_loop']['curr_rate_hz']) print('Loop count: ' + str(robot_client.status['server']['control_loop']['num_loops'])) print('Loop overruns: ' + str(robot_client.status['server']['control_loop']['missed_loops'])) + print('Running on user: ' +f'{StretchBodyServer.get_server_owning_user()}') + + if daemon_is_installed(): + status_daemon() + print("") print("Use `stretch_body_server --print` to view the latest logs.") print("") @@ -282,69 +287,36 @@ def main(): exit(0) if args.daemon: - if not start_daemon(): - raise RuntimeError("Could not start the Stretch Body Server system service") - tail_log_file(log_file) - return - - - if args.install_daemon: + if not install_and_start_daemon(): + raise RuntimeError("Could not start the Stretch Body Server system service.") + time.sleep(2.0) # Wait for daemon to start + elif args.install_daemon: if not install_daemon(): - raise RuntimeError("Could not install the Stretch Body Server system service") - - - if args.uninstall_daemon: + raise RuntimeError("Could not install the Stretch Body Server system service.") + return print("Stretch Body Server daemon is installed. Server will start on boot. Run `--daemon` to start immediately.") + elif args.uninstall_daemon: if not uninstall_daemon(): - raise RuntimeError("Could not uninstall the Stretch Body Server system service") + raise RuntimeError("Could not uninstall the Stretch Body Server system service.") + return print("Uninstalled Stretch Body Server daemon. Server will not start on boot.") - if args.launch or args.restart: - if args.restart: - if restart_daemon(): # if the daemon is running, restart it. Otherwise, launch in the terminal. - print("\nTailing logs, press Ctrl+C to exit (server will keep running in the background):") - tail_log_file(log_file) - return - if not robot_client.startup(): + elif args.launch: + return _launch(robot_client, args) + elif args.restart: + if not daemon_is_installed(): + if is_server_active(robot_client): print("Stopping existing server...") - robot_client.kill_server() + with robot_client: + robot_client.kill_server() time.sleep(2.0) # Wait for shutdown else: print("No instances of the server was found. Launching a new instance.") - else: # args.launch is true - if is_server_active(robot_client): - if StretchBodyServer.is_server_owned_by_current_user(): - print(f""" -=============================================== - -StretchBodyClient: A server is already running. -StretchBodyClient: You can run `stretch_body_server --kill` to forcefully end the running session. -StretchBodyClient: You can run `stretch_body_server --restart` to forcefully restart the running session. - -=============================================== - -""") - else: - print(f""" -=============================================== - -StretchBodyClient: A server is already running, but it was started by a different user ({StretchBodyServer.get_server_owning_user()}). -StretchBodyClient: You can run `stretch_body_server --kill` to forcefully end the other user's session. - -=============================================== - -""") - return + return _launch(robot_client, args) - try: - if args.profile: - os.environ['STRETCH_PROFILE'] = '1' - robot_server.run_server() - except Exception as e: - logger.error(f"Unexpected error while running stretch body server: {e}") - finally: - archive_session_logs() - - return + if not install_and_start_daemon(): + raise RuntimeError(f"Could not restart the Stretch Body Server system service.") + + time.sleep(2.0) # Wait for daemon to start - if args.kill: + elif args.kill: if not is_server_active(robot_client): print("No server is running.") return @@ -399,11 +371,51 @@ def main(): robot_client.stop() else: - print("No active server found. Printing tail of archived logs:\n") print_last_log_from_archive() + + if daemon_is_installed(): + status_daemon() + + print(f"\n\n----------------------\n{Fore.RED + Style.BRIGHT}No active server found. Run `stretch_body_server --restart` to start it.{Style.RESET_ALL} Printed tail of archived logs{' and background-daemon status ' if daemon_is_installed() else ' '}above.\n") return +def _launch(robot_client, args): + + if is_server_active(robot_client): + if StretchBodyServer.is_server_owned_by_current_user(): + print(f""" +=============================================== + +StretchBodyClient: A server is already running. +StretchBodyClient: You can run `stretch_body_server --kill` to forcefully end the running session. +StretchBodyClient: You can run `stretch_body_server --restart` to forcefully restart the running session. + +=============================================== + +""") + else: + print(f""" +=============================================== + +StretchBodyClient: A server is already running, but it was started by a different user ({StretchBodyServer.get_server_owning_user()}). +StretchBodyClient: You can run `stretch_body_server --kill` to forcefully end the other user's session. + +=============================================== + +""") + return False + + try: + if args.profile: + os.environ['STRETCH_PROFILE'] = '1' + robot_server.run_server() + except Exception as e: + logger.error(f"Unexpected error while running stretch body server: {e}") + finally: + archive_session_logs() + + return True def get_service_file_content(log_level: str) -> str: """Creates the linux service file that gets copied to ~/.config/systemd/user/""" @@ -469,7 +481,7 @@ def _manage_daemon(action:str): if action == "stop": action_with_suffix = "stopping" elif action == "status": action_with_suffix = "getting status of" log_level = RobotParams._robot_params['logging']['root']['level'] - logger.info(f"{action_with_suffix.capitalize()} Stretch Body Server systemd service...") + # logger.info(f"{action_with_suffix.capitalize()} Stretch Body Server systemd service...") user_systemd_dir = Path.home() / ".config" / "systemd" / "user" service_file_path = user_systemd_dir / "stretch_body_server.service" @@ -481,42 +493,38 @@ def _manage_daemon(action:str): f.write(get_service_file_content(log_level)) subprocess.run(["systemctl", "--user", "daemon-reload"], check=True) subprocess.run(["systemctl", "--user", "enable", "stretch_body_server.service"], check=True) + print(f"Installed the service to {service_file_path}") elif action == "uninstall": subprocess.run(["systemctl", "--user", "stop", "stretch_body_server.service"], check=False) subprocess.run(["systemctl", "--user", "disable", "stretch_body_server.service"], check=False) if service_file_path.exists(): service_file_path.unlink() + print(f"Removed the service from {service_file_path}") subprocess.run(["systemctl", "--user", "daemon-reload"], check=True) elif action == "start": - user_systemd_dir.mkdir(parents=True, exist_ok=True) - with open(service_file_path, "w") as f: - f.write(get_service_file_content(log_level)) - subprocess.run(["systemctl", "--user", "daemon-reload"], check=True) subprocess.run(["systemctl", "--user", "start", "stretch_body_server.service"], check=True) elif action == "stop": subprocess.run(["systemctl", "--user", "stop", "stretch_body_server.service"], check=True) elif action == "restart": - user_systemd_dir.mkdir(parents=True, exist_ok=True) - with open(service_file_path, "w") as f: - f.write(get_service_file_content(log_level)) - subprocess.run(["systemctl", "--user", "daemon-reload"], check=True) subprocess.run(["systemctl", "--user", "restart", "stretch_body_server.service"], check=True) elif action == "status": + print(f"\n\n----- Stretch Body Server Background Service Status -----") + print(f"Installed at: {service_file_path}\n\n") subprocess.run(["systemctl", "--user", "status", "stretch_body_server.service", "--no-pager"], check=False) subprocess.run(["journalctl", "--user", "-u", "stretch_body_server.service", "-n", "20", "--no-pager"], check=False) - logger.info(f"{action_with_suffix.capitalize()} Stretch Body Server systemd service: SUCCESS") + # logger.info(f"{action_with_suffix.capitalize()} Stretch Body Server systemd service: SUCCESS") return True except Exception as e: logger.error(f"Error while {action_with_suffix} Stretch Body Server system service: {e}") return False -def start_daemon() -> bool: +def install_and_start_daemon() -> bool: """Installs the systemctl service, and then calls restart to (re)start it, if there isn't already an active non-daemon server running. Note: It calls restart in case a service is already running.""" if not _manage_daemon("install"): return False @@ -557,6 +565,22 @@ def daemon_is_running() -> bool: except FileNotFoundError: print("Error: 'systemctl' command not found. Are you on a systemd Linux machine?") return False + +def daemon_is_installed() -> bool: + try: + result = subprocess.run( + ['systemctl', '--user', 'is-enabled', 'stretch_body_server.service'], + capture_output=True, + text=True, + check=False + ) + if result.returncode == 0 and result.stdout.strip() != 'not-found': + return True + else: + return False + except FileNotFoundError: + print("Error: 'systemctl' command not found. Are you on a systemd Linux machine?") + return False def uninstall_daemon() -> bool: return _manage_daemon("uninstall") diff --git a/stretch4_body/tools/stretch_configure_tool.py b/stretch4_body/tools/stretch_configure_tool.py index a3a59a0..52028d8 100644 --- a/stretch4_body/tools/stretch_configure_tool.py +++ b/stretch4_body/tools/stretch_configure_tool.py @@ -95,22 +95,27 @@ def main(quick, auto_detect): detected_tool = None if not quick: try: - try: - from stretch4_body.robot.robot_client import PowerPeriphClient as PowerPeriph - is_client = True - except ImportError: - from stretch4_body.subsystem.power_periph import PowerPeriph - is_client = False + direct = False + from stretch4_body.robot.robot_client import PowerPeriphClient as PowerPeriph p = PowerPeriph() - p.startup() + + if not p.startup(): + # If the client can't connect, try without the client: + direct = True + + from stretch4_body.subsystem.power_periph import PowerPeriph + + p = PowerPeriph() + + if not p.startup(): + return print("Failed to connect to the robot's power management. Please run `stretch_system_check`.") if not auto_detect: if click.confirm('Turn off power to the peripheral?', default=True): print('Powering off eoa...') p.actuator_control('eoa', enable=False) - if is_client: - p.push_command() + p.push_command() try: click.pause('Connect the tool then press any key to continue...') @@ -121,8 +126,7 @@ def main(quick, auto_detect): print('Powering on eoa...') p.actuator_control('eoa', enable=True) - if is_client: - p.push_command() + p.push_command() time.sleep(2.0) # Wait for motors to boot # Auto-detect tool ID @@ -229,13 +233,24 @@ def main(quick, auto_detect): write_fleet_yaml(user_params_fn, _user_params, fleet_dir, user_params_header) print(f"Saved to {fleet_dir}{user_params_fn}") - - if click.confirm('\nWould you like to restart the stretch_body_server and home the end_of_arm?', default=True): - print('Restarting stretch_body_server...') - p_restart = subprocess.Popen(['stretch_body_server', '--restart'], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) - - print('Waiting for stretch_body_server to come back online...') - from stretch4_body.robot.robot_client import RobotClient + + is_do_home = False + if direct: + is_do_home = click.confirm("\nWould you like to home the end_of_arm?", default=True) + else: + is_do_home= click.confirm('\nWould you like to restart the stretch_body_server and home the end_of_arm?', default=True) + + if not quick and is_do_home: + p_restart = None + if not direct: + print('Restarting stretch_body_server...') + p_restart = subprocess.Popen(['stretch_body_server', '--restart'], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + + print('Waiting for stretch_body_server to come back online...') + if direct: + from stretch4_body.robot.robot import Robot as RobotClient + else: + from stretch4_body.robot.robot_client import RobotClient r = RobotClient() connected = False for i in range(20): # try for 20 seconds @@ -253,10 +268,13 @@ def main(quick, auto_detect): print(f"Error during homing: {e}") finally: r.stop() - p_restart.terminate() + if p_restart is not None: + p_restart.terminate() + print("Done! You are ready to use the tool.") else: print("Failed to connect to robot server after restart. Please try homing manually.") - p_restart.terminate() + if p_restart is not None: + p_restart.terminate() else: print("""Done! You may need to home the robot or restart services for the tool to be recognized. diff --git a/test/test_client_connection.py b/test/test_client_connection.py new file mode 100644 index 0000000..b1085b5 --- /dev/null +++ b/test/test_client_connection.py @@ -0,0 +1,424 @@ +#!/usr/bin/env python3 + +import os +import socket +import threading +import time +import pytest +import zmq + +from stretch4_body.core.client_server import StretchBodyServer +from stretch4_body.robot.robot_client import RobotClient + +# Keep reference to the original ZMQ methods for delegation +_original_bind = zmq.Socket.bind +_original_connect = zmq.Socket.connect + + +def find_free_port(): + """Helper to find a free TCP port.""" + s = socket.socket() + s.bind(('', 0)) + port = s.getsockname()[1] + s.close() + return port + + +@pytest.fixture(autouse=True) +def mock_network_env(monkeypatch, tmp_path): + """ + Isolate the test environment using monkeypatch. + Creates unique temporary directory for IPC socket files + and allocates free ephemeral TCP ports to prevent any collision + with system-wide server instances. + """ + # Force is_user_in_group to always return True for tests + monkeypatch.setattr("stretch4_body.core.client_server.is_user_in_group", lambda group: True) + monkeypatch.setattr("stretch4_body.utils.file_access_utils.is_user_in_group", lambda group: True) + + # Isolated temporary folder for socket paths + socket_dir = str(tmp_path / "stretch_zmq") + monkeypatch.setattr("stretch4_body.core.client_server.SERVER_ZMQ_SOCKET_DIR", socket_dir) + monkeypatch.setattr("stretch4_body.core.client_server.PORT_ADMIN", f"{socket_dir}/port_admin") + monkeypatch.setattr("stretch4_body.core.client_server.PORT_COMMAND", f"{socket_dir}/port_command") + monkeypatch.setattr("stretch4_body.core.client_server.PORT_STATUS", f"{socket_dir}/port_status") + + # Allocate ephemeral ports + admin_port = find_free_port() + cmd_port = find_free_port() + status_port = find_free_port() + + def mock_bind_impl(self, addr): + if "23114" in addr: + addr = addr.replace("23114", str(admin_port)) + elif "23115" in addr: + addr = addr.replace("23115", str(cmd_port)) + elif "23116" in addr: + addr = addr.replace("23116", str(status_port)) + return _original_bind(self, addr) + + def mock_connect_impl(self, addr): + if "23114" in addr: + addr = addr.replace("23114", str(admin_port)) + elif "23115" in addr: + addr = addr.replace("23115", str(cmd_port)) + elif "23116" in addr: + addr = addr.replace("23116", str(status_port)) + return _original_connect(self, addr) + + monkeypatch.setattr(zmq.Socket, "bind", mock_bind_impl) + monkeypatch.setattr(zmq.Socket, "connect", mock_connect_impl) + + +@pytest.fixture +def running_robot_server(): + """ + Yields a started StretchBodyServer instance and runs a mockup status publisher + periodically in a background thread to allow RobotClient to pull status successfully. + """ + server = StretchBodyServer() + assert server.startup() is True + + stop_event = threading.Event() + last_cmds = [] + + mock_status = { + "server": { + "state": "RUNNING", + "control_loop": {"target_rate_hz": 25.0, "curr_rate_hz": 25.0, "num_loops": 100, "missed_loops": 0}, + "cpu": {} + }, + "routines": { + "last_routine_id": None, + "last_routine_successful": True, + "active_routine": None + }, + "safety_layer": { + "safe_motion_manager": {"active": False}, + "sentry_manager": {"active": False} + }, + "robot": {}, + "power_periph": {}, + "arm": {"motor": {"pos_calibrated": True}}, + "lift": {"motor": {"pos_calibrated": True}}, + "omnibase": {}, + "end_of_arm": {}, + "line_sensor_loop": {} + } + + def cb_cmd(cmd_dict): + last_cmds.append(cmd_dict) + return [] + + def cb_admin(msg): + pass + + def loop(): + while not stop_event.is_set(): + try: + server.dispatch_admin_messages(cb_admin) + server.dispatch_command_messages(cb_cmd, is_routine_active=False) + server.publish_status(mock_status) + except Exception: + pass + time.sleep(0.01) + + t = threading.Thread(target=loop, daemon=True) + t.start() + + yield server, mock_status, last_cmds + + stop_event.set() + t.join(timeout=1.0) + server.stop() + + +def test_robot_client_basic_startup_shutdown(running_robot_server): + """ + Verify basic startup and shutdown sequence for RobotClient when the server is active. + """ + r = RobotClient() + assert r.startup(verbose=False) is True + assert r.connected is True + + r.stop() + assert r.connected is False + +def test_robot_client_double_startup(running_robot_server): + """ + Verify basic startup and shutdown sequence for RobotClient when the server is active. + """ + r = RobotClient() + assert r.startup(verbose=False) is True + assert r.startup(verbose=False) is True + assert r.connected is True + + r.stop() + assert r.connected is False + +def test_multiple_stop_and_startup(running_robot_server): + r = RobotClient() + + for i in range(10): + assert r.startup(verbose=False) is True + assert r.connected is True + + r.arm.move_by(0) + r.push_command() + + time.sleep(1/10) + + r.stop() + assert r.connected is False + + +def test_robot_client_context_manager(running_robot_server): + """ + Verify that RobotClient works correctly as a Python context manager. + """ + with RobotClient() as r: + assert r is not None + assert r.connected is True + + # After exiting the block, the client should automatically shut down + assert r.connected is False + + +def test_robot_client_offline_graceful(): + """ + Verify that RobotClient fails to startup gracefully if no server is running. + """ + r = RobotClient() + # Startup should return False without hanging or raising exceptions + assert r.startup(verbose=False) is False + assert r.connected is False + + +def test_robot_client_reconnection(running_robot_server): + """ + Verify subsequent connect, disconnect, and reconnect behaviors of RobotClient in the same session. + """ + # First connection + r1 = RobotClient() + assert r1.startup(verbose=False) is True + r1.stop() + assert r1.connected is False + + # Second connection + r2 = RobotClient() + assert r2.startup(verbose=False) is True + r2.stop() + assert r2.connected is False + + +def test_robot_client_multiple_clients(running_robot_server): + """ + Verify that multiple RobotClients can connect to the server at the same time. + """ + r1 = RobotClient(client_id="robot_client_1") + r2 = RobotClient(client_id="robot_client_2") + + assert r1.startup(verbose=False) is True + assert r2.startup(verbose=False) is True + + assert r1.connected is True + assert r2.connected is True + + r1.stop() + r2.stop() + + +def test_robot_client_command_queuing(running_robot_server, monkeypatch): + """ + Verify command queuing and dispatcher push from RobotClient to the server. + """ + server, mock_status, last_cmds = running_robot_server + + # Bypass pusher lock file to allow clean pushing in tests without lock file issues + monkeypatch.setattr("stretch4_body.utils.freeable_file_lock.FreeableFileLock.acquire", lambda self: True) + + with RobotClient() as r: + assert r is not None + # Check that we can command subsystems + if hasattr(r, 'power_periph'): + r.power_periph.trigger_beep() + if hasattr(r, 'omnibase'): + r.omnibase.translate_by(0.1, 0.0) + + # Allow handshakes to complete, then push command + time.sleep(0.1) + r.push_command() + + # Check server received commands + time.sleep(0.1) + assert len(last_cmds) >= 1 + + # Verify the format of sent command dictionary + cmd_sent = last_cmds[-1] + assert "omnibase" in cmd_sent or "power_periph" in cmd_sent +def test_robot_client_hang_on_server_disconnect(running_robot_server): + """ + Test that if a RobotClient is running and the server suddenly disappears/stops publishing, + a blocking pull_status() loop would block. We verify this thread-safety scenario + by running pull_status in a background thread and verifying it doesn't return + once the server is stopped, and we can handle it safely. + """ + server, mock_status, last_cmds = running_robot_server + + r = RobotClient() + assert r.startup(verbose=False) is True + + # 1. First pull_status succeeds when server is active + assert r.pull_status(blocking=True) is True + + # 2. Stop the server (simulating a crash or disconnect) + server.stop() + + # Drain any queued status messages remaining in the client's ZMQ buffer + for _ in range(100): + if not r.pull_status(blocking=False): + break + + # 3. Try to pull status in a thread + status_received = [] + def pull_worker(): + try: + res = r.pull_status(blocking=True) + status_received.append(res) + except Exception as e: + status_received.append(e) + + t = threading.Thread(target=pull_worker, daemon=True) + t.start() + + # Wait to see if it finishes or hangs + t.join(timeout=0.5) + + # Since the server is offline, pull_status(blocking=True) should be stuck waiting/hanging, + # so the thread should still be alive, and status_received should be empty! + print(f"\n--- DEBUG: status_received = {status_received} ---") + assert t.is_alive() is True + assert len(status_received) == 0 + + # Stop the client cleanly + r.stop() + + +def test_robot_client_multithreaded_gamepad_teleop_simulation(running_robot_server, monkeypatch): + """ + Simulates calibrate_intrinsics_robot_move.py scenario: + - Main thread starts RobotClient() -> `r_main` + - Background thread starts GamePadTeleop-like client -> `r_teleop` + - Both pull status and send commands concurrently. + - Test that they operate correctly concurrently, and stop gracefully without hangs. + """ + server, mock_status, last_cmds = running_robot_server + + # Bypass pusher lock file to allow clean pushing in concurrent threads without lock file issues + monkeypatch.setattr("stretch4_body.utils.freeable_file_lock.FreeableFileLock.acquire", lambda self: True) + + r_main = RobotClient(client_id="main_client") + r_teleop = RobotClient(client_id="teleop_client") + + assert r_main.startup(verbose=False) is True + assert r_teleop.startup(verbose=False) is True + + stop_event = threading.Event() + teleop_errors = [] + + def teleop_loop(): + while not stop_event.is_set(): + try: + # Teleop thread pulls status and pushes a command periodically + r_teleop.pull_status(blocking=True) + if hasattr(r_teleop, 'omnibase'): + r_teleop.omnibase.set_velocity(0.1, 0.0, 0.1) + r_teleop.push_command(ignore_control_lock=True) + time.sleep(0.01) + except Exception as e: + teleop_errors.append(e) + break + + t_teleop = threading.Thread(target=teleop_loop, daemon=True) + t_teleop.start() + + # Main thread also does operations concurrently + for _ in range(5): + assert r_main.pull_status(blocking=True) is True + time.sleep(0.02) + + # Clean stop + stop_event.set() + t_teleop.join(timeout=1.0) + + r_main.stop() + r_teleop.stop() + + assert len(teleop_errors) == 0 + + +def test_client_stop_no_hang_on_offline_server_with_pending_messages(running_robot_server): + """ + Verifies that calling stop() on a client with pending messages (e.g., when the + server is offline and cannot receive them) does not hang. + Previously, without linger=0 and context.destroy(linger=0), this would hang indefinitely. + """ + server, mock_status, last_cmds = running_robot_server + + r = RobotClient(client_id="hang_test_client") + assert r.startup(verbose=False) is True + + # Stop the server so it is offline + server.stop() + + # Wait a bit for the server ports to close + time.sleep(0.1) + + # Queue some commands. This places them in the socket's outbound buffer. + # Because the server is offline, they cannot be delivered. + if hasattr(r, 'omnibase'): + r.omnibase.set_velocity(0.1, 0.0, 0.1) + r.push_command(ignore_control_lock=True) + + # Calling stop() must complete instantly (within 1.5 seconds). + # If there was a linger, this would block indefinitely. + start_time = time.monotonic() + r.stop() + duration = time.monotonic() - start_time + + assert duration < 1.5, f"stop() took {duration:.3f}s, which indicates a linger/hang!" + + +def test_client_gc_no_hang_on_offline_server_with_pending_messages(running_robot_server): + """ + Verifies that garbage collection/destruction of a RobotClient with pending + messages (when the server is offline) does not hang. + Previously, GC of the context during __del__ would block indefinitely inside Context.term(). + """ + server, mock_status, last_cmds = running_robot_server + + r = RobotClient(client_id="gc_test_client") + assert r.startup(verbose=False) is True + + # Stop the server so it is offline + server.stop() + time.sleep(0.1) + + # Send a command to queue message in outbound buffer + if hasattr(r, 'omnibase'): + r.omnibase.set_velocity(0.1, 0.0, 0.1) + r.push_command(ignore_control_lock=True) + + # Trigger deletion/garbage collection of the client. + # It must complete instantly. + start_time = time.monotonic() + del r + import gc + gc.collect() + duration = time.monotonic() - start_time + + assert duration < 1.5, f"Garbage collection took {duration:.3f}s, indicating a destructor/Context.term() hang!" + + +if __name__ == '__main__': + pytest.main([__file__]) diff --git a/test/test_command_overwriting.py b/test/test_command_overwriting.py index 135dadf..c3becef 100644 --- a/test/test_command_overwriting.py +++ b/test/test_command_overwriting.py @@ -55,3 +55,6 @@ def test_cross_joint_overwriting(capsys): assert len(r.cmd_dict) == 2 assert f"{joint1}.end_of_arm" in r.cmd_dict assert f"{joint2}.end_of_arm" in r.cmd_dict + +if __name__ == '__main__': + pytest.main([__file__]) diff --git a/test/test_require_connection.py b/test/test_require_connection.py index a7698ab..7bdf185 100644 --- a/test/test_require_connection.py +++ b/test/test_require_connection.py @@ -17,7 +17,7 @@ def test_client_throws_exception(self): client._do_send_cmd({'test': 1}) with self.assertRaises(NotConnectedError): - client._do_send_recv_admin_str(b"ping") + client.do_send_recv_admin_str(b"ping") @patch('stretch4_body.core.robot_params.RobotParams.get_params') def test_subsystem_client_throws_exception(self, mock_get_params): @@ -36,8 +36,8 @@ def test_subsystem_client_throws_exception(self, mock_get_params): with self.assertRaises(NotConnectedError): sub.kill_server() - with self.assertRaises(NotConnectedError): - sub.pause_control_loop() + # with self.assertRaises(NotConnectedError): + # sub.pause_control_loop() if __name__ == '__main__': unittest.main() diff --git a/test/test_server_connection.py b/test/test_server_connection.py new file mode 100644 index 0000000..c265101 --- /dev/null +++ b/test/test_server_connection.py @@ -0,0 +1,353 @@ +#!/usr/bin/env python3 + +import os +import socket +import threading +import time +import pytest +import zmq + +from stretch4_body.core.client_server import ( + StretchBodyServer, + StretchBodyClient, + NotConnectedError, + LEASE_TIMEOUT +) + +# Keep reference to the original ZMQ methods for delegation +_original_bind = zmq.Socket.bind +_original_connect = zmq.Socket.connect + + +def find_free_port(): + """Helper to find a free TCP port.""" + s = socket.socket() + s.bind(('', 0)) + port = s.getsockname()[1] + s.close() + return port + + +@pytest.fixture(autouse=True) +def mock_network_env(monkeypatch, tmp_path): + """ + Isolate the test environment using monkeypatch. + Creates unique temporary directory for IPC socket files + and allocates free ephemeral TCP ports to prevent any collision + with system-wide server instances. + """ + # Force is_user_in_group to always return True for tests + monkeypatch.setattr("stretch4_body.core.client_server.is_user_in_group", lambda group: True) + monkeypatch.setattr("stretch4_body.utils.file_access_utils.is_user_in_group", lambda group: True) + + # Isolated temporary folder for socket paths + socket_dir = str(tmp_path / "stretch_zmq") + monkeypatch.setattr("stretch4_body.core.client_server.SERVER_ZMQ_SOCKET_DIR", socket_dir) + monkeypatch.setattr("stretch4_body.core.client_server.PORT_ADMIN", f"{socket_dir}/port_admin") + monkeypatch.setattr("stretch4_body.core.client_server.PORT_COMMAND", f"{socket_dir}/port_command") + monkeypatch.setattr("stretch4_body.core.client_server.PORT_STATUS", f"{socket_dir}/port_status") + + # Allocate ephemeral ports + admin_port = find_free_port() + cmd_port = find_free_port() + status_port = find_free_port() + + def mock_bind_impl(self, addr): + if "23114" in addr: + addr = addr.replace("23114", str(admin_port)) + elif "23115" in addr: + addr = addr.replace("23115", str(cmd_port)) + elif "23116" in addr: + addr = addr.replace("23116", str(status_port)) + return _original_bind(self, addr) + + def mock_connect_impl(self, addr): + if "23114" in addr: + addr = addr.replace("23114", str(admin_port)) + elif "23115" in addr: + addr = addr.replace("23115", str(cmd_port)) + elif "23116" in addr: + addr = addr.replace("23116", str(status_port)) + return _original_connect(self, addr) + + monkeypatch.setattr(zmq.Socket, "bind", mock_bind_impl) + monkeypatch.setattr(zmq.Socket, "connect", mock_connect_impl) + + +@pytest.fixture +def running_server(): + """ + Yields a started StretchBodyServer instance and runs its message dispatch loop + in a background thread. Automatically stops the server on cleanup. + """ + server = StretchBodyServer() + assert server.startup() is True + + last_cmds = [] + last_admin_msgs = [] + stop_event = threading.Event() + + def cb_cmd(cmd_dict): + last_cmds.append(cmd_dict) + return [] + + def cb_admin(msg): + last_admin_msgs.append(msg) + + def loop(): + while not stop_event.is_set(): + try: + server.dispatch_admin_messages(cb_admin) + server.dispatch_command_messages(cb_cmd, is_routine_active=False) + except Exception: + pass + time.sleep(0.005) + + t = threading.Thread(target=loop, daemon=True) + t.start() + + yield server, last_cmds, last_admin_msgs + + stop_event.set() + t.join(timeout=1.0) + server.stop() + + +def test_basic_connect_disconnect(running_server): + """ + Test that StretchBodyClient can connect, ping, send a command, + and disconnect cleanly from StretchBodyServer. + """ + server, last_cmds, last_admin_msgs = running_server + + client = StretchBodyClient() + assert client.startup(verbose=False) is True + assert client.connected is True + + # Test admin ping command + assert client.do_send_recv_admin_str(b"ping") == b"ping" + + # Test basic command sending + test_cmd = {"device": "arm", "val": 42} + client._do_send_cmd(test_cmd) + + # Allow background dispatch to process the command + time.sleep(0.05) + assert len(last_cmds) == 1 + assert last_cmds[0] == test_cmd + + # Clean disconnect + client.stop() + assert client.connected is False + + +def test_reconnection(running_server): + """ + Test subsequent client connect and disconnect actions in the same session. + """ + server, last_cmds, last_admin_msgs = running_server + + # First connection + client1 = StretchBodyClient(name="client_1") + assert client1.startup(verbose=False) is True + assert client1.connected is True + client1.stop() + assert client1.connected is False + + # Immediate second connection + client2 = StretchBodyClient(name="client_2") + assert client2.startup(verbose=False) is True + assert client2.connected is True + client2.stop() + + +def test_multiple_clients(running_server): + """ + Test multiple clients connecting to the same server simultaneously. + """ + server, last_cmds, last_admin_msgs = running_server + + c1 = StretchBodyClient(name="c1") + c2 = StretchBodyClient(name="c2") + + assert c1.startup(verbose=False) is True + assert c2.startup(verbose=False) is True + + assert c1.do_send_recv_admin_str(b"ping") == b"ping" + assert c2.do_send_recv_admin_str(b"ping") == b"ping" + + c1.stop() + c2.stop() + + +def test_lease_priority_and_expiry(running_server): + """ + Verify the server's command lease prioritization, rejection, + priority override, and expiration logic. + """ + server, last_cmds, last_admin_msgs = running_server + + c_low = StretchBodyClient(name="client_low") + c_high = StretchBodyClient(name="client_high") + + assert c_low.startup(verbose=False) is True + assert c_high.startup(verbose=False) is True + + # 1. c_low sends low-priority command (priority=1) -> gets lease + c_low._do_send_cmd({"cmd": "low_cmd"}, priority=1) + time.sleep(0.05) + assert len(last_cmds) == 1 + assert last_cmds[-1] == {"cmd": "low_cmd"} + assert server.lease_holder_id == c_low.client_id + assert server.lease_holder_priority == 1 + + # 2. c_high sends equal or lower priority command (priority=1) -> rejected/ignored + c_high._do_send_cmd({"cmd": "rejected_cmd"}, priority=1) + time.sleep(0.05) + # The last_cmds list shouldn't grow, since it should be rejected because c_low holds lease + assert len(last_cmds) == 1 + + # 3. c_high sends HIGHER priority command (priority=2) -> overrides lease + c_high._do_send_cmd({"cmd": "high_cmd"}, priority=2) + time.sleep(0.05) + assert len(last_cmds) == 2 + assert last_cmds[-1] == {"cmd": "high_cmd"} + assert server.lease_holder_id == c_high.client_id + assert server.lease_holder_priority == 2 + + # 4. Wait for lease to expire (LEASE_TIMEOUT is 1.1s) + # Let's wait slightly more than 1.1s + time.sleep(LEASE_TIMEOUT + 0.1) + + # 5. Now any client can grab the lease again. c_low sends a low priority command -> succeeds! + c_low._do_send_cmd({"cmd": "fresh_cmd"}, priority=1) + time.sleep(0.05) + assert len(last_cmds) == 3 + assert last_cmds[-1] == {"cmd": "fresh_cmd"} + assert server.lease_holder_id == c_low.client_id + + c_low.stop() + c_high.stop() + + +def test_command_conflation(running_server): + """ + Test that when commands are sent rapidly, sequence numbers increment + and the server processes them. + """ + server, last_cmds, last_admin_msgs = running_server + + client = StretchBodyClient() + assert client.startup(verbose=False) is True + # Allow ZMQ PUB/SUB subscription handshake to complete before publishing + time.sleep(0.1) + + # Send 5 commands rapidly + for i in range(5): + client._do_send_cmd({"cmd": f"msg_{i}"}) + + time.sleep(0.1) + # Since CONFLATE is on, we might not receive all of them, but we should receive at least the last one. + assert len(last_cmds) >= 1 + assert last_cmds[-1] == {"cmd": "msg_4"} + + client.stop() + + +def test_multiple_servers(tmp_path, monkeypatch): + """ + Verify that two independent servers can start on completely different paths + without interfering with each other's sockets or locks. + """ + # Server 1 environment + dir1 = str(tmp_path / "server1") + os.makedirs(dir1, exist_ok=True) + + # Server 2 environment + dir2 = str(tmp_path / "server2") + os.makedirs(dir2, exist_ok=True) + + # Allocate non-conflicting dynamic ports + s1_admin = find_free_port() + s1_cmd = find_free_port() + s1_status = find_free_port() + + s2_admin = find_free_port() + s2_cmd = find_free_port() + s2_status = find_free_port() + + # Define Server 1 with its own unique ports and paths + server1 = StretchBodyServer() + + # Patch server1 socket bindings + def s1_bind_impl(self, addr): + if "23114" in addr: + addr = f"tcp://*:{s1_admin}" + elif "23115" in addr: + addr = f"tcp://*:{s1_cmd}" + elif "23116" in addr: + addr = f"tcp://*:{s1_status}" + return _original_bind(self, addr) + + # Define Server 2 with its own unique ports and paths + server2 = StretchBodyServer() + + def s2_bind_impl(self, addr): + if "23114" in addr: + addr = f"tcp://*:{s2_admin}" + elif "23115" in addr: + addr = f"tcp://*:{s2_cmd}" + elif "23116" in addr: + addr = f"tcp://*:{s2_status}" + return _original_bind(self, addr) + + # Startup server 1 + with monkeypatch.context() as mctx: + mctx.setattr("stretch4_body.core.client_server.SERVER_ZMQ_SOCKET_DIR", dir1) + mctx.setattr("stretch4_body.core.client_server.PORT_ADMIN", f"{dir1}/port_admin") + mctx.setattr("stretch4_body.core.client_server.PORT_COMMAND", f"{dir1}/port_command") + mctx.setattr("stretch4_body.core.client_server.PORT_STATUS", f"{dir1}/port_status") + mctx.setattr(zmq.Socket, "bind", s1_bind_impl) + assert server1.startup() is True + + # Startup server 2 simultaneously + with monkeypatch.context() as mctx: + mctx.setattr("stretch4_body.core.client_server.SERVER_ZMQ_SOCKET_DIR", dir2) + mctx.setattr("stretch4_body.core.client_server.PORT_ADMIN", f"{dir2}/port_admin") + mctx.setattr("stretch4_body.core.client_server.PORT_COMMAND", f"{dir2}/port_command") + mctx.setattr("stretch4_body.core.client_server.PORT_STATUS", f"{dir2}/port_status") + mctx.setattr(zmq.Socket, "bind", s2_bind_impl) + assert server2.startup() is True + + # Stop both servers cleanly + server1.stop() + server2.stop() + + +def test_connection_failures_gracefully(): + """ + Verify client handling when the server is offline or not running. + """ + client = StretchBodyClient() + # Startup should return False gracefully and print a message without crashes + assert client.startup(verbose=False) is False + assert client.connected is False + + +def test_require_connection_decorator(): + """ + Verify that calling connection-requiring methods on an unconnected client + raises NotConnectedError. + """ + client = StretchBodyClient() + # We force the is_valid attribute to be False to raise NotConnectedError + # (or leave it undefined, which also triggers NotConnectedError) + with pytest.raises(NotConnectedError): + client._do_send_cmd({"cmd": "noop"}) + + with pytest.raises(NotConnectedError): + client._do_recv_status() + + +if __name__ == '__main__': + pytest.main([__file__])