Capture ImportTestRun for migration and add test

aiidateam · May 27, 2024 · 37f4806 · 37f4806
1 parent b0b0606
commit 37f4806
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/src/aiida/cmdline/commands/cmd_archive.py b/src/aiida/cmdline/commands/cmd_archive.py
@@ -485,6 +485,7 @@ def _import_archive_and_migrate(
 
     archive_format = get_format()
     filepath = ctx.obj['config'].get_option('storage.sandbox') or None
+    dry_run_success = f'import dry-run of archive {archive} completed. Profile storage unmodified.'
 
     with SandboxFolder(filepath=filepath) as temp_folder:
         archive_path = archive
@@ -516,16 +517,17 @@ def _import_archive_and_migrate(
                 echo.echo_report('proceeding with import of migrated archive')
                 try:
                     _import_archive(archive_path, archive_format=archive_format, **import_kwargs)
+                except ImportTestRun:
+                    echo.echo_success(dry_run_success)
+                    return
                 except Exception as sub_exception:
                     _echo_exception(
                         f'an exception occurred while trying to import the migrated archive {archive}', sub_exception
                     )
             else:
                 _echo_exception(f'an exception occurred while trying to import the archive {archive}', exception)
         except ImportTestRun:
-            echo.echo_success(
-                f'Import dry-run of archive {archive} terminated successfully. Profile storage unmodified.'
-            )
+            echo.echo_success(dry_run_success)
             return
 
         except Exception as exception:

diff --git a/tests/cmdline/commands/test_archive_import.py b/tests/cmdline/commands/test_archive_import.py
@@ -49,6 +49,17 @@ def test_import_archive(run_cli_command, newest_archive):
     run_cli_command(cmd_archive.import_archive, options)
 
 
+@pytest.mark.parametrize('archive', (
+    get_archive_file('arithmetic.add.aiida', filepath='calcjob'),
+    get_archive_file('export_0.9_simple.aiida', filepath=ARCHIVE_PATH),
+))
+def test_import_dry_run(run_cli_command, archive):
+    """Test import dry-run"""
+    result = run_cli_command(cmd_archive.import_archive, [archive, '--dry-run'])
+    print(result.output)
+    assert f'import dry-run of archive {archive} completed' in result.output
+
+
 def test_import_to_group(run_cli_command, newest_archive):
     """Test import to existing Group and that Nodes are added correctly for multiple imports of the same,
     as well as separate, archives.