monkeypatching...

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",

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.")

aiida_testing/mock_code/_env_keys.py → aiida_testing/mock_code/_extra_keys.py

@@ -6,14 +6,21 @@
 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'
     EXECUTABLE_PATH = 'AIIDA_MOCK_EXECUTABLE_PATH'
     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'

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)