refactor: always use multiprocess spawn

Author: hfudev

pytest-embedded/pytest_embedded/__init__.py

@@ -1,8 +1,52 @@
 """A pytest plugin that designed for embedded testing."""
 
-from .app import App
-from .dut import Dut
-from .dut_factory import DutFactory
+import multiprocessing.queues
+from multiprocessing.managers import BaseManager
+from typing import AnyStr
+
+MP_CTX = multiprocessing.get_context('spawn')
+
+
+class MessageQueueManager(BaseManager):
+    pass
+
+
+class MessageQueue(multiprocessing.queues.Queue):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, ctx=MP_CTX, **kwargs)
+
+    def put(self, obj, **kwargs):
+        from .utils import to_bytes
+
+        if not isinstance(obj, str | bytes):
+            super().put(obj, **kwargs)
+            return
+
+        if obj == '' or obj == b'':
+            return
+
+        _b = to_bytes(obj)
+        try:
+            super().put(_b)
+        except Exception:  # queue might be closed
+            pass
+
+    def write(self, s: AnyStr):
+        self.put(s)
+
+    def flush(self):
+        pass
+
+    def isatty(self):
+        return True
+
+
+MessageQueueManager.register('MessageQueue', MessageQueue)
+
+
+from .app import App  # noqa: E402
+from .dut import Dut  # noqa: E402
+from .dut_factory import DutFactory  # noqa: E402
 
 __all__ = ['App', 'Dut', 'DutFactory']
 

pytest-embedded/pytest_embedded/dut.py

@@ -1,6 +1,5 @@
 import functools
 import logging
-import multiprocessing
 import os.path
 import re
 from collections.abc import Callable
@@ -10,7 +9,7 @@
 import pexpect
 
 from .app import App
-from .log import PexpectProcess
+from .log import MessageQueue, PexpectProcess
 from .unity import UNITY_SUMMARY_LINE_REGEX, TestSuite
 from .utils import Meta, _InjectMixinCls, remove_asci_color_code, to_bytes, to_list
 
@@ -29,7 +28,7 @@ class Dut(_InjectMixinCls):
     def __init__(
         self,
         pexpect_proc: PexpectProcess,
-        msg_queue: multiprocessing.Queue,
+        msg_queue: MessageQueue,
         app: App,
         pexpect_logfile: str,
         test_case_name: str,

pytest-embedded/pytest_embedded/dut_factory.py

@@ -19,7 +19,7 @@
     from pytest_embedded_serial import Serial
     from pytest_embedded_wokwi import Wokwi
 
-from . import App, Dut
+from . import MP_CTX, App, Dut, MessageQueueManager
 from .log import MessageQueue, PexpectProcess
 from .utils import FIXTURES_SERVICES, ClassCliOptions, to_str
 
@@ -28,11 +28,6 @@ def _drop_none_kwargs(kwargs: dict[t.Any, t.Any]):
     return {k: v for k, v in kwargs.items() if v is not None}
 
 
-if sys.platform == 'darwin':
-    _ctx = multiprocessing.get_context('fork')
-else:
-    _ctx = multiprocessing.get_context()
-
 _stdout = sys.__stdout__
 
 
@@ -45,10 +40,6 @@ def _drop_none_kwargs(kwargs: dict[t.Any, t.Any]):
 PARAMETRIZED_FIXTURES_CACHE = {}
 
 
-def msg_queue_gn() -> MessageQueue:
-    return MessageQueue()
-
-
 def _listen(q: MessageQueue, filepath: str, with_timestamp: bool = True, count: int = 1, total: int = 1) -> None:
     shall_add_prefix = True
     while True:
@@ -94,7 +85,7 @@ def _listener_gn(msg_queue, _pexpect_logfile, with_timestamp, dut_index, dut_tot
         'total': dut_total,
     }
 
-    return _ctx.Process(
+    return MP_CTX.Process(
         target=_listen,
         args=(
             msg_queue,
@@ -742,7 +733,8 @@ def create(
         layout = []
         try:
             global PARAMETRIZED_FIXTURES_CACHE
-            msg_queue = msg_queue_gn()
+
+            msg_queue = MessageQueueManager().MessageQueue()
             layout.append(msg_queue)
 
             _pexpect_logfile = os.path.join(

pytest-embedded/pytest_embedded/log.py

@@ -1,56 +1,18 @@
 import errno
 import logging
-import multiprocessing
 import os
 import subprocess
-import sys
 import tempfile
 import textwrap
 import uuid
-from multiprocessing import queues
 from typing import AnyStr
 
 import pexpect.fdpexpect
 from pexpect import EOF, TIMEOUT
 from pexpect.utils import poll_ignore_interrupts, select_ignore_interrupts
 
-from .utils import Meta, remove_asci_color_code, to_bytes, to_str, utcnow_str
-
-if sys.platform == 'darwin':
-    _ctx = multiprocessing.get_context('fork')
-else:
-    _ctx = multiprocessing.get_context()
-
-
-class MessageQueue(queues.Queue):
-    def __init__(self, *args, **kwargs):
-        if 'ctx' not in kwargs:
-            kwargs['ctx'] = _ctx
-
-        super().__init__(*args, **kwargs)
-
-    def put(self, obj, **kwargs):
-        if not isinstance(obj, str | bytes):
-            super().put(obj, **kwargs)
-            return
-
-        if obj == '' or obj == b'':
-            return
-
-        _b = to_bytes(obj)
-        try:
-            super().put(_b, **kwargs)
-        except:  # noqa # queue might be closed
-            pass
-
-    def write(self, s: AnyStr):
-        self.put(s)
-
-    def flush(self):
-        pass
-
-    def isatty(self):
-        return True
+from pytest_embedded import MP_CTX, MessageQueue
+from pytest_embedded.utils import Meta, remove_asci_color_code, to_bytes, to_str, utcnow_str
 
 
 class PexpectProcess(pexpect.fdpexpect.fdspawn):
@@ -144,18 +106,18 @@ def live_print_call(*args, msg_queue: MessageQueue | None = None, expect_returnc
         raise subprocess.CalledProcessError(process.returncode, process.args)
 
 
-class _PopenRedirectProcess(_ctx.Process):
+class _PopenRedirectProcess(MP_CTX.Process):
     def __init__(self, msg_queue: MessageQueue, logfile: str):
-        self._q = msg_queue
-
-        self.logfile = logfile
-
-        super().__init__(target=self._forward_io, daemon=True)  # killed by the main process
+        super().__init__(target=self._forward_io, args=(msg_queue, logfile), daemon=True)
 
-    def _forward_io(self) -> None:
-        with open(self.logfile) as fr:
+    @staticmethod
+    def _forward_io(msg_queue, logfile) -> None:
+        with open(logfile) as fr:
             while True:
-                self._q.put(fr.read())
+                try:
+                    msg_queue.put(fr.read())  # msg_queue may be closed
+                except Exception:
+                    break
 
 
 class DuplicateStdoutPopen(subprocess.Popen):

pytest-embedded/pytest_embedded/plugin.py

@@ -26,6 +26,7 @@
 from _pytest.main import Session
 from _pytest.python import Function
 
+from . import MessageQueueManager
 from .app import App
 from .dut import Dut
 from .dut_factory import (
@@ -36,7 +37,6 @@
     app_fn,
     dut_gn,
     gdb_gn,
-    msg_queue_gn,
     openocd_gn,
     pexpect_proc_fn,
     qemu_gn,
@@ -460,6 +460,10 @@ def _close_or_terminate(obj):
                 if isinstance(obj, subprocess.Popen | multiprocessing.process.BaseProcess):
                     obj.terminate()
                     obj.kill()
+                elif isinstance(obj, MessageQueue):
+                    obj.close()
+                    obj.cancel_join_thread()
+                    obj.join_thread()
                 elif isinstance(obj, io.IOBase):
                     try:
                         obj.close()
@@ -630,6 +634,16 @@ def port_app_cache() -> dict[str, str]:
     return {}
 
 
+@pytest.fixture(scope='session')
+def _mp_manager():
+    manager = MessageQueueManager()
+    manager.start()
+
+    yield manager
+
+    manager.shutdown()
+
+
 @pytest.fixture
 def test_case_tempdir(test_case_name: str, session_tempdir: str) -> str:
     """Function scoped temp dir for pytest-embedded"""
@@ -668,8 +682,8 @@ def _pexpect_logfile(test_case_tempdir, logfile_extension, dut_index, dut_total)
 
 @pytest.fixture
 @multi_dut_generator_fixture
-def msg_queue() -> MessageQueue:  # kwargs passed by `multi_dut_generator_fixture()`
-    return msg_queue_gn()
+def msg_queue(_mp_manager) -> MessageQueue:  # kwargs passed by `multi_dut_generator_fixture()`
+    return _mp_manager.MessageQueue()
 
 
 @pytest.fixture

pytest-embedded/tests/test_base.py

@@ -290,7 +290,6 @@ def test_expect_all_failed(dut):
     result.assert_outcomes(passed=10)
 
 
-@pytest.mark.xfail(reason='unstable')
 def test_expect_from_timeout(testdir):
     testdir.makepyfile(r"""
         import threading