dirty commit

all tests but one already pass

let's see whether we can avoid copy-pasting the function next

aiida_testing/mock_code/_cli.py

@@ -26,48 +26,7 @@ def run() -> None:
     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}'.")
+    pass
 
 
 def get_hash() -> 'hashlib._Hash':
@@ -104,12 +63,12 @@ def strip_submit_content(aiidasubmit_content_bytes: bytes) -> bytes:
     return '\n'.join(lines).encode()
 
 
-def replace_submit_file(executable_path: str) -> None:
+def replace_submit_file(executable_path: str, working_directory='.') -> 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:
+    with open(Path(working_directory) / SUBMIT_FILE, 'r') as submit_file:
         submit_file_content = submit_file.read()
 
     submit_file_res_lines = []
@@ -122,7 +81,7 @@ def replace_submit_file(executable_path: str) -> None:
             )
         else:
             submit_file_res_lines.append(line)
-    with open(SUBMIT_FILE, 'w') as submit_file:
+    with open(Path(working_directory) / SUBMIT_FILE, 'w') as submit_file:
         submit_file.write('\n'.join(submit_file_res_lines))
 
 

aiida_testing/mock_code/_fixtures.py

@@ -20,11 +20,8 @@
 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_submission"
 )
 
 
@@ -191,6 +188,113 @@ def _get_mock_code(
         )
 
         code.store()
+
+        code.set_extra(EnvKeys.LABEL.value, label)
+        code.set_extra(EnvKeys.DATA_DIR.value, str(data_dir_abspath))
+        code.set_extra(EnvKeys.EXECUTABLE_PATH.value, str(code_executable_path))
+        code.set_extra(EnvKeys.IGNORE_FILES.value, ignore_files)
+        code.set_extra(EnvKeys.IGNORE_PATHS.value, ignore_paths)
+        code.set_extra(EnvKeys.REGENERATE_DATA.value, _regenerate_test_data)
+
         return code
 
     return _get_mock_code
+
+
+@pytest.fixture(scope='function', autouse=True)
+def patch_calculation_submission(monkeypatch):
+    """Patch execmanager.submit_calculation such as to take data from test data directory.
+    """
+    from aiida_testing.mock_code._env_keys import EnvKeys
+    from aiida_testing.mock_code._cli import get_hash, replace_submit_file
+    from aiida.engine.daemon import execmanager
+    import shutil
+    import sys
+    from pathlib import Path
+
+    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(EnvKeys.LABEL.value)
+        data_dir = code.get_extra(EnvKeys.DATA_DIR.value)
+        executable_path = code.get_extra(EnvKeys.EXECUTABLE_PATH.value)
+        ignore_files = code.get_extra(EnvKeys.IGNORE_FILES.value)
+        ignore_paths = code.get_extra(EnvKeys.IGNORE_PATHS.value)
+        regenerate_data = code.get_extra(EnvKeys.REGENERATE_DATA.value)
+
+        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)
+
+        #import pdb; pdb.set_trace()
+        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
+            workdir = calculation.get_remote_workdir()
+            replace_submit_file(executable_path=executable_path, working_directory=workdir)
+            #subprocess.call(['bash', SUBMIT_FILE])
+
+            ### Start copy of execmanager.submit_calculation
+            transport.chdir(workdir)
+            #func(calculation, transport)
+            job_id = calculation.get_job_id()
+
+            # If the `job_id` attribute is already set, that means this function was already executed once and the scheduler
+            # submit command was successful as the job id it returned was set on the node. This scenario can happen when the
+            # daemon runner gets shutdown right after accomplishing the submission task, but before it gets the chance to
+            # finalize the state transition of the `CalcJob` to the `UPDATE` transport task. Since the job is already submitted
+            # we do not want to submit it a second time, so we simply return the existing job id here.
+            if job_id is not None:
+                return job_id
+
+            scheduler = calculation.computer.get_scheduler()
+            scheduler.set_transport(transport)
+
+            submit_script_filename = calculation.get_option('submit_script_filename')
+            workdir = calculation.get_remote_workdir()
+            job_id = scheduler.submit_from_script(workdir, submit_script_filename)
+            calculation.set_job_id(job_id)
+
+            return job_id
+
+            ### End copy of execmanager.submit_calculation
+
+            ## Note: this backup will have to be done later, since the calculation may not be finished here
+            # # 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}'.")
+
+            # return a non-existing jobid
+            return -1
+
+    monkeypatch.setattr(execmanager, 'submit_calculation', mock_submit_calculation)

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:

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