diff --git a/tests/mock_code/data/mock-diff-4b5ca93b23993979a24f795147487f1c/_scheduler-stderr.txt b/_scheduler-stderr.txt similarity index 100% rename from tests/mock_code/data/mock-diff-4b5ca93b23993979a24f795147487f1c/_scheduler-stderr.txt rename to _scheduler-stderr.txt diff --git a/tests/mock_code/data/mock-diff-4b5ca93b23993979a24f795147487f1c/_scheduler-stdout.txt b/_scheduler-stdout.txt similarity index 100% rename from tests/mock_code/data/mock-diff-4b5ca93b23993979a24f795147487f1c/_scheduler-stdout.txt rename to _scheduler-stdout.txt diff --git a/aiida_testing/mock_code/__init__.py b/aiida_testing/mock_code/__init__.py index 799bf14..fd369f9 100644 --- a/aiida_testing/mock_code/__init__.py +++ b/aiida_testing/mock_code/__init__.py @@ -6,7 +6,8 @@ from ._fixtures import * -# Note: This is necessary for the sphinx doc - otherwise it does not find aiida_testing.mock_code.mock_code_factory +# Note: This is necessary for the sphinx doc - otherwise it does not find +# aiida_testing.mock_code.mock_code_factory __all__ = ( "pytest_addoption", "testing_config_action", diff --git a/aiida_testing/mock_code/_cli.py b/aiida_testing/mock_code/_cli.py index c91f780..51fb1cb 100644 --- a/aiida_testing/mock_code/_cli.py +++ b/aiida_testing/mock_code/_cli.py @@ -1,166 +1,12 @@ -#!/usr/bin/env python # -*- coding: utf-8 -*- -""" -Implements the executable for running a mock AiiDA code. +"""Defines the fallback executable that is used in a Code when no other +executable is configured. """ -import os import sys -import shutil -import hashlib -import subprocess -import typing as ty -import fnmatch -from pathlib import Path - -from ._env_keys import EnvKeys - -SUBMIT_FILE = '_aiidasubmit.sh' - - -def run() -> None: - """ - Run the mock AiiDA code. If the corresponding result exists, it is - simply copied over to the current working directory. Otherwise, - the code will replace the executable in the aiidasubmit file, - launch the "real" code, and then copy the results into the data - directory. - """ - # Get environment variables - label = os.environ[EnvKeys.LABEL.value] - data_dir = os.environ[EnvKeys.DATA_DIR.value] - executable_path = os.environ[EnvKeys.EXECUTABLE_PATH.value] - ignore_files = os.environ[EnvKeys.IGNORE_FILES.value].split(':') - ignore_paths = os.environ[EnvKeys.IGNORE_PATHS.value].split(':') - regenerate_data = os.environ[EnvKeys.REGENERATE_DATA.value] == 'True' - - hash_digest = get_hash().hexdigest() - - res_dir = Path(data_dir) / f"mock-{label}-{hash_digest}" - - if regenerate_data and res_dir.exists(): - shutil.rmtree(res_dir) - - if not res_dir.exists(): - if not executable_path: - sys.exit("No existing output, and no executable specified.") - - # replace executable path in submit file and run calculation - replace_submit_file(executable_path=executable_path) - subprocess.call(['bash', SUBMIT_FILE]) - - # back up results to data directory - os.makedirs(res_dir) - copy_files( - src_dir=Path('.'), - dest_dir=res_dir, - ignore_files=ignore_files, - ignore_paths=ignore_paths - ) - - else: - # copy outputs from data directory to working directory - for path in res_dir.iterdir(): - if path.is_dir(): - shutil.rmtree(path.name, ignore_errors=True) - shutil.copytree(path, path.name) - elif path.is_file(): - shutil.copyfile(path, path.name) - else: - sys.exit(f"Can not copy '{path.name}'.") - - -def get_hash() -> 'hashlib._Hash': - """ - Get the MD5 hash for the current working directory. - """ - md5sum = hashlib.md5() - # Here the order needs to be consistent, thus globbing - # with 'sorted'. - for path in sorted(Path('.').glob('**/*')): - if path.is_file() and not path.match('.aiida/**'): - with open(path, 'rb') as file_obj: - file_content_bytes = file_obj.read() - if path.name == SUBMIT_FILE: - file_content_bytes = strip_submit_content(file_content_bytes) - md5sum.update(path.name.encode()) - md5sum.update(file_content_bytes) - return md5sum - -def strip_submit_content(aiidasubmit_content_bytes: bytes) -> bytes: - """ - Helper function to strip content which changes between - test runs from the aiidasubmit file. - """ - aiidasubmit_content = aiidasubmit_content_bytes.decode() - lines: ty.Iterable[str] = aiidasubmit_content.splitlines() - # Strip lines containing the aiida_testing.mock_code environment variables. - lines = (line for line in lines if 'export AIIDA_MOCK' not in line) - # Remove abspath of the aiida-mock-code, but keep cmdline - # arguments. - lines = (line.split("aiida-mock-code'")[-1] for line in lines) - return '\n'.join(lines).encode() - - -def replace_submit_file(executable_path: str) -> None: - """ - Replace the executable specified in the AiiDA submit file, and - strip the AIIDA_MOCK environment variables. - """ - with open(SUBMIT_FILE, 'r') as submit_file: - submit_file_content = submit_file.read() - - submit_file_res_lines = [] - for line in submit_file_content.splitlines(): - if 'export AIIDA_MOCK' in line: - continue - if 'aiida-mock-code' in line: - submit_file_res_lines.append( - f"'{executable_path}' " + line.split("aiida-mock-code'")[1] - ) - else: - submit_file_res_lines.append(line) - with open(SUBMIT_FILE, 'w') as submit_file: - submit_file.write('\n'.join(submit_file_res_lines)) - - -def copy_files( - src_dir: Path, dest_dir: Path, ignore_files: ty.Iterable[str], ignore_paths: ty.Iterable[str] -) -> None: - """Copy files from source to destination directory while ignoring certain files/folders. - - :param src_dir: Source directory - :param dest_dir: Destination directory - :param ignore_files: A list of file names (UNIX shell style patterns allowed) which are not copied to the - destination. - :param ignore_paths: A list of paths (UNIX shell style patterns allowed) which are not copied to the destination. +def run(): + """Dummy executable that is passed to a Code when no config is set. """ - exclude_paths: ty.Set = {filepath for path in ignore_paths for filepath in src_dir.glob(path)} - exclude_files = {path.relative_to(src_dir) for path in exclude_paths if path.is_file()} - exclude_dirs = {path.relative_to(src_dir) for path in exclude_paths if path.is_dir()} - - # Here we rely on getting the directory name before - # accessing its content, hence using os.walk. - for dirpath, _, filenames in os.walk(src_dir): - relative_dir = Path(dirpath).relative_to(src_dir) - dirs_to_check = list(relative_dir.parents) + [relative_dir] - - if relative_dir.parts and relative_dir.parts[0] == ('.aiida'): - continue - - if any(exclude_dir in dirs_to_check for exclude_dir in exclude_dirs): - continue - - for filename in filenames: - if any(fnmatch.fnmatch(filename, expr) for expr in ignore_files): - continue - - if relative_dir / filename in exclude_files: - continue - - os.makedirs(dest_dir / relative_dir, exist_ok=True) - - relative_file_path = relative_dir / filename - shutil.copyfile(src_dir / relative_file_path, dest_dir / relative_file_path) + sys.exit("No executable specified in the aiida-testing config, and no existing result found.") diff --git a/aiida_testing/mock_code/_env_keys.py b/aiida_testing/mock_code/_extra_keys.py similarity index 56% rename from aiida_testing/mock_code/_env_keys.py rename to aiida_testing/mock_code/_extra_keys.py index 2c7d84c..9eb54ce 100644 --- a/aiida_testing/mock_code/_env_keys.py +++ b/aiida_testing/mock_code/_extra_keys.py @@ -6,10 +6,9 @@ from enum import Enum -class EnvKeys(Enum): +class CodeExtraKeys(Enum): """ - An enum containing the environment variables defined for - the mock code execution. + An enum containing the keys to be used in the `Code` extras. """ LABEL = 'AIIDA_MOCK_LABEL' DATA_DIR = 'AIIDA_MOCK_DATA_DIR' @@ -17,3 +16,11 @@ class EnvKeys(Enum): IGNORE_FILES = 'AIIDA_MOCK_IGNORE_FILES' IGNORE_PATHS = 'AIIDA_MOCK_IGNORE_PATHS' REGENERATE_DATA = 'AIIDA_MOCK_REGENERATE_DATA' + + +class CalculationExtraKeys(Enum): + """ + An enum containing the keys to be used in the `Calculation` extras. + """ + NEEDS_COPY_TO_RES_DIR = 'AIIDA_MOCK_NEEDS_COPY_TO_RES_DIR' + RES_DIR = 'AIIDA_MOCK_RES_DIR' diff --git a/aiida_testing/mock_code/_fixtures.py b/aiida_testing/mock_code/_fixtures.py index 0601bf7..696a84c 100644 --- a/aiida_testing/mock_code/_fixtures.py +++ b/aiida_testing/mock_code/_fixtures.py @@ -3,9 +3,10 @@ Defines a pytest fixture for creating mock AiiDA codes. """ +import os +import sys import uuid import shutil -import inspect import pathlib import typing as ty import warnings @@ -14,17 +15,16 @@ import click import pytest +from aiida.engine.daemon import execmanager from aiida.orm import Code -from ._env_keys import EnvKeys +from ._extra_keys import CodeExtraKeys, CalculationExtraKeys +from ._helpers import get_hash, copy_files from .._config import Config, CONFIG_FILE_NAME, ConfigActions __all__ = ( - "pytest_addoption", - "testing_config_action", - "mock_regenerate_test_data", - "testing_config", - "mock_code_factory", + "pytest_addoption", "testing_config_action", "mock_regenerate_test_data", "testing_config", + "mock_code_factory", "patch_calculation_execution" ) @@ -77,8 +77,9 @@ def testing_config(testing_config_action): # pylint: disable=redefined-outer-na @pytest.fixture(scope='function') def mock_code_factory( - aiida_localhost, testing_config, testing_config_action, mock_regenerate_test_data -): # pylint: disable=redefined-outer-name + aiida_localhost, testing_config, testing_config_action, mock_regenerate_test_data, + patch_calculation_execution +): # pylint: disable=redefined-outer-name, unused-argument """ Fixture to create a mock AiiDA Code. @@ -87,7 +88,7 @@ def mock_code_factory( """ - def _get_mock_code( + def _get_mock_code( # pylint: disable=too-many-arguments label: str, entry_point: str, data_dir_abspath: ty.Union[str, pathlib.Path], @@ -97,7 +98,7 @@ def _get_mock_code( _config: dict = testing_config, _config_action: str = testing_config_action, _regenerate_test_data: bool = mock_regenerate_test_data, - ): # pylint: disable=too-many-arguments + ) -> Code: """ Creates a mock AiiDA code. If the same inputs have been run previously, the results are copied over from the corresponding sub-directory of @@ -172,25 +173,107 @@ def _get_mock_code( if _config_action == ConfigActions.GENERATE.value: mock_code_config[label] = code_executable_path + if code_executable_path in {'TO_SPECIFY', 'NOT_FOUND'}: + remote_executable_path = mock_executable_path + else: + remote_executable_path = code_executable_path + code = Code( input_plugin_name=entry_point, - remote_computer_exec=[aiida_localhost, mock_executable_path] + remote_computer_exec=[aiida_localhost, remote_executable_path] ) code.label = code_label - code.set_prepend_text( - inspect.cleandoc( - f""" - export {EnvKeys.LABEL.value}="{label}" - export {EnvKeys.DATA_DIR.value}="{data_dir_abspath}" - export {EnvKeys.EXECUTABLE_PATH.value}="{code_executable_path}" - export {EnvKeys.IGNORE_FILES.value}="{':'.join(ignore_files)}" - export {EnvKeys.IGNORE_PATHS.value}="{':'.join(ignore_paths)}" - export {EnvKeys.REGENERATE_DATA.value}={'True' if _regenerate_test_data else 'False'} - """ - ) - ) code.store() + + code.set_extra(CodeExtraKeys.LABEL.value, label) + code.set_extra(CodeExtraKeys.DATA_DIR.value, str(data_dir_abspath)) + code.set_extra(CodeExtraKeys.EXECUTABLE_PATH.value, str(code_executable_path)) + code.set_extra(CodeExtraKeys.IGNORE_FILES.value, ignore_files) + code.set_extra(CodeExtraKeys.IGNORE_PATHS.value, ignore_paths) + code.set_extra(CodeExtraKeys.REGENERATE_DATA.value, _regenerate_test_data) + return code return _get_mock_code + + +@pytest.fixture(scope='function', autouse=True) +def patch_calculation_execution(monkeypatch): + """Patch execmanager.submit_calculation such as to take data from test data directory. + """ + + unpatched_submit_calculation = execmanager.submit_calculation + unpatched_retrieve_calculation = execmanager.retrieve_calculation + + def mock_submit_calculation(calculation, transport): + """ + Run the mock AiiDA code. If the corresponding result exists, it is + simply copied over to the current working directory. Otherwise, + the code will replace the executable in the aiidasubmit file, + launch the "real" code, and then copy the results into the data + directory. + :param calculation: + :param transport: + :return: + """ + code = calculation.inputs.code + label = code.get_extra(CodeExtraKeys.LABEL.value) + data_dir = code.get_extra(CodeExtraKeys.DATA_DIR.value) + executable_path = code.get_extra(CodeExtraKeys.EXECUTABLE_PATH.value) + + regenerate_data = code.get_extra(CodeExtraKeys.REGENERATE_DATA.value) + + workdir = pathlib.Path(calculation.get_remote_workdir()) + hash_digest = get_hash(workdir, code=code).hexdigest() + res_dir = pathlib.Path(data_dir) / f"mock-{label}-{hash_digest}" + + calculation.set_extra(CalculationExtraKeys.RES_DIR.value, str(res_dir.absolute())) + + if regenerate_data and res_dir.exists(): + shutil.rmtree(res_dir) + + if not res_dir.exists(): + if not executable_path: + sys.exit("No existing output, and no executable specified.") + + calculation.set_extra(CalculationExtraKeys.NEEDS_COPY_TO_RES_DIR.value, True) + res_jobid = unpatched_submit_calculation(calculation, transport) + + else: + # copy outputs from data directory to working directory + for path in res_dir.iterdir(): + out_path = workdir / path.name + if path.is_dir(): + shutil.rmtree(out_path, ignore_errors=True) + shutil.copytree(path, out_path) + elif path.is_file(): + shutil.copyfile(path, out_path) + else: + sys.exit(f"Can not copy '{path.name}'.") + + # return a non-existing jobid + res_jobid = -1 + return res_jobid + + def mock_retrieve_calculation(calculation, transport, retrieved_temporary_folder): + # back up results to data directory + if calculation.get_extra(CalculationExtraKeys.NEEDS_COPY_TO_RES_DIR.value, False): + code = calculation.inputs.code + + ignore_files = code.get_extra(CodeExtraKeys.IGNORE_FILES.value) + ignore_paths = code.get_extra(CodeExtraKeys.IGNORE_PATHS.value) + + res_dir = calculation.get_extra(CalculationExtraKeys.RES_DIR.value) + os.makedirs(res_dir) + copy_files( + src_dir=pathlib.Path(calculation.get_remote_workdir()), + dest_dir=res_dir, + ignore_files=ignore_files, + ignore_paths=ignore_paths + ) + + unpatched_retrieve_calculation(calculation, transport, retrieved_temporary_folder) + + monkeypatch.setattr(execmanager, 'submit_calculation', mock_submit_calculation) + monkeypatch.setattr(execmanager, 'retrieve_calculation', mock_retrieve_calculation) diff --git a/aiida_testing/mock_code/_helpers.py b/aiida_testing/mock_code/_helpers.py new file mode 100644 index 0000000..5bc40b8 --- /dev/null +++ b/aiida_testing/mock_code/_helpers.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +""" +Implements the executable for running a mock AiiDA code. +""" + +import os +import shutil +import pathlib +import hashlib +import typing as ty +import fnmatch +from pathlib import Path + +from aiida import orm + +SUBMIT_FILE = '_aiidasubmit.sh' + + +def get_hash(dirpath: pathlib.Path, code: orm.Code) -> 'hashlib._Hash': + """ + Get the MD5 hash for the current working directory. + """ + md5sum = hashlib.md5() + # Here the order needs to be consistent, thus globbing + # with 'sorted'. + for path in sorted(dirpath.glob('**/*')): + if path.is_file() and not path.match('.aiida/**'): + with open(path, 'rb') as file_obj: + file_content_bytes = file_obj.read() + if path.name == SUBMIT_FILE: + file_content_bytes = strip_submit_content(file_content_bytes, code=code) + md5sum.update(path.name.encode()) + md5sum.update(file_content_bytes) + + return md5sum + + +def strip_submit_content(aiidasubmit_content_bytes: bytes, code: orm.Code) -> bytes: + """ + Helper function to strip content which changes between + test runs from the aiidasubmit file. + """ + aiidasubmit_content = aiidasubmit_content_bytes.decode() + replaced_content = aiidasubmit_content.replace(f"'{code.get_remote_exec_path()}'", '') + # code.label) + return replaced_content.encode() + + +def copy_files( + src_dir: Path, dest_dir: Path, ignore_files: ty.Iterable[str], ignore_paths: ty.Iterable[str] +) -> None: + """Copy files from source to destination directory while ignoring certain files/folders. + + :param src_dir: Source directory + :param dest_dir: Destination directory + :param ignore_files: A list of file names (UNIX shell style patterns allowed) which are not copied to the + destination. + :param ignore_paths: A list of paths (UNIX shell style patterns allowed) which are not copied to the destination. + """ + exclude_paths: ty.Set = {filepath for path in ignore_paths for filepath in src_dir.glob(path)} + exclude_files = {path.relative_to(src_dir) for path in exclude_paths if path.is_file()} + exclude_dirs = {path.relative_to(src_dir) for path in exclude_paths if path.is_dir()} + + # Here we rely on getting the directory name before + # accessing its content, hence using os.walk. + for dirpath, _, filenames in os.walk(src_dir): + relative_dir = Path(dirpath).relative_to(src_dir) + dirs_to_check = list(relative_dir.parents) + [relative_dir] + + if relative_dir.parts and relative_dir.parts[0] == ('.aiida'): + continue + + if any(exclude_dir in dirs_to_check for exclude_dir in exclude_dirs): + continue + + for filename in filenames: + if any(fnmatch.fnmatch(filename, expr) for expr in ignore_files): + continue + + if relative_dir / filename in exclude_files: + continue + + os.makedirs(dest_dir / relative_dir, exist_ok=True) + + relative_file_path = relative_dir / filename + shutil.copyfile(src_dir / relative_file_path, dest_dir / relative_file_path) diff --git a/mypy.ini b/mypy.ini index 6686218..a60c7c9 100644 --- a/mypy.ini +++ b/mypy.ini @@ -4,9 +4,11 @@ python_version = 3.6 ; Strictness settings -disallow_any_unimported = True +disallow_any_unimported = False disallow_subclassing_any = True +# disallow_untyped_defs = True +disallow_incomplete_defs = True disallow_untyped_calls = True disallow_untyped_decorators = True diff --git a/setup.cfg b/setup.cfg index e22285e..8cc5fcc 100644 --- a/setup.cfg +++ b/setup.cfg @@ -32,6 +32,7 @@ include_package_data = true install_requires = aiida-core>=1.0.0<2.0.0 pytest>=3.6 + pytest-mock pyyaml~=5.1.2 voluptuous~=0.11.7 packages = find: @@ -50,8 +51,6 @@ pre_commit = pylint==2.6.0 mypy==0.740 - [options.entry_points] console_scripts = aiida-mock-code = aiida_testing.mock_code._cli:run - diff --git a/tests/conftest.py b/tests/conftest.py index 520f4ec..b69ab6d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,5 +2,4 @@ """ Configuration file for pytest tests of aiida-testing. """ - pytest_plugins = ['aiida.manage.tests.pytest_fixtures', 'aiida_testing.mock_code'] # pylint: disable=invalid-name diff --git a/tests/mock_code/data/mock-diff-broken-4b5ca93b23993979a24f795147487f1c/_scheduler-stderr.txt b/tests/mock_code/data/mock-diff-16fd960a5a26cfc51d24655b2db7834f/_scheduler-stderr.txt similarity index 100% rename from tests/mock_code/data/mock-diff-broken-4b5ca93b23993979a24f795147487f1c/_scheduler-stderr.txt rename to tests/mock_code/data/mock-diff-16fd960a5a26cfc51d24655b2db7834f/_scheduler-stderr.txt diff --git a/tests/mock_code/data/mock-diff-broken-4b5ca93b23993979a24f795147487f1c/_scheduler-stdout.txt b/tests/mock_code/data/mock-diff-16fd960a5a26cfc51d24655b2db7834f/_scheduler-stdout.txt similarity index 100% rename from tests/mock_code/data/mock-diff-broken-4b5ca93b23993979a24f795147487f1c/_scheduler-stdout.txt rename to tests/mock_code/data/mock-diff-16fd960a5a26cfc51d24655b2db7834f/_scheduler-stdout.txt diff --git a/tests/mock_code/data/mock-diff-4b5ca93b23993979a24f795147487f1c/patch.diff b/tests/mock_code/data/mock-diff-16fd960a5a26cfc51d24655b2db7834f/patch.diff similarity index 100% rename from tests/mock_code/data/mock-diff-4b5ca93b23993979a24f795147487f1c/patch.diff rename to tests/mock_code/data/mock-diff-16fd960a5a26cfc51d24655b2db7834f/patch.diff diff --git a/tests/mock_code/data/mock-diff-broken-4b5ca93b23993979a24f795147487f1c/patch.diff b/tests/mock_code/data/mock-diff-broken-16fd960a5a26cfc51d24655b2db7834f/patch.diff similarity index 100% rename from tests/mock_code/data/mock-diff-broken-4b5ca93b23993979a24f795147487f1c/patch.diff rename to tests/mock_code/data/mock-diff-broken-16fd960a5a26cfc51d24655b2db7834f/patch.diff diff --git a/tests/mock_code/test_diff.py b/tests/mock_code/test_diff.py index 472cedd..1484742 100644 --- a/tests/mock_code/test_diff.py +++ b/tests/mock_code/test_diff.py @@ -160,7 +160,7 @@ def test_regenerate_test_data(mock_code_factory, generate_diff_inputs, datadir): label='diff', data_dir_abspath=TEST_DATA_DIR, entry_point=CALC_ENTRY_POINT, - ignore_paths=('_aiidasubmit.sh', ), + ignore_paths=('_aiidasubmit.sh', 'file?.txt'), _regenerate_test_data=True, ) diff --git a/tests/mock_code/test_ignore_paths.py b/tests/mock_code/test_ignore_paths.py index afb7266..d6e4d52 100644 --- a/tests/mock_code/test_ignore_paths.py +++ b/tests/mock_code/test_ignore_paths.py @@ -6,7 +6,7 @@ from pathlib import Path import pytest -from aiida_testing.mock_code._cli import copy_files +from aiida_testing.mock_code._helpers import copy_files OUTPUT_PATHS = ( Path('file1.txt'),