From 37f48064d0b34577fd40c4f2e45065a46d3bc123 Mon Sep 17 00:00:00 2001
From: Julian Geiger <julian.geiger@psi.ch>
Date: Mon, 27 May 2024 15:00:22 +0200
Subject: [PATCH] Capture `ImportTestRun` for migration and add test

---
 src/aiida/cmdline/commands/cmd_archive.py     |  8 +++++---
 tests/cmdline/commands/test_archive_import.py | 11 +++++++++++
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/aiida/cmdline/commands/cmd_archive.py b/src/aiida/cmdline/commands/cmd_archive.py
index 07458cb79d..2db9fbcd62 100644
--- 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,6 +517,9 @@ 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
@@ -523,9 +527,7 @@ def _import_archive_and_migrate(
             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
index 92b6690433..683486017f 100644
--- 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.