diff --git a/src/aiida/brokers/zeromq/broker.py b/src/aiida/brokers/zeromq/broker.py index b72d3bbdf7..25c97d8bd1 100644 --- a/src/aiida/brokers/zeromq/broker.py +++ b/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'] diff --git a/src/aiida/tools/pytest_fixtures/__init__.py b/src/aiida/tools/pytest_fixtures/__init__.py index e19d4c455e..2cd6a6828f 100644 --- a/src/aiida/tools/pytest_fixtures/__init__.py +++ b/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', diff --git a/src/aiida/tools/pytest_fixtures/broker.py b/src/aiida/tools/pytest_fixtures/broker.py new file mode 100644 index 0000000000..9ad0648c3e --- /dev/null +++ b/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 diff --git a/src/aiida/tools/pytest_fixtures/configuration.py b/src/aiida/tools/pytest_fixtures/configuration.py index 0c716eb79c..f1b13d2538 100644 --- a/src/aiida/tools/pytest_fixtures/configuration.py +++ b/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, diff --git a/tests/brokers/test_zeromq_broker.py b/tests/brokers/test_zeromq_broker.py index 85604cce9d..967eae3bdd 100644 --- a/tests/brokers/test_zeromq_broker.py +++ b/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 diff --git a/tests/brokers/test_zeromq_communicator.py b/tests/brokers/test_zeromq_communicator.py index 171c0838fd..bf684062ea 100644 --- a/tests/brokers/test_zeromq_communicator.py +++ b/tests/brokers/test_zeromq_communicator.py @@ -13,14 +13,26 @@ import time from concurrent.futures import Future from pathlib import Path -from unittest.mock import MagicMock import kiwipy import pytest from aiida.brokers.zeromq.broker import ZeromqBroker from aiida.brokers.zeromq.communicator import ZeromqCommunicator -from tests.conftest import _patch_zmq_broker_service_filepaths, _run_zeromq_broker_server + + +@pytest.fixture(scope='module') +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() def get_router_endpoint(broker: ZeromqBroker) -> str: @@ -29,26 +41,12 @@ def get_router_endpoint(broker: ZeromqBroker) -> str: return f'ipc://{sockets_path}/router.sock' -@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') - 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 - - class TestZeromqCommunicatorLifecycle: """Tests for communicator initialization and lifecycle.""" - def test_init(self, zeromq_broker_with_server): + def test_init(self, aiida_broker): """Test communicator initialization.""" - communicator = ZeromqCommunicator(router_endpoint=get_router_endpoint(zeromq_broker_with_server)) + communicator = ZeromqCommunicator(router_endpoint=get_router_endpoint(aiida_broker)) communicator.start() try: @@ -58,16 +56,16 @@ def test_init(self, zeromq_broker_with_server): assert communicator.is_closed() is True - def test_context_manager(self, zeromq_broker_with_server): + def test_context_manager(self, aiida_broker): """Test communicator as context manager.""" - with ZeromqCommunicator(router_endpoint=get_router_endpoint(zeromq_broker_with_server)) as communicator: + with ZeromqCommunicator(router_endpoint=get_router_endpoint(aiida_broker)) as communicator: assert communicator.is_closed() is False assert communicator.is_closed() is True - def test_close_idempotent(self, zeromq_broker_with_server): + def test_close_idempotent(self, aiida_broker): """Test close is idempotent.""" - comm = ZeromqCommunicator(router_endpoint=get_router_endpoint(zeromq_broker_with_server)) + comm = ZeromqCommunicator(router_endpoint=get_router_endpoint(aiida_broker)) comm.start() comm.close() comm.close() # should not raise @@ -84,8 +82,8 @@ class TestZeromqCommunicatorMessaging: """Tests for communicator messaging operations with a real zeromq_broker.""" @pytest.fixture - def zeromq_comm(self, zeromq_broker_with_server): - comm = ZeromqCommunicator(router_endpoint=get_router_endpoint(zeromq_broker_with_server)) + def zeromq_comm(self, aiida_broker): + comm = ZeromqCommunicator(router_endpoint=get_router_endpoint(aiida_broker)) comm.start() yield comm @@ -141,11 +139,11 @@ class TestZeromqCommunicatorRoundTrip: """Integration tests for full task, RPC, and broadcast round-trips.""" @pytest.fixture - def sender_and_worker(self, zeromq_broker_with_server): - sender = ZeromqCommunicator(router_endpoint=get_router_endpoint(zeromq_broker_with_server), client_id='sender') + def sender_and_worker(self, aiida_broker): + sender = ZeromqCommunicator(router_endpoint=get_router_endpoint(aiida_broker), client_id='sender') sender.start() - worker = ZeromqCommunicator(router_endpoint=get_router_endpoint(zeromq_broker_with_server), client_id='worker') + worker = ZeromqCommunicator(router_endpoint=get_router_endpoint(aiida_broker), client_id='worker') worker.start() yield sender, worker diff --git a/tests/conftest.py b/tests/conftest.py index d5bf761ec7..fcbf3879e2 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -19,24 +19,18 @@ import logging import os import pathlib -import shutil -import signal import subprocess import sys -import time import types import typing as t import warnings -from contextlib import contextmanager from enum import Enum from pathlib import Path -from unittest.mock import MagicMock, patch import click import pytest from aiida import get_profile, orm -from aiida.brokers import ZeromqBroker from aiida.common.folders import Folder from aiida.common.links import LinkType from aiida.manage import get_manager @@ -71,78 +65,6 @@ class TestBrokerBackend(Enum): NONE = 'none' -@contextmanager -def _patch_zmq_broker_service_filepaths(profile, service_dir: Path): - """Patch only the ZeroMQ broker-service filepaths for a test profile.""" - config = get_config() - original_filepaths = config.filepaths - - def filepaths(current_profile): - result = copy.deepcopy(original_filepaths(current_profile)) - - if current_profile is profile: - result['broker_service'] = {'dir': str(service_dir), 'log': str(service_dir / 'broker.log')} - - return result - - with patch.object(config, 'filepaths', side_effect=filepaths): - yield - - -@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, tmp_path): - yield ZeromqBroker(profile) - - -@contextmanager -def _run_zeromq_broker_server(zeromq_broker: ZeromqBroker, timeout: float = 10.0): - """Run a ZeroMQ broker service subprocess for the duration of a context.""" - if zeromq_broker.check_service_reachable(): - raise ValueError('Broker server already running') - - zeromq_broker.service_dir.mkdir(parents=True, exist_ok=True) - - process = subprocess.Popen( - [sys.executable, '-m', 'aiida.brokers.zeromq.service', '--service-dir', str(zeromq_broker.service_dir)], - start_new_session=True, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - stdin=subprocess.DEVNULL, - ) - - try: - start_time = time.time() - while time.time() - start_time < timeout: - if process.poll() is not None: - msg = f'ZeroMQ broker exited before becoming ready with code {process.returncode}' - raise RuntimeError(msg) - if zeromq_broker.check_service_reachable(): - break - time.sleep(0.1) - else: - msg = f'ZeroMQ broker did not become ready within {timeout}s' - raise TimeoutError(msg) - - yield process - finally: - try: - if process.poll() is None: - process.send_signal(signal.SIGINT) - try: - process.wait(timeout=timeout) - except subprocess.TimeoutExpired: - process.kill() - process.wait() - finally: - shutil.rmtree(zeromq_broker.service_dir, ignore_errors=True) - - def pytest_collection_modifyitems(items, config): """Automatically generate markers for certain tests. @@ -335,13 +257,7 @@ def aiida_profile(pytestconfig, aiida_config, aiida_profile_factory, config_psql with aiida_profile_factory( aiida_config, storage_backend=storage, storage_config=config, broker_backend=broker ) as profile: - # Start ZeroMQ broker service if needed (tests don't use circus) - broker_instance = get_manager().get_broker() - if isinstance(broker_instance, ZeromqBroker): - with _run_zeromq_broker_server(broker_instance): - yield profile - else: - yield profile + yield profile @pytest.fixture() @@ -926,7 +842,6 @@ def factory( def run_cli_command_subprocess(command, parameters, user_input, profile_name, suppress_warnings): """Run CLI command through ``subprocess``.""" import subprocess - import sys env = os.environ.copy() command_path = cli_command_map()[command] diff --git a/tests/tools/pytest_fixtures/test_configuration.py b/tests/tools/pytest_fixtures/test_configuration.py index aabd950d9b..77deb48174 100644 --- a/tests/tools/pytest_fixtures/test_configuration.py +++ b/tests/tools/pytest_fixtures/test_configuration.py @@ -2,6 +2,8 @@ from pathlib import Path +import pytest + from aiida.manage.configuration import get_config, load_config from aiida.manage.configuration.settings import DEFAULT_CONFIG_FILE_NAME @@ -55,3 +57,10 @@ def test_aiida_profile_tmp(aiida_profile, aiida_profile_tmp): assert isinstance(aiida_profile_tmp, Profile) assert aiida_profile_tmp.is_test_profile assert aiida_profile_tmp.uuid != aiida_profile.uuid + + +def test_aiida_profile_factory_unsupported_broker(aiida_config_tmp, aiida_profile_factory): + """Test that ``aiida_profile_factory`` raises for a broker backend without a default configuration.""" + with pytest.raises(ValueError, match='Unsupported broker backend: core.unsupported'): + with aiida_profile_factory(aiida_config_tmp, broker_backend='core.unsupported'): + pass