Add ZeroMQ pytest fixtures to start the broker service inside pytest

src/aiida/brokers/zeromq/broker.py

@@ -11,7 +11,6 @@
 import psutil
 
 from aiida.brokers.broker import Broker, BrokerConfigField, BrokerServiceStatus
-from aiida.common.exceptions import ConfigurationError
 from aiida.common.log import AIIDA_LOGGER
 
 from .communicator import ZeromqCommunicator
@@ -44,8 +43,9 @@ class ZeromqBroker(Broker):
             help='Whether the lifecycle of the broker service is managed by the daemon.',
             default=True,
             param_type='bool',
-            # Running the broker service outside of the daemon is not yet supported, so the setting is not
-            # configurable and always stored with its default.
+            # Regular profiles always let the daemon supervise the service, so the setting is not exposed on the CLI and
+            # is stored with its default. Running the service outside of the daemon is used by the pytest fixtures,
+            # which set this to ``False`` and manage the service lifecycle themselves.
             expose_cli=False,
         ),
     )
@@ -54,16 +54,6 @@ def __init__(self, profile: Profile) -> None:
         super().__init__(profile)
         self._communicator: ZeromqCommunicator | None = None
 
-        # The broker service determines this location, not this client. Currently the service is always managed by
-        # the daemon, so the daemon client is the authority for the directory.
-        if not profile.process_control_config.get('supervised_by_daemon', True):
-            msg = (
-                'The ZeroMQ broker service is not managed by the daemon (`supervised_by_daemon` is false in the broker '
-                'settings), so the location of its state files is unknown. Running the broker service outside of the '
-                'daemon is not yet supported.'
-            )
-            raise ConfigurationError(msg)
-
         from aiida.manage.configuration import get_config
 
         zmq_broker_service_dir = get_config().filepaths(profile)['broker_service']['dir']

src/aiida/tools/pytest_fixtures/__init__.py

@@ -3,6 +3,7 @@
 
 # fmt: off
 
+from .broker import run_aiida_broker_service, run_aiida_broker_service_for_profile
 from .configuration import (
     aiida_config,
     aiida_config_factory,
@@ -50,6 +51,8 @@
     'daemon_client',
     'entry_points',
     'postgres_cluster',
+    'run_aiida_broker_service',
+    'run_aiida_broker_service_for_profile',
     'ssh_key',
     'started_daemon_client',
     'stopped_daemon_client',

src/aiida/tools/pytest_fixtures/broker.py

@@ -0,0 +1,102 @@
+###########################################################################
+# Copyright (c), The AiiDA team. All rights reserved.                     #
+# This file is part of the AiiDA code.                                    #
+#                                                                         #
+# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
+# For further information on the license, see the LICENSE.txt file        #
+# For further information please visit http://www.aiida.net               #
+###########################################################################
+"""Fixtures to provide the message broker required by a test profile."""
+
+from __future__ import annotations
+
+import contextlib
+import shutil
+import signal
+import subprocess
+import sys
+import time
+import typing as t
+
+import pytest
+
+if t.TYPE_CHECKING:
+    from aiida.brokers import ZeromqBroker
+    from aiida.manage.configuration.profile import Profile
+
+
+@pytest.fixture(scope='session')
+def run_aiida_broker_service_for_profile() -> t.Callable[..., t.ContextManager[ZeromqBroker]]:
+    """Return a context manager that runs a ZeroMQ broker service for a given profile.
+
+    This exposes the same mechanism that the :func:`aiida_broker` fixture uses to launch the ZeroMQ service directly,
+    independently of the daemon. It is intended for tests that provide their own profile (e.g. one whose broker service
+    is rooted in a temporary directory) and need its service running for the duration of a block, without going through
+    the session-scoped :func:`aiida_broker` fixture.
+
+    Usage::
+
+        def test(run_aiida_broker_service_for_profile, some_profile):
+            with run_aiida_broker_service_for_profile(some_profile) as broker:
+                ...
+
+    :returns: A context manager that takes a :class:`~aiida.manage.configuration.profile.Profile` (and an optional
+        ``timeout``) and runs the ZeroMQ broker service for the duration of the context, yielding the
+        :class:`~aiida.brokers.zeromq.broker.ZeromqBroker` constructed from that profile.
+    """
+
+    @contextlib.contextmanager
+    def run_broker_service(profile: Profile, timeout: float = 60.0):
+        from aiida.brokers import ZeromqBroker
+
+        broker = ZeromqBroker(profile)
+
+        if broker.check_service_reachable():
+            msg = f'The ZeroMQ broker service is already running at `{broker.service_dir}`.'
+            raise RuntimeError(msg)
+
+        broker.service_dir.mkdir(exist_ok=False)
+
+        process = subprocess.Popen(
+            [sys.executable, '-m', 'aiida.brokers.zeromq.service', '--service-dir', str(broker.service_dir)],
+            start_new_session=True,
+            stdin=subprocess.DEVNULL,
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
+
+        try:
+            start_time = time.monotonic()
+            while not broker.check_service_reachable():
+                if process.poll() is not None:
+                    msg = f'The ZeroMQ broker service exited before becoming reachable with code {process.returncode}.'
+                    raise RuntimeError(msg)
+                if time.monotonic() - start_time > timeout:
+                    msg = f'The ZeroMQ broker service did not become reachable within {timeout} seconds.'
+                    raise TimeoutError(msg)
+                time.sleep(0.1)
+
+            yield broker
+        finally:
+            if process.poll() is None:
+                process.send_signal(signal.SIGINT)
+                try:
+                    process.wait(timeout=timeout)
+                except subprocess.TimeoutExpired:
+                    process.kill()
+                    process.wait()
+
+            shutil.rmtree(broker.service_dir, ignore_errors=True)
+
+    return run_broker_service
+
+
+@pytest.fixture(scope='session', autouse=True)
+def run_aiida_broker_service(aiida_profile: Profile, run_aiida_broker_service_for_profile) -> t.Iterator[None]:
+    """Run the ZeroMQ broker service for the currently loaded test profile for the duration of the session."""
+    if aiida_profile.process_control_backend != 'core.zeromq':
+        yield None
+        return
+
+    with run_aiida_broker_service_for_profile(aiida_profile):
+        yield None

src/aiida/tools/pytest_fixtures/configuration.py

@@ -121,14 +121,20 @@ def factory(
         storage_config = storage_config or {'filepath': str(pathlib.Path(config.dirpath) / name / 'storage')}
 
         if broker_backend and broker_config is None:
-            broker_config = {
-                'broker_protocol': 'amqp',
-                'broker_username': 'guest',
-                'broker_password': 'guest',
-                'broker_host': '127.0.0.1',
-                'broker_port': 5672,
-                'broker_virtual_host': '',
-            }
+            if broker_backend == 'core.rabbitmq':
+                broker_config = {
+                    'broker_protocol': 'amqp',
+                    'broker_username': 'guest',
+                    'broker_password': 'guest',
+                    'broker_host': '127.0.0.1',
+                    'broker_port': 5672,
+                    'broker_virtual_host': '',
+                }
+            elif broker_backend == 'core.zeromq':
+                pass
+            else:
+                msg = f'Unsupported broker backend: {broker_backend}'
+                raise ValueError(msg)
 
         profile = create_profile(
             config,

tests/brokers/test_zeromq_broker.py

@@ -10,6 +10,7 @@
 
 from __future__ import annotations
 
+import copy
 import json
 from unittest.mock import MagicMock, PropertyMock, patch
 
@@ -18,36 +19,47 @@
 
 from aiida.brokers.zeromq.broker import ZeromqBroker, ZeromqIncomingTask
 from aiida.brokers.zeromq.queue import PersistentQueue
-from tests.conftest import _patch_zmq_broker_service_filepaths, _run_zeromq_broker_server
+from aiida.manage.configuration import get_config
 
 
 @pytest.fixture(scope='module')
-def zeromq_broker_with_server(tmp_path_factory):
-    """Create a ZMQ broker instance with a ZMQ server running in the background."""
-    service_dir = tmp_path_factory.mktemp('zeromq-broker')
+def aiida_broker():
+    """Returns the broker for the aiida_profile session fixture with a running service."""
+    from aiida.manage import get_manager
+
+    broker = get_manager().get_broker()
+
+    if not isinstance(broker, ZeromqBroker):
+        pytest.skip('Requires a profile with a ZeroMQ broker.')
+
+    yield broker
+    broker.close()
+
+
+@pytest.fixture
+def zeromq_broker(tmp_path):
+    """Create a ZMQ broker instance rooted in ``tmp_path``."""
     profile = MagicMock()
-    profile.process_control_config = {'supervised_by_daemon': True}
     profile.name = 'test-profile'
-    with _patch_zmq_broker_service_filepaths(profile, service_dir):
-        broker = ZeromqBroker(profile)
-        with _run_zeromq_broker_server(broker):
-            yield broker
 
+    config = get_config()
+    original_filepaths = config.filepaths
 
-def test_get_default_config():
-    """Test that the default broker settings declare the service as managed by the daemon."""
-    assert ZeromqBroker.get_default_config() == {'supervised_by_daemon': True}
+    def filepaths(current_profile):
+        result = copy.deepcopy(original_filepaths(current_profile))
 
+        if current_profile is profile:
+            result['broker_service'] = {'dir': str(tmp_path), 'log': str(tmp_path / 'broker.log')}
 
-def test_not_supervised_by_daemon_raises():
-    """Test that broker construction fails when the service is not managed by the daemon."""
-    from aiida.common.exceptions import ConfigurationError
+        return result
 
-    profile = MagicMock()
-    profile.process_control_config = {'supervised_by_daemon': False}
+    with patch.object(config, 'filepaths', side_effect=filepaths):
+        yield ZeromqBroker(profile)
 
-    with pytest.raises(ConfigurationError, match='not managed by the daemon'):
-        ZeromqBroker(profile)
+
+def test_get_default_config():
+    """Test that the default broker settings declare the service as managed by the daemon."""
+    assert ZeromqBroker.get_default_config() == {'supervised_by_daemon': True}
 
 
 class TestZeromqBrokerStatusQueries:
@@ -64,9 +76,9 @@ def test_probe_service_status_no_file(self, zeromq_broker):
             'error': f'Status file `{zeromq_broker.service_dir / "broker.status"}` does not exist.',
         }
 
-    def test_str_running(self, zeromq_broker_with_server):
+    def test_str_running(self, aiida_broker):
         """Test __str__ when running."""
-        s = str(zeromq_broker_with_server)
+        s = str(aiida_broker)
         assert 'ZeroMQ Broker' in s
         assert 'PID' in s
 
@@ -169,20 +181,16 @@ def test_get_communicator_warns_while_waiting_for_endpoint(self, zeromq_broker,
             communicator_cls.assert_called_once_with(router_endpoint=endpoint, task_timeout=None)
             communicator.start.assert_called_once()
 
-    def test_get_communicator_and_close(self, zeromq_broker_with_server, monkeypatch):
-        """Test get_communicator and close."""
-        try:
-            with patch('aiida.manage.configuration.get_config_option', return_value=None):
-                monkeypatch.setattr('aiida.brokers.zeromq.broker.BROKER_READY_TIMEOUT', 0.5)
-                comm = zeromq_broker_with_server.get_communicator()
-                assert comm is not None
-                assert not comm.is_closed()
-
-                # Second call returns cached instance
-                comm2 = zeromq_broker_with_server.get_communicator()
-                assert comm2 is comm
-        finally:
-            zeromq_broker_with_server.close()
+    def test_get_communicator(self, aiida_broker, monkeypatch):
+        """Test get_communicator."""
+        monkeypatch.setattr('aiida.brokers.zeromq.broker.BROKER_READY_TIMEOUT', 0.5)
+        comm = aiida_broker.get_communicator()
+        assert comm is not None
+        assert not comm.is_closed()
+
+        # Second call returns cached instance
+        comm2 = aiida_broker.get_communicator()
+        assert comm2 is comm
 
     def test_get_communicator_timeout(self, zeromq_broker, monkeypatch):
         """Test get_communicator raises on timeout when zeromq_broker not running."""
@@ -237,10 +245,10 @@ def test_processing_context_manager(self, tmp_path):
 class TestZeromqBrokerIntegration:
     """Integration tests for the ZeroMQ broker with AiiDA."""
 
-    def test_broker_lifecycle(self, zeromq_broker_with_server):
+    def test_broker_lifecycle(self, aiida_broker):
         """Test the zeromq_broker lifecycle."""
-        assert zeromq_broker_with_server.check_service_reachable()
+        assert aiida_broker.check_service_reachable()
 
-        status = zeromq_broker_with_server.probe_service_status()
+        status = aiida_broker.probe_service_status()
         assert status['connected'] is True
         assert 'pid' in status