Repository: Delgan/loguru (python)

Issue:
Add new "context" option to specify the multiprocessing context internally used by handlers
A better fix should be implemented for #746, because changing the default context globally using `set_start_method()` is not always suitable.

Invalid use case example on Linux (will hang forever because child processes crash):

```python
import multiprocessing
import sys

from loguru import logger


def set_logger(logger_):
    global logger
    logger = logger_


def log_test(num):
    print("before log")
    logger.info(f"hi {num}!")
    print("do we reach this?")
    logger.complete()


def main():
    logger.remove()
    logger.add(sys.stderr, enqueue=True)

    with multiprocessing.get_context("spawn").Pool(
        4, initializer=set_logger, initargs=(logger,)
    ) as pool:
        pool.map(log_test, list(range(10)))


if __name__ == "__main__":
    main()
```

Proposed solution using a new `context` attribute:

```python
import multiprocessing
import sys

from loguru import logger


def set_logger(logger_):
    global logger
    logger = logger_


def log_test(num):
    print("before log")
    logger.info(f"hi {num}!")
    print("do we reach this?")
    logger.complete()


def main():
    context = multiprocessing.get_context("spawn")
    logger.remove()
    logger.add(sys.stderr, enqueue=True, context=context)

    with context.Pool(4, initializer=set_logger, initargs=(logger,)) as pool:
        pool.map(log_test, list(range(10)))


if __name__ == "__main__":
    main()
```

Pull request description:
Add new "context" option for multiprocessing while adding handler
Fixes #851.

Patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 230ff2e..d267126 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -1,3 +1,9 @@
+`Unreleased`_
+=============
+
+- Add a new ``context`` optional argument to ``logger.add()`` specifying ``multiprocessing`` context (like ``"spawn"`` or ``"fork"``) to be used internally instead of the default one (`#851 <https://github.com/Delgan/loguru/issues/851>`_).
+
+
 `0.7.0`_ (2023-04-10)
 =====================
 
diff --git a/docs/resources/recipes.rst b/docs/resources/recipes.rst
index 531115e..6e53499 100644
--- a/docs/resources/recipes.rst
+++ b/docs/resources/recipes.rst
@@ -1027,53 +1027,18 @@ The |multiprocessing| library is also commonly used to start a pool of workers u
 
 Independently of the operating system, note that the process in which a handler is added with ``enqueue=True`` is in charge of the queue internally used. This means that you should avoid to ``.remove()`` such handler from the parent process is any child is likely to continue using it. More importantly, note that a |Thread| is started internally to consume the queue. Therefore, it is recommended to call |complete| before leaving |Process| to make sure the queue is left in a stable state.
 
-Another thing to keep in mind when dealing with multiprocessing is the fact that handlers created with ``enqueue=True`` create a queue internally in the current multiprocessing context. If they are passed through to a subprocesses instantiated within a different context (e.g. with :code:`multiprocessing.get_context("spawn")` on linux, where the default context is :code:`"fork"`) it will most likely result in crashing the subprocess. This is also noted in the `python multiprocessing docs <https://docs.python.org/3/library/multiprocessing.html#contexts-and-start-methods>`_.
+Another thing to keep in mind when dealing with multiprocessing is the fact that handlers created with ``enqueue=True`` create a queue internally in the default multiprocessing context. If they are passed through to a subprocesses instantiated within a different context (e.g. with ``multiprocessing.get_context("spawn")`` on linux, where the default context is ``"fork"``) it will most likely result in crashing the subprocess. This is also noted in the `python multiprocessing docs <https://docs.python.org/3/library/multiprocessing.html#contexts-and-start-methods>`_. To prevent any problems, you should specify the context to be used by Loguru while adding the handler. This can be done by passing the ``context`` argument to the ``add()`` method::
 
-So, running this on linux where the default context is ``fork`` this will not work since the handler is added in a different context:
-
-
-.. code::
-
-    # main.py
     import multiprocessing
     from loguru import logger
     import workers_a
-    import workers_b
 
     if __name__ == "__main__":
-        logger.remove()
-        logger.add("file.log", enqueue=True)
-
-        worker = workers_a.Worker()
-        with multiprocessing.get_context("spawn").Pool(4, initializer=worker.set_logger, initargs=(logger, )) as pool:
-            results = pool.map(worker.work, [1, 10, 100])
-
-        with multiprocessing.get_context("spawn").Pool(4, initializer=workers_b.set_logger, initargs=(logger, )) as pool:
-            results = pool.map(workers_b.work, [1, 10, 100])
-
-        logger.info("Done")
-
-To fix this you can set the multiprocessing context globally so that the handler is created in the same context as the subprocesses run in:
-
-.. code::
-
-    # main.py
-    import multiprocessing
-    from loguru import logger
-    import workers_a
-    import workers_b
-
-    if __name__ == "__main__":
-        multiprocessing.set_start_method("spawn")
+        context = multiprocessing.get_context("spawn")
 
         logger.remove()
-        logger.add("file.log", enqueue=True)
+        logger.add("file.log", enqueue=True, context=context)
 
         worker = workers_a.Worker()
-        with multiprocessing.Pool(4, initializer=worker.set_logger, initargs=(logger, )) as pool:
+        with context.Pool(4, initializer=worker.set_logger, initargs=(logger, )) as pool:
             results = pool.map(worker.work, [1, 10, 100])
-
-        with multiprocessing.Pool(4, initializer=workers_b.set_logger, initargs=(logger, )) as pool:
-            results = pool.map(workers_b.work, [1, 10, 100])
-
-        logger.info("Done")
diff --git a/loguru/_defaults.py b/loguru/_defaults.py
index 8f4d29b..d3d8de7 100644
--- a/loguru/_defaults.py
+++ b/loguru/_defaults.py
@@ -42,6 +42,7 @@ LOGURU_SERIALIZE = env("LOGURU_SERIALIZE", bool, False)
 LOGURU_BACKTRACE = env("LOGURU_BACKTRACE", bool, True)
 LOGURU_DIAGNOSE = env("LOGURU_DIAGNOSE", bool, True)
 LOGURU_ENQUEUE = env("LOGURU_ENQUEUE", bool, False)
+LOGURU_CONTEXT = env("LOGURU_CONTEXT", str, None)
 LOGURU_CATCH = env("LOGURU_CATCH", bool, True)
 
 LOGURU_TRACE_NO = env("LOGURU_TRACE_NO", int, 5)
diff --git a/loguru/_handler.py b/loguru/_handler.py
index d4563e5..e6545a9 100644
--- a/loguru/_handler.py
+++ b/loguru/_handler.py
@@ -1,6 +1,5 @@
 import functools
 import json
-import multiprocessing
 import os
 import threading
 from contextlib import contextmanager
@@ -41,6 +40,7 @@ class Handler:
         colorize,
         serialize,
         enqueue,
+        multiprocessing_context,
         error_interceptor,
         exception_formatter,
         id_,
@@ -55,6 +55,7 @@ class Handler:
         self._colorize = colorize
         self._serialize = serialize
         self._enqueue = enqueue
+        self._multiprocessing_context = multiprocessing_context
         self._error_interceptor = error_interceptor
         self._exception_formatter = exception_formatter
         self._id = id_
@@ -86,9 +87,9 @@ class Handler:
                 self._decolorized_format = self._formatter.strip()
 
         if self._enqueue:
-            self._queue = multiprocessing.SimpleQueue()
-            self._confirmation_event = multiprocessing.Event()
-            self._confirmation_lock = multiprocessing.Lock()
+            self._queue = self._multiprocessing_context.SimpleQueue()
+            self._confirmation_event = self._multiprocessing_context.Event()
+            self._confirmation_lock = self._multiprocessing_context.Lock()
             self._owner_process_pid = os.getpid()
             self._thread = Thread(
                 target=self._queued_writer, daemon=True, name="loguru-writer-%d" % self._id
diff --git a/loguru/_logger.py b/loguru/_logger.py
index bdc1fb4..d5c5f39 100644
--- a/loguru/_logger.py
+++ b/loguru/_logger.py
@@ -37,6 +37,7 @@
 .. |logging| replace:: :mod:`logging`
 .. |signal| replace:: :mod:`signal`
 .. |contextvars| replace:: :mod:`contextvars`
+.. |multiprocessing| replace:: :mod:`multiprocessing`
 .. |Thread.run| replace:: :meth:`Thread.run()<threading.Thread.run()>`
 .. |Exception| replace:: :class:`Exception`
 .. |AbstractEventLoop| replace:: :class:`AbstractEventLoop<asyncio.AbstractEventLoop>`
@@ -62,6 +63,9 @@
 .. _coroutine function: https://docs.python.org/3/glossary.html#term-coroutine-function
 .. |re.Pattern| replace:: ``re.Pattern``
 .. _re.Pattern: https://docs.python.org/3/library/re.html#re-objects
+.. |multiprocessing.Context| replace:: ``multiprocessing.Context``
+.. _multiprocessing.Context:
+   https://docs.python.org/3/library/multiprocessing.html#contexts-and-start-methods
 
 .. |better_exceptions| replace:: ``better_exceptions``
 .. _better_exceptions: https://github.com/Qix-/better-exceptions
@@ -82,7 +86,8 @@ import sys
 import warnings
 from collections import namedtuple
 from inspect import isclass, iscoroutinefunction, isgeneratorfunction
-from multiprocessing import current_process
+from multiprocessing import current_process, get_context
+from multiprocessing.context import BaseContext
 from os.path import basename, splitext
 from threading import current_thread
 
@@ -236,6 +241,7 @@ class Logger:
         backtrace=_defaults.LOGURU_BACKTRACE,
         diagnose=_defaults.LOGURU_DIAGNOSE,
         enqueue=_defaults.LOGURU_ENQUEUE,
+        context=_defaults.LOGURU_CONTEXT,
         catch=_defaults.LOGURU_CATCH,
         **kwargs
     ):
@@ -270,6 +276,10 @@ class Logger:
             Whether the messages to be logged should first pass through a multiprocessing-safe queue
             before reaching the sink. This is useful while logging to a file through multiple
             processes. This also has the advantage of making logging calls non-blocking.
+        context : |multiprocessing.Context| or |str|, optional
+            A context object or name that will be used for all tasks involving internally the
+            |multiprocessing| module, in particular when ``enqueue=True``. If ``None``, the default
+            context is used.
         catch : |bool|, optional
             Whether errors occurring while sink handles logs messages should be automatically
             caught. If ``True``, an exception message is displayed on |sys.stderr| but the exception
@@ -340,7 +350,7 @@ class Logger:
         The ``sink`` handles incoming log messages and proceed to their writing somewhere and
         somehow. A sink can take many forms:
 
-        - A |file-like object|_ like ``sys.stderr`` or ``open("somefile.log", "w")``. Anything with
+        - A |file-like object|_ like ``sys.stderr`` or ``open("file.log", "w")``. Anything with
           a ``.write()`` method is considered as a file-like object. Custom handlers may also
           implement ``flush()`` (called after each logged message), ``stop()`` (called at sink
           termination) and ``complete()`` (awaited by the eponymous method).
@@ -948,6 +958,14 @@ class Logger:
         if not isinstance(encoding, str):
             encoding = "ascii"
 
+        if context is None or isinstance(context, str):
+            context = get_context(context)
+        elif not isinstance(context, BaseContext):
+            raise TypeError(
+                "Invalid context, it should be a string or a multiprocessing context, "
+                "not: '%s'" % type(context).__name__
+            )
+
         with self._core.lock:
             exception_formatter = ExceptionFormatter(
                 colorize=colorize,
@@ -968,6 +986,7 @@ class Logger:
                 colorize=colorize,
                 serialize=serialize,
                 enqueue=enqueue,
+                multiprocessing_context=context,
                 id_=handler_id,
                 error_interceptor=error_interceptor,
                 exception_formatter=exception_formatter,


Test patch:
diff --git a/tests/test_add_option_context.py b/tests/test_add_option_context.py
new file mode 100644
index 0000000..ee6d7c1
--- /dev/null
+++ b/tests/test_add_option_context.py
@@ -0,0 +1,66 @@
+import multiprocessing
+import os
+from unittest.mock import MagicMock
+
+import pytest
+
+from loguru import logger
+
+
+def get_handler_context():
+    # No better way to test correct value than to access the private attribute.
+    handler = next(iter(logger._core.handlers.values()))
+    return handler._multiprocessing_context
+
+
+def test_default_context():
+    logger.add(lambda _: None, context=None)
+    assert get_handler_context() == multiprocessing.get_context(None)
+
+
+@pytest.mark.skipif(os.name == "nt", reason="Windows does not support forking")
+@pytest.mark.parametrize("context_name", ["fork", "forkserver"])
+def test_fork_context_as_string(context_name):
+    logger.add(lambda _: None, context=context_name)
+    assert get_handler_context() == multiprocessing.get_context(context_name)
+
+
+def test_spawn_context_as_string():
+    logger.add(lambda _: None, context="spawn")
+    assert get_handler_context() == multiprocessing.get_context("spawn")
+
+
+@pytest.mark.skipif(os.name == "nt", reason="Windows does not support forking")
+@pytest.mark.parametrize("context_name", ["fork", "forkserver"])
+def test_fork_context_as_object(context_name):
+    context = multiprocessing.get_context(context_name)
+    logger.add(lambda _: None, context=context)
+    assert get_handler_context() == context
+
+
+def test_spawn_context_as_object():
+    context = multiprocessing.get_context("spawn")
+    logger.add(lambda _: None, context=context)
+    assert get_handler_context() == context
+
+
+def test_context_effectively_used():
+    default_context = multiprocessing.get_context()
+    mocked_context = MagicMock(spec=default_context, wraps=default_context)
+    logger.add(lambda _: None, context=mocked_context, enqueue=True)
+    logger.complete()
+    assert mocked_context.Lock.called
+
+
+def test_invalid_context_name():
+    with pytest.raises(ValueError, match=r"cannot find context for"):
+        logger.add(lambda _: None, context="foobar")
+
+
+@pytest.mark.parametrize("context", [42, object()])
+def test_invalid_context_object(context):
+    with pytest.raises(
+        TypeError,
+        match=r"Invalid context, it should be a string or a multiprocessing context",
+    ):
+        logger.add(lambda _: None, context=context)
diff --git a/tests/test_coroutine_sink.py b/tests/test_coroutine_sink.py
index b8fc832..373b7a2 100644
--- a/tests/test_coroutine_sink.py
+++ b/tests/test_coroutine_sink.py
@@ -7,7 +7,6 @@ import threading
 
 import pytest
 
-import loguru
 from loguru import logger
 
 
@@ -590,14 +589,13 @@ class Writer:
 
 
 def test_complete_with_sub_processes(monkeypatch, capsys):
-    ctx = multiprocessing.get_context("spawn")
-    monkeypatch.setattr(loguru._handler, "multiprocessing", ctx)
+    spawn_context = multiprocessing.get_context("spawn")
 
     loop = asyncio.new_event_loop()
     writer = Writer()
-    logger.add(writer.write, format="{message}", enqueue=True, loop=loop)
+    logger.add(writer.write, context=spawn_context, format="{message}", enqueue=True, loop=loop)
 
-    process = ctx.Process(target=subworker, args=[logger])
+    process = spawn_context.Process(target=subworker, args=[logger])
     process.start()
     process.join()
 
diff --git a/tests/test_multiprocessing.py b/tests/test_multiprocessing.py
index cb16466..cd25cd6 100644
--- a/tests/test_multiprocessing.py
+++ b/tests/test_multiprocessing.py
@@ -1,5 +1,6 @@
 import asyncio
 import copy
+import multiprocessing
 import os
 import platform
 import sys
@@ -8,26 +9,17 @@ import time
 
 import pytest
 
-import loguru
 from loguru import logger
 
 
 @pytest.fixture
-def fork_context(monkeypatch):
-    import multiprocessing
-
-    context = multiprocessing.get_context("fork")
-    monkeypatch.setattr(loguru._handler, "multiprocessing", context)
-    yield context
+def fork_context():
+    yield multiprocessing.get_context("fork")
 
 
 @pytest.fixture
-def spawn_context(monkeypatch):
-    import multiprocessing
-
-    context = multiprocessing.get_context("spawn")
-    monkeypatch.setattr(loguru._handler, "multiprocessing", context)
-    yield context
+def spawn_context():
+    yield multiprocessing.get_context("spawn")
 
 
 def do_something(i):
@@ -105,7 +97,7 @@ class Writer:
 def test_apply_spawn(spawn_context):
     writer = Writer()
 
-    logger.add(writer, format="{message}", enqueue=True, catch=False)
+    logger.add(writer, context=spawn_context, format="{message}", enqueue=True, catch=False)
 
     with spawn_context.Pool(1, set_logger, [logger]) as pool:
         for i in range(3):
@@ -123,7 +115,7 @@ def test_apply_spawn(spawn_context):
 def test_apply_fork(fork_context):
     writer = Writer()
 
-    logger.add(writer, format="{message}", enqueue=True, catch=False)
+    logger.add(writer, context=fork_context, format="{message}", enqueue=True, catch=False)
 
     with fork_context.Pool(1, set_logger, [logger]) as pool:
         for i in range(3):
@@ -141,7 +133,7 @@ def test_apply_fork(fork_context):
 def test_apply_inheritance(fork_context):
     writer = Writer()
 
-    logger.add(writer, format="{message}", enqueue=True, catch=False)
+    logger.add(writer, context=fork_context, format="{message}", enqueue=True, catch=False)
 
     with fork_context.Pool(1) as pool:
         for i in range(3):
@@ -158,7 +150,7 @@ def test_apply_inheritance(fork_context):
 def test_apply_async_spawn(spawn_context):
     writer = Writer()
 
-    logger.add(writer, format="{message}", enqueue=True, catch=False)
+    logger.add(writer, context=spawn_context, format="{message}", enqueue=True, catch=False)
 
     with spawn_context.Pool(1, set_logger, [logger]) as pool:
         for i in range(3):
@@ -177,7 +169,7 @@ def test_apply_async_spawn(spawn_context):
 def test_apply_async_fork(fork_context):
     writer = Writer()
 
-    logger.add(writer, format="{message}", enqueue=True, catch=False)
+    logger.add(writer, context=fork_context, format="{message}", enqueue=True, catch=False)
 
     with fork_context.Pool(1, set_logger, [logger]) as pool:
         for i in range(3):
@@ -196,7 +188,7 @@ def test_apply_async_fork(fork_context):
 def test_apply_async_inheritance(fork_context):
     writer = Writer()
 
-    logger.add(writer, format="{message}", enqueue=True, catch=False)
+    logger.add(writer, context=fork_context, format="{message}", enqueue=True, catch=False)
 
     with fork_context.Pool(1) as pool:
         for i in range(3):
@@ -214,7 +206,7 @@ def test_apply_async_inheritance(fork_context):
 def test_process_spawn(spawn_context):
     writer = Writer()
 
-    logger.add(writer, format="{message}", enqueue=True, catch=False)
+    logger.add(writer, context=spawn_context, format="{message}", enqueue=True, catch=False)
 
     process = spawn_context.Process(target=subworker, args=(logger,))
     process.start()
@@ -232,7 +224,7 @@ def test_process_spawn(spawn_context):
 def test_process_fork(fork_context):
     writer = Writer()
 
-    logger.add(writer, format="{message}", enqueue=True, catch=False)
+    logger.add(writer, context=fork_context, format="{message}", enqueue=True, catch=False)
 
     process = fork_context.Process(target=subworker, args=(logger,))
     process.start()
@@ -250,7 +242,7 @@ def test_process_fork(fork_context):
 def test_process_inheritance(fork_context):
     writer = Writer()
 
-    logger.add(writer, format="{message}", enqueue=True, catch=False)
+    logger.add(writer, context=fork_context, format="{message}", enqueue=True, catch=False)
 
     process = fork_context.Process(target=subworker_inheritance)
     process.start()
@@ -267,7 +259,7 @@ def test_process_inheritance(fork_context):
 def test_remove_in_child_process_spawn(spawn_context):
     writer = Writer()
 
-    logger.add(writer, format="{message}", enqueue=True, catch=False)
+    logger.add(writer, context=spawn_context, format="{message}", enqueue=True, catch=False)
 
     process = spawn_context.Process(target=subworker_remove, args=(logger,))
     process.start()
@@ -285,7 +277,7 @@ def test_remove_in_child_process_spawn(spawn_context):
 def test_remove_in_child_process_fork(fork_context):
     writer = Writer()
 
-    logger.add(writer, format="{message}", enqueue=True, catch=False)
+    logger.add(writer, context=fork_context, format="{message}", enqueue=True, catch=False)
 
     process = fork_context.Process(target=subworker_remove, args=(logger,))
     process.start()
@@ -303,7 +295,7 @@ def test_remove_in_child_process_fork(fork_context):
 def test_remove_in_child_process_inheritance(fork_context):
     writer = Writer()
 
-    logger.add(writer, format="{message}", enqueue=True, catch=False)
+    logger.add(writer, context=fork_context, format="{message}", enqueue=True, catch=False)
 
     process = fork_context.Process(target=subworker_remove_inheritance)
     process.start()
@@ -324,7 +316,7 @@ def test_remove_in_main_process_spawn(spawn_context):
     writer = Writer()
     barrier = spawn_context.Barrier(2)
 
-    logger.add(writer, format="{message}", enqueue=True, catch=False)
+    logger.add(writer, context=spawn_context, format="{message}", enqueue=True, catch=False)
 
     process = spawn_context.Process(target=subworker_barrier, args=(logger, barrier))
     process.start()
@@ -343,7 +335,7 @@ def test_remove_in_main_process_fork(fork_context):
     writer = Writer()
     barrier = fork_context.Barrier(2)
 
-    logger.add(writer, format="{message}", enqueue=True, catch=False)
+    logger.add(writer, context=fork_context, format="{message}", enqueue=True, catch=False)
 
     process = fork_context.Process(target=subworker_barrier, args=(logger, barrier))
     process.start()
@@ -362,7 +354,7 @@ def test_remove_in_main_process_inheritance(fork_context):
     writer = Writer()
     barrier = fork_context.Barrier(2)
 
-    logger.add(writer, format="{message}", enqueue=True, catch=False)
+    logger.add(writer, context=fork_context, format="{message}", enqueue=True, catch=False)
 
     process = fork_context.Process(target=subworker_barrier_inheritance, args=(barrier,))
     process.start()
@@ -382,7 +374,9 @@ def test_await_complete_spawn(capsys, spawn_context):
 
     loop = asyncio.new_event_loop()
 
-    logger.add(writer, format="{message}", loop=loop, enqueue=True, catch=False)
+    logger.add(
+        writer, context=spawn_context, format="{message}", loop=loop, enqueue=True, catch=False
+    )
 
     process = spawn_context.Process(target=subworker_complete, args=(logger,))
     process.start()
@@ -407,7 +401,9 @@ def test_await_complete_fork(capsys, fork_context):
 
     loop = asyncio.new_event_loop()
 
-    logger.add(writer, format="{message}", loop=loop, enqueue=True, catch=False)
+    logger.add(
+        writer, context=fork_context, format="{message}", loop=loop, enqueue=True, catch=False
+    )
 
     process = fork_context.Process(target=subworker_complete, args=(logger,))
     process.start()
@@ -432,7 +428,9 @@ def test_await_complete_inheritance(capsys, fork_context):
 
     loop = asyncio.new_event_loop()
 
-    logger.add(writer, format="{message}", loop=loop, enqueue=True, catch=False)
+    logger.add(
+        writer, context=fork_context, format="{message}", loop=loop, enqueue=True, catch=False
+    )
 
     process = fork_context.Process(target=subworker_complete_inheritance)
     process.start()
@@ -455,9 +453,9 @@ def test_not_picklable_sinks_spawn(spawn_context, tmp_path, capsys):
     stream = sys.stderr
     output = []
 
-    logger.add(filepath, format="{message}", enqueue=True, catch=False)
-    logger.add(stream, format="{message}", enqueue=True)
-    logger.add(lambda m: output.append(m), format="{message}", enqueue=True)
+    logger.add(filepath, context=spawn_context, format="{message}", enqueue=True, catch=False)
+    logger.add(stream, context=spawn_context, format="{message}", enqueue=True)
+    logger.add(lambda m: output.append(m), context=spawn_context, format="{message}", enqueue=True)
 
     process = spawn_context.Process(target=subworker, args=[logger])
     process.start()
@@ -482,9 +480,15 @@ def test_not_picklable_sinks_fork(capsys, tmp_path, fork_context):
     stream = sys.stderr
     output = []
 
-    logger.add(filepath, format="{message}", enqueue=True, catch=False)
-    logger.add(stream, format="{message}", enqueue=True, catch=False)
-    logger.add(lambda m: output.append(m), format="{message}", enqueue=True, catch=False)
+    logger.add(filepath, context=fork_context, format="{message}", enqueue=True, catch=False)
+    logger.add(stream, context=fork_context, format="{message}", enqueue=True, catch=False)
+    logger.add(
+        lambda m: output.append(m),
+        context=fork_context,
+        format="{message}",
+        enqueue=True,
+        catch=False,
+    )
 
     process = fork_context.Process(target=subworker, args=[logger])
     process.start()
@@ -509,9 +513,15 @@ def test_not_picklable_sinks_inheritance(capsys, tmp_path, fork_context):
     stream = sys.stderr
     output = []
 
-    logger.add(filepath, format="{message}", enqueue=True, catch=False)
-    logger.add(stream, format="{message}", enqueue=True, catch=False)
-    logger.add(lambda m: output.append(m), format="{message}", enqueue=True, catch=False)
+    logger.add(filepath, context=fork_context, format="{message}", enqueue=True, catch=False)
+    logger.add(stream, context=fork_context, format="{message}", enqueue=True, catch=False)
+    logger.add(
+        lambda m: output.append(m),
+        context=fork_context,
+        format="{message}",
+        enqueue=True,
+        catch=False,
+    )
 
     process = fork_context.Process(target=subworker_inheritance)
     process.start()
@@ -554,7 +564,7 @@ def test_no_deadlock_if_internal_lock_in_use(tmp_path, enqueue, deepcopied, fork
     def worker():
         logger_.info("Child")
 
-    logger_.add(slow_sink, format="{message}", enqueue=enqueue, catch=False)
+    logger_.add(slow_sink, context=fork_context, format="{message}", enqueue=enqueue, catch=False)
 
     thread = threading.Thread(target=main)
     thread.start()
@@ -577,7 +587,7 @@ def test_no_deadlock_if_internal_lock_in_use(tmp_path, enqueue, deepcopied, fork
 @pytest.mark.parametrize("enqueue", [True, False])
 def test_no_deadlock_if_external_lock_in_use(enqueue, capsys, fork_context):
     # Can't reproduce the bug on pytest (even if stderr is not wrapped), but let it anyway
-    logger.add(sys.stderr, enqueue=enqueue, catch=True, format="{message}")
+    logger.add(sys.stderr, context=fork_context, enqueue=enqueue, catch=True, format="{message}")
     num = 100
 
     for i in range(num):
@@ -597,7 +607,7 @@ def test_no_deadlock_if_external_lock_in_use(enqueue, capsys, fork_context):
 @pytest.mark.skipif(os.name == "nt", reason="Windows does not support forking")
 @pytest.mark.skipif(platform.python_implementation() == "PyPy", reason="PyPy is too slow")
 def test_complete_from_multiple_child_processes(capsys, fork_context):
-    logger.add(lambda _: None, enqueue=True, catch=False)
+    logger.add(lambda _: None, context=fork_context, enqueue=True, catch=False)
     num = 100
 
     barrier = fork_context.Barrier(num)