Merge pull request #238 from robotpy/separate-process

Add option to launch each test in separate process

Author: virtuald

pyfrc/mains/cli_test.py

@@ -1,3 +1,4 @@
+import logging
 import os
 from os.path import abspath
 import inspect
@@ -7,6 +8,7 @@
 
 import wpilib
 
+import tomli
 import pytest
 
 from ..util import yesno
@@ -16,6 +18,8 @@
 # TODO: setting the plugins so that the end user can invoke pytest directly
 # could be a useful thing. Will have to consider that later.
 
+logger = logging.getLogger("test")
+
 
 class _TryAgain(Exception):
     pass
@@ -37,6 +41,12 @@ def __init__(self, parser=None):
                 action="store_true",
                 help="Use pyfrc's builtin tests if no tests are specified",
             )
+            parser.add_argument(
+                "--isolated",
+                default=None,
+                action="store_true",
+                help="Run each test in a separate robot process. Set `tool.robotpy.pyfrc.isolated` to true in your pyproject.toml to enable by default",
+            )
             parser.add_argument(
                 "--coverage-mode",
                 default=False,
@@ -55,16 +65,47 @@ def run(
         project_path: pathlib.Path,
         robot_class: typing.Type[wpilib.RobotBase],
         builtin: bool,
+        isolated: typing.Optional[bool],
         coverage_mode: bool,
+        verbose: bool,
         pytest_args: typing.List[str],
     ):
+        if isolated is None:
+            pyproject_path = project_path / "pyproject.toml"
+            if pyproject_path.exists():
+                with open(pyproject_path, "rb") as fp:
+                    d = tomli.load(fp)
+
+                try:
+                    v = d["tool"]["robotpy"]["pyfrc"]["isolated"]
+                except KeyError:
+                    pass
+                else:
+                    if not isinstance(v, bool):
+                        raise ValueError(
+                            f"tool.robotpy.pyfrc.isolated must be a boolean value (got {v})"
+                        )
+
+                    isolated = v
+
+        if isolated is None:
+            isolated = False
+
+        if not isolated:
+            logger.info(
+                "Isolated test mode not enabled, consider using it if your tests hang"
+            )
+            logger.info("- See 'robotpy test --help' for details")
+
         try:
             return self._run_test(
                 main_file,
                 project_path,
                 robot_class,
                 builtin,
+                isolated,
                 coverage_mode,
+                verbose,
                 pytest_args,
             )
         except _TryAgain:
@@ -73,7 +114,9 @@ def run(
                 project_path,
                 robot_class,
                 builtin,
+                isolated,
                 coverage_mode,
+                verbose,
                 pytest_args,
             )
 
@@ -83,7 +126,9 @@ def _run_test(
         project_path: pathlib.Path,
         robot_class: typing.Type[wpilib.RobotBase],
         builtin: bool,
+        isolated: bool,
         coverage_mode: bool,
+        verbose: bool,
         pytest_args: typing.List[str],
     ):
         # find test directory, change current directory so pytest can find the tests
@@ -92,14 +137,15 @@ def _run_test(
         curdir = pathlib.Path.cwd().absolute()
 
         self.try_dirs = [
-            (project_path / "tests").absolute(),
-            (project_path / ".." / "tests").absolute(),
+            ((project_path / "tests").absolute(), False),
+            ((project_path / ".." / "tests").absolute(), True),
         ]
 
-        for d in self.try_dirs:
+        for d, chdir in self.try_dirs:
             if d.exists():
-                test_directory = d
-                os.chdir(test_directory)
+                builtin = False
+                if chdir:
+                    os.chdir(d)
                 break
         else:
             if not builtin:
@@ -112,10 +158,22 @@ def _run_test(
             pytest_args.insert(0, abspath(inspect.getfile(basic)))
 
         try:
-            retv = pytest.main(
-                pytest_args,
-                plugins=[pytest_plugin.PyFrcPlugin(robot_class, main_file)],
-            )
+            if isolated:
+                from ..test_support import pytest_dist_plugin
+
+                retv = pytest.main(
+                    pytest_args,
+                    plugins=[
+                        pytest_dist_plugin.DistPlugin(
+                            robot_class, main_file, builtin, verbose
+                        )
+                    ],
+                )
+            else:
+                retv = pytest.main(
+                    pytest_args,
+                    plugins=[pytest_plugin.PyFrcPlugin(robot_class, main_file, False)],
+                )
         finally:
             os.chdir(curdir)
 
@@ -132,7 +190,7 @@ def _no_tests(
     ):
         print()
         print("Looked for tests at:")
-        for d in self.try_dirs:
+        for d, _ in self.try_dirs:
             print("-", d)
         print()
         print(

pyfrc/test_support/pytest_dist_plugin.py

@@ -0,0 +1,177 @@
+import multiprocessing
+import os
+import pathlib
+import sys
+import time
+
+from typing import Type
+
+import pytest
+
+import robotpy.main
+import robotpy.logconfig
+import wpilib
+
+
+from .pytest_plugin import PyFrcPlugin
+
+
+def _enable_faulthandler():
+    #
+    # In the event of a segfault, faulthandler will dump the currently
+    # active stack so you can figure out what went wrong.
+    #
+    # Additionally, on non-Windows platforms we register a SIGUSR2
+    # handler -- if you send the robot process a SIGUSR2, then
+    # faulthandler will dump all of your current stacks. This can
+    # be really useful for figuring out things like deadlocks.
+    #
+
+    import logging
+
+    logger = logging.getLogger("faulthandler")
+
+    try:
+        # These should work on all platforms
+        import faulthandler
+
+        faulthandler.enable()
+    except Exception as e:
+        logger.warning("Could not enable faulthandler: %s", e)
+        return
+
+    try:
+        import signal
+
+        faulthandler.register(signal.SIGUSR2)
+        logger.info("registered SIGUSR2 for PID %s", os.getpid())
+    except Exception:
+        return
+
+
+def _run_test(item_nodeid, config_args, robot_class, robot_file, verbose, pipe):
+    """This function runs in a subprocess"""
+    robotpy.logconfig.configure_logging(verbose)
+    _enable_faulthandler()
+
+    # This is used by getDeployDirectory, so make sure it gets fixed
+    robotpy.main.robot_py_path = robot_file
+
+    # keep the plugin around because it has a reference to the robot
+    # and we don't want it to die and deadlock
+    plugin = PyFrcPlugin(robot_class, robot_file, True)
+
+    ec = pytest.main(
+        [item_nodeid, "--no-header", *config_args],
+        plugins=[plugin],
+    )
+
+    # ensure output is printed out
+    # .. TODO could implement pytest_runtestloop and send the
+    #    test result back to the parent and print it there?
+    sys.stdout.flush()
+
+    # Don't let the process die, let the parent kill us to avoid
+    # python interpreter badness
+    pipe.send(ec)
+
+    # ensure that the gc doesn't collect the plugin..
+    while plugin:
+        time.sleep(100)
+
+
+def _run_test_in_new_process(
+    test_function, config, robot_class, robot_file, builtin_tests, verbose
+):
+    """Run a test function in a new process."""
+
+    config_args = config.invocation_params.args
+    if builtin_tests:
+        item_nodeid = f"{config_args[0]}::{test_function.name}"
+        config_args = config_args[1:]
+    else:
+        item_nodeid = test_function.nodeid
+
+    parent, child = multiprocessing.Pipe()
+
+    process = multiprocessing.Process(
+        target=_run_test,
+        args=(item_nodeid, config_args, robot_class, robot_file, verbose, child),
+    )
+    process.start()
+    try:
+        ec = parent.recv()
+    finally:
+        process.kill()
+
+    if ec != 0:
+        pytest.fail(
+            f"Test failed in subprocess: {item_nodeid} (exit code {ec})",
+            pytrace=False,
+        )
+
+
+def _make_runtest(item, config, robot_class, robot_file, builtin_tests, verbose):
+    def isolated_runtest():
+        _run_test_in_new_process(
+            item, config, robot_class, robot_file, builtin_tests, verbose
+        )
+
+    return isolated_runtest
+
+
+class DistPlugin:
+
+    def __init__(
+        self,
+        robot_class: Type[wpilib.RobotBase],
+        robot_file: pathlib.Path,
+        builtin_tests: bool,
+        verbose: bool,
+    ) -> None:
+        self._robot_class = robot_class
+        self._robot_file = robot_file
+        self._builtin_tests = builtin_tests
+        self._verbose = verbose
+
+    @pytest.hookimpl(tryfirst=True)
+    def pytest_collection_modifyitems(
+        self,
+        session: pytest.Session,
+        config: pytest.Config,
+        items: list[pytest.Function],
+    ):
+        """Modify collected test items to run each in a new process."""
+
+        multiprocessing.set_start_method("spawn")
+
+        for item in items:
+            # Overwrite the runtest protocol for each item
+            item.runtest = _make_runtest(
+                item,
+                config,
+                self._robot_class,
+                self._robot_file,
+                self._builtin_tests,
+                self._verbose,
+            )
+
+    #
+    # These fixtures match the ones in PyFrcPlugin but these have no effect
+    #
+
+    @pytest.fixture(scope="function", autouse=True)
+    def robot(self):
+        pass
+
+    @pytest.fixture(scope="function")
+    def control(self, reraise, robot):
+        pass
+
+    @pytest.fixture()
+    def robot_file(self):
+        pass
+
+    @pytest.fixture()
+    def robot_path(self):
+        pass

pyfrc/test_support/pytest_plugin.py

@@ -31,7 +31,14 @@ class PyFrcPlugin:
     be passed to your test function.
     """
 
-    def __init__(self, robot_class: Type[wpilib.RobotBase], robot_file: pathlib.Path):
+    def __init__(
+        self,
+        robot_class: Type[wpilib.RobotBase],
+        robot_file: pathlib.Path,
+        isolated: bool,
+    ):
+        self.isolated = isolated
+
         # attach physics
         physics, robot_class = PhysicsInterface._create_and_attach(
             robot_class,
@@ -103,6 +110,13 @@ def robot(self):
         # Tests only get a proxy to ensure cleanup is more reliable
         yield weakref.proxy(robot)
 
+        # If running in separate processes, no need to do cleanup
+        if self.isolated:
+            # .. and funny enough, in isolated mode we *don't* want the
+            # robot to be cleaned up, as that can deadlock
+            self._saved_robot = robot
+            return
+
         # reset engine to ensure it gets cleaned up too
         # -> might be holding wpilib objects, or the robot
         if self._physics:

setup.cfg

@@ -29,6 +29,7 @@ install_requires =
 
     wpilib>=2025.1.1,<2026
     robotpy-cli~=2024.0
+    tomli
 setup_requires =
     setuptools_scm > 6
 python_requires = >=3.9