11import multiprocessing
2- import time
2+ import os
33import socket
44import sys
5+ import time
56import traceback
6- import os
7+ from contextlib import suppress
8+ from multiprocessing .context import ForkServerContext , SpawnContext
9+ from typing import Optional , Type , Union # Add Type and Union here
10+
711import pytest
8- from typing import Optional # Add this import at the top
9- from rcs .envs .creators import SimEnvCreator
10- from rcs .envs .utils import (
11- default_mujoco_cameraset_cfg ,
12- default_sim_gripper_cfg ,
13- default_sim_robot_cfg ,
14- )
1512from rcs .envs .base import ControlMode , RelativeTo
16- from rcs .rpc .server import RcsServer
13+ from rcs .envs .creators import SimEnvCreator
14+ from rcs .envs .utils import default_sim_gripper_cfg , default_sim_robot_cfg
1715from rcs .rpc .client import RcsClient
16+ from rcs .rpc .server import RcsServer
1817
1918HOST = "127.0.0.1"
2019
20+
2121def get_free_port () -> int :
2222 with socket .socket (socket .AF_INET , socket .SOCK_STREAM ) as s :
2323 s .bind ((HOST , 0 ))
2424 return s .getsockname ()[1 ]
2525
26+
2627def wait_for_port (
2728 host : str ,
2829 port : int ,
2930 timeout : float ,
3031 server_proc : Optional [multiprocessing .Process ] = None ,
31- err_q : Optional [multiprocessing .Queue ] = None
32+ err_q : Optional [multiprocessing .Queue ] = None ,
3233) -> None :
3334 start = time .time ()
3435 last_exc = None
@@ -44,21 +45,17 @@ def wait_for_port(
4445 if server_proc is not None and not server_proc .is_alive ():
4546 server_err = None
4647 if err_q is not None :
47- try :
48+ with suppress ( Exception ) :
4849 server_err = err_q .get_nowait ()
49- except Exception :
50- pass
5150 msg = f"Server process exited early (exitcode={ server_proc .exitcode } )."
5251 if server_err :
5352 msg += f"\n Server traceback:\n { server_err } "
5453 raise RuntimeError (msg )
5554 time .sleep (0.2 )
5655 server_err = None
5756 if err_q is not None :
58- try :
57+ with suppress ( Exception ) :
5958 server_err = err_q .get_nowait ()
60- except Exception :
61- pass
6259 msg = f"Timed out waiting for { host } :{ port } to open."
6360 if last_exc :
6461 msg += f" Last socket error: { last_exc } "
@@ -68,6 +65,7 @@ def wait_for_port(
6865 msg += f"\n Server traceback:\n { server_err } "
6966 raise TimeoutError (msg )
7067
68+
7169def run_server (host : str , port : int , err_q : multiprocessing .Queue ) -> None :
7270 try :
7371 env = SimEnvCreator ()(
@@ -76,7 +74,7 @@ def run_server(host: str, port: int, err_q: multiprocessing.Queue) -> None:
7674 robot_cfg = default_sim_robot_cfg (),
7775 gripper_cfg = default_sim_gripper_cfg (),
7876 # Disabled to avoid rendering problem in python subprocess.
79- #cameras=default_mujoco_cameraset_cfg(),
77+ # cameras=default_mujoco_cameraset_cfg(),
8078 max_relative_movement = 0.1 ,
8179 relative_to = RelativeTo .LAST_STEP ,
8280 )
@@ -90,20 +88,22 @@ def run_server(host: str, port: int, err_q: multiprocessing.Queue) -> None:
9088 time .sleep (1 )
9189 except Exception :
9290 tb = "" .join (traceback .format_exception (* sys .exc_info ()))
93- try :
91+ with suppress ( Exception ) :
9492 err_q .put (tb )
95- except Exception :
96- pass
9793 sys .exit (1 )
9894
99- def _mp_context () -> multiprocessing .context .BaseContext :
95+
96+ def _mp_context () -> Union [SpawnContext , ForkServerContext ]:
10097 # Prefer spawn to avoid fork-related issues with GL/MuJoCo/threaded libs
10198 methods = multiprocessing .get_all_start_methods ()
10299 if "spawn" in methods :
103100 return multiprocessing .get_context ("spawn" )
104101 if "forkserver" in methods :
105102 return multiprocessing .get_context ("forkserver" )
106- return multiprocessing .get_context (methods [0 ])
103+
104+ msg = "No suitable multiprocessing context found."
105+ raise RuntimeError (msg )
106+
107107
108108def _external_server_from_env () -> tuple [str , int ] | None :
109109 # Set RCS_TEST_HOST and RCS_TEST_PORT to reuse an already running server.
@@ -119,6 +119,7 @@ def _external_server_from_env() -> tuple[str, int] | None:
119119 return HOST , 50055
120120 return None
121121
122+
122123def test_run_server_starts_and_stops ():
123124 # Skip if reusing an external server
124125 ext = _external_server_from_env ()
@@ -130,17 +131,24 @@ def test_run_server_starts_and_stops():
130131 server_proc = ctx .Process (target = run_server , args = (HOST , port , err_q ))
131132 server_proc .start ()
132133 try :
133- wait_for_port (HOST , port , timeout = 120.0 , server_proc = server_proc , err_q = err_q )
134+ wait_for_port (HOST , port , timeout = 120.0 , server_proc = server_proc , err_q = err_q ) # type: ignore
134135 assert server_proc .is_alive (), "Server process did not start as expected."
135136 finally :
136137 if server_proc .is_alive ():
137138 server_proc .terminate ()
138139 server_proc .join (timeout = 5 )
139140 assert not server_proc .is_alive (), "Server process did not terminate as expected."
140141
142+
141143class TestRcsClientServer :
144+ client : RcsClient
145+ host : str = HOST
146+ port : int = 0
147+ server_proc = None
148+ err_q : Optional [multiprocessing .Queue ] = None
149+
142150 @classmethod
143- def setup_class (cls ):
151+ def setup_class (cls : Type [ "TestRcsClientServer" ] ):
144152 ext = _external_server_from_env ()
145153 if ext :
146154 cls .host , cls .port = ext
@@ -156,11 +164,11 @@ def setup_class(cls):
156164 cls .server_proc = ctx .Process (target = run_server , args = (cls .host , cls .port , cls .err_q ))
157165 cls .server_proc .start ()
158166 # Wait until the server is actually listening or fail early if it crashed
159- wait_for_port (cls .host , cls .port , timeout = 180.0 , server_proc = cls .server_proc , err_q = cls .err_q )
167+ wait_for_port (cls .host , cls .port , timeout = 180.0 , server_proc = cls .server_proc , err_q = cls .err_q ) # type: ignore
160168 cls .client = RcsClient (host = cls .host , port = cls .port )
161169
162170 @classmethod
163- def teardown_class (cls ):
171+ def teardown_class (cls : Type [ "TestRcsClientServer" ] ):
164172 try :
165173 if getattr (cls , "client" , None ):
166174 cls .client .close ()
@@ -188,8 +196,14 @@ def test_unwrapped(self):
188196 _ = self .client .unwrapped
189197
190198 def test_close (self ):
191- self .client .close ()
199+ if self .client is not None :
200+ self .client .close ()
192201 # Reconnect for further tests
193- wait_for_port (self .__class__ .host , self .__class__ .port , timeout = 15.0 ,
194- server_proc = self .__class__ .server_proc , err_q = self .__class__ .err_q )
202+ wait_for_port (
203+ self .__class__ .host ,
204+ self .__class__ .port ,
205+ timeout = 15.0 ,
206+ server_proc = self .__class__ .server_proc , # type: ignore
207+ err_q = self .__class__ .err_q ,
208+ )
195209 self .__class__ .client = RcsClient (host = self .__class__ .host , port = self .__class__ .port )
0 commit comments