Simple tools to interact with processes, needs more work

Signed-off-by: Ivan Santiago Paunovic <[email protected]>

Author: ivanpauno

launch/launch/actions/execute_process.py

@@ -196,6 +196,7 @@ def __init__(
         ]] = None,
         respawn: bool = False,
         respawn_delay: Optional[float] = None,
+        cached_output: bool = False,
         **kwargs
     ) -> None:
         """
@@ -294,6 +295,8 @@ def __init__(
         :param: respawn if 'True', relaunch the process that abnormally died.
             Defaults to 'False'.
         :param: respawn_delay a delay time to relaunch the died process if respawn is 'True'.
+        :param: cached_output if `True`, both stdout and stderr will be cached.
+            Use get_stdout() and get_stderr() to read the buffered output.
         """
         super().__init__(**kwargs)
         self.__cmd = [normalize_to_list_of_substitutions(x) for x in cmd]
@@ -340,6 +343,7 @@ def __init__(
         self.__sigkill_timer = None  # type: Optional[TimerAction]
         self.__stdout_buffer = io.StringIO()
         self.__stderr_buffer = io.StringIO()
+        self.__cached_output = cached_output
 
         self.__executed = False
 
@@ -591,59 +595,32 @@ def __on_process_stdin(
         cast(ProcessStdin, event)
         return None
 
-    def __on_process_stdout(
-        self, event: ProcessIO
+    def __on_process_output(
+        self, event: ProcessIO, buffer, logger
     ) -> Optional[SomeActionsType]:
         to_write = event.text.decode(errors='replace')
-        if self.__stdout_buffer.closed:
-            # __stdout_buffer was probably closed by __flush_buffers on shutdown.  Output without
+        if buffer.closed:
+            # buffer was probably closed by __flush_buffers on shutdown.  Output without
             # buffering.
-            self.__stdout_logger.info(
-                self.__output_format.format(line=to_write, this=self)
-            )
-        else:
-            self.__stdout_buffer.write(to_write)
-            self.__stdout_buffer.seek(0)
-            last_line = None
-            for line in self.__stdout_buffer:
-                if line.endswith(os.linesep):
-                    self.__stdout_logger.info(
-                        self.__output_format.format(line=line[:-len(os.linesep)], this=self)
-                    )
-                else:
-                    last_line = line
-                    break
-            self.__stdout_buffer.seek(0)
-            self.__stdout_buffer.truncate(0)
-            if last_line is not None:
-                self.__stdout_buffer.write(last_line)
-
-    def __on_process_stderr(
-        self, event: ProcessIO
-    ) -> Optional[SomeActionsType]:
-        to_write = event.text.decode(errors='replace')
-        if self.__stderr_buffer.closed:
-            # __stderr buffer was probably closed by __flush_buffers on shutdown.  Output without
-            # buffering.
-            self.__stderr_logger.info(
+            buffer.info(
                 self.__output_format.format(line=to_write, this=self)
             )
         else:
-            self.__stderr_buffer.write(to_write)
-            self.__stderr_buffer.seek(0)
+            buffer.write(to_write)
+            buffer.seek(0)
             last_line = None
-            for line in self.__stderr_buffer:
+            for line in buffer:
                 if line.endswith(os.linesep):
-                    self.__stderr_logger.info(
+                    logger.info(
                         self.__output_format.format(line=line[:-len(os.linesep)], this=self)
                     )
                 else:
                     last_line = line
                     break
-            self.__stderr_buffer.seek(0)
-            self.__stderr_buffer.truncate(0)
+            buffer.seek(0)
+            buffer.truncate(0)
             if last_line is not None:
-                self.__stderr_buffer.write(last_line)
+                buffer.write(last_line)
 
     def __flush_buffers(self, event, context):
         line = self.__stdout_buffer.getvalue()
@@ -669,6 +646,35 @@ def __flush_buffers(self, event, context):
             self.__stderr_buffer.seek(0)
             self.__stderr_buffer.truncate(0)
 
+    def __on_process_output_cached(
+        self, event: ProcessIO, buffer, logger
+    ) -> Optional[SomeActionsType]:
+        to_write = event.text.decode(errors='replace')
+        last_cursor = buffer.tell()
+        self.__stdout_buffer.seek(0, 2)  # go to end of buffer
+        buffer.write(to_write)
+        buffer.seek(last_cursor)
+        new_cursor = last_cursor
+        for line in buffer:
+            if not line.endswith(os.linesep):
+                break
+            new_cursor = buffer.tell()
+            logger.info(
+                self.__output_format.format(line=line[:-len(os.linesep)], this=self)
+            )
+        buffer.seek(new_cursor)
+
+    def __flush_cached_buffers(self, event, context):
+        for line in self.__stdout_buffer:
+            self.__stdout_buffer.info(
+                self.__output_format.format(line=line, this=self)
+            )
+
+        for line in self.__stderr_buffer:
+            self.__stderr_logger.info(
+                self.__output_format.format(line=line, this=self)
+            )
+
     def __on_shutdown(self, event: Event, context: LaunchContext) -> Optional[SomeActionsType]:
         due_to_sigint = cast(Shutdown, event).due_to_sigint
         return self._shutdown_process(
@@ -903,6 +909,13 @@ def execute(self, context: LaunchContext) -> Optional[List[LaunchDescriptionEnti
             # If shutdown starts before execution can start, don't start execution.
             return None
 
+        if self.__cached_output:
+            on_output_method = self.__on_process_output_cached
+            flush_buffers_method = self.__flush_cached_buffers
+        else:
+            on_output_method = self.__on_process_output
+            flush_buffers_method = self.__flush_buffers
+
         event_handlers = [
             EventHandler(
                 matcher=lambda event: is_a_subclass(event, ShutdownProcess),
@@ -915,8 +928,10 @@ def execute(self, context: LaunchContext) -> Optional[List[LaunchDescriptionEnti
             OnProcessIO(
                 target_action=self,
                 on_stdin=self.__on_process_stdin,
-                on_stdout=self.__on_process_stdout,
-                on_stderr=self.__on_process_stderr
+                on_stdout=lambda event: on_output_method(
+                    event, self.__stdout_buffer, self.__stdout_logger),
+                on_stderr=lambda event: on_output_method(
+                    event, self.__stderr_buffer, self.__stderr_logger),
             ),
             OnShutdown(
                 on_shutdown=self.__on_shutdown,
@@ -927,7 +942,7 @@ def execute(self, context: LaunchContext) -> Optional[List[LaunchDescriptionEnti
             ),
             OnProcessExit(
                 target_action=self,
-                on_exit=self.__flush_buffers,
+                on_exit=flush_buffers_method,
             ),
         ]
         for event_handler in event_handlers:
@@ -985,3 +1000,30 @@ def shell(self):
     def prefix(self):
         """Getter for prefix."""
         return self.__prefix
+
+    def get_stdout(self):
+        """
+        Get cached stdout.
+
+        :raises RuntimeError: if cached_output is false.
+        """
+        if not self.__cached_output:
+            raise RuntimeError(f"cached output must be true to be able to get stdout, proc '{self.__name}'")
+        return self.__stdout_buffer.getvalue()
+
+    def get_stderr(self):
+        """
+        Get cached stdout.
+
+        :raises RuntimeError: if cached_output is false.
+        """
+        if not self.__cached_output:
+            raise RuntimeError(f"cached output must be true to be able to get stderr, proc '{self.__name}'")
+        return self.__stderr_buffer.getvalue()
+
+    @property
+    def return_code(self):
+        """Gets the process return code, None if it hasn't finished."""
+        if self._subprocess_transport is None:
+            return None
+        return self._subprocess_transport.get_returncode()

launch/launch/launch_context.py

@@ -174,9 +174,17 @@ def would_handle_event(self, event: Event) -> bool:
         """Check whether an event would be handled or not."""
         return any(handler.matches(event) for handler in self._event_handlers)
 
-    def register_event_handler(self, event_handler: BaseEventHandler) -> None:
-        """Register a event handler."""
-        self._event_handlers.appendleft(event_handler)
+    def register_event_handler(self, event_handler: BaseEventHandler, append = False) -> None:
+        """
+        Register a event handler.
+        
+        :param append: if 'true', the new event handler will be executed after the previously
+            registered ones. If not, it will prepend the old handlers.
+        """
+        if append:
+            self._event_handlers.append(event_handler)
+        else:
+            self._event_handlers.appendleft(event_handler)
 
     def unregister_event_handler(self, event_handler: BaseEventHandler) -> None:
         """Unregister an event handler."""

launch_testing/launch_testing/pytest/fixture.py

@@ -51,6 +51,19 @@ def launch_service(event_loop):
     return launch_service
 
 
+def get_launch_context_fixture(*, scope='function', overridable=True):
+    """Return a launch service fixture."""
+
+    @pytest.fixture(scope=scope)
+    def launch_context(launch_service):
+        """Create an instance of the launch service for each test case."""
+        return launch_service.context
+    if overridable:
+        launch_context._launch_testing_overridable_fixture = True
+        launch_context._launch_testing_fixture_scope = scope
+    return launch_context
+
+
 def get_event_loop_fixture(*, scope='function', overridable=True):
     """Return an event loop fixture."""
 
@@ -98,7 +111,8 @@ def fixture(
         mod_locals = vars(mod)
         for name, getter in (
             ('launch_service', get_launch_service_fixture),
-            ('event_loop', get_event_loop_fixture)
+            ('event_loop', get_event_loop_fixture),
+            ('launch_context', get_launch_context_fixture),
         ):
             if name in mod_locals:
                 obj = mod_locals[name]

launch_testing/launch_testing/pytest/plugin.py

@@ -27,6 +27,7 @@
 
 from .fixture import finalize_launch_service
 from .fixture import get_launch_service_fixture
+from .fixture import get_launch_context_fixture
 
 """
 launch_testing native pytest based implementation.
@@ -524,5 +525,8 @@ def inner(**kwargs):
     return inner
 
 
-"""Launch service fixture."""
 launch_service = get_launch_service_fixture(overridable=False)
+"""Launch service fixture."""
+
+launch_context = get_launch_context_fixture(overridable=False)
+"""Launch context fixture."""

launch_testing/launch_testing/tools/pytest_process.py

@@ -0,0 +1,117 @@
+# Copyright 2021 Open Source Robotics Foundation, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import contextlib
+import threading
+import time
+
+import launch
+from launch import event_handlers
+
+
+@contextlib.contextmanager
+def register_event_handler(context, event_handler):
+    # Code to acquire resource, e.g.:
+    try:
+        yield context.register_event_handler(event_handler, append=True)
+    finally:
+        context.unregister_event_handler(event_handler)
+
+def _get_on_process_start(execute_process_action, pyevent):
+    event_handlers.OnProcessStart(
+        target_action=execute_process_action, on_start=lambda _1, _2: pyevent.set())
+
+async def _wait_for_event(
+    launch_context, execute_process_action, get_launch_event_handler, timeout=None
+):
+    pyevent = asyncio.Event()
+    event_handler = get_launch_event_handler(execute_process_action, pyevent)
+    with register_event_handler(launch_context, event_handler):
+        await asyncio.wait_for(pyevent.wait(), timeout)
+
+async def _wait_for_event_with_condition(
+    launch_context, execute_process_action, get_launch_event_handler, condition, timeout=None
+):
+    pyevent = asyncio.Event()
+    event_handler = get_launch_event_handler(execute_process_action, pyevent)
+    cond_value = condition()
+    with register_event_handler(launch_context, event_handler):
+        start = time.time()
+        now = start
+        while not cond_value and (timeout is None or now < start + timeout):
+            await asyncio.wait_for(pyevent.wait(), start - now + timeout)
+            pyevent.clear()
+            cond_value = condition()
+            now = time.time()
+    return cond_value
+
+def _wait_for_event_sync(
+    launch_context, execute_process_action, get_launch_event_handler, timeout=None
+):
+    pyevent = threading.Event()
+    event_handler = get_launch_event_handler(execute_process_action, pyevent)
+    with register_event_handler(launch_context, event_handler):
+        pyevent.wait(timeout)
+
+def _wait_for_event_with_condition_sync(
+    launch_context, execute_process_action, get_launch_event_handler, condition, timeout=None
+):
+    pyevent = threading.Event()
+    event_handler = get_launch_event_handler(execute_process_action, pyevent)
+    cond_value = condition()
+    with register_event_handler(launch_context, event_handler):
+        start = time.time()
+        now = start
+        while not cond_value and (timeout is None or now < start + timeout):
+            pyevent.wait(start - now + timeout)
+            pyevent.clear()
+            cond_value = condition()
+            now = time.time()
+    return cond_value
+
+def _get_stdout_event_handler(action, pyevent):
+    return event_handlers.OnProcessIO(
+        target_action=action, on_stdout=lambda _1: pyevent.set())
+
+async def wait_for_output(
+    launch_context, execute_process_action, validate_output, timeout=None
+):
+    def condition():
+        try:
+            return validate_output(execute_process_action.get_stdout())
+        except AssertionError:
+            return False
+    success = await _wait_for_event_with_condition(
+        launch_context, execute_process_action, _get_stdout_event_handler, condition, timeout)
+    if not success:
+        # Validate the output again, this time not catching assertion errors.
+        # This allows the user to use asserts directly, errors will be nicely rendeded by pytest.
+        return validate_output(execute_process_action.get_stdout())
+
+
+def wait_for_output_sync(
+    launch_context, execute_process_action, validate_output, timeout=None
+):
+    def condition():
+        try:
+            return validate_output(execute_process_action.get_stdout())
+        except AssertionError:
+            return False
+    success = _wait_for_event_with_condition_sync(
+        launch_context, execute_process_action, _get_stdout_event_handler, condition, timeout)
+    if not success:
+        # Validate the output again, this time not catching assertion errors.
+        # This allows the user to use asserts directly, errors will be nicely rendeded by pytest.
+        return validate_output(execute_process_action.get_stdout()) in (None, True)