From 6627b21fe308cf4d727d1761b687d6315e8fc4fa Mon Sep 17 00:00:00 2001
From: Ali <khsrali@gmail.com>
Date: Fri, 13 Dec 2024 12:35:32 +0100
Subject: [PATCH] update asyncssh dependency

---
 environment.yml                           |  2 +-
 pyproject.toml                            |  2 +-
 src/aiida/transports/plugins/ssh_async.py | 23 ++++++++++++++++-------
 tests/cmdline/commands/test_computer.py   | 10 +++++++++-
 tests/transports/test_all_plugins.py      |  4 +++-
 5 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/environment.yml b/environment.yml
index 8862430c4..9ce52da6f 100644
--- a/environment.yml
+++ b/environment.yml
@@ -8,7 +8,7 @@ dependencies:
 - python~=3.9
 - alembic~=1.2
 - archive-path~=0.4.2
-- asyncssh
+- asyncssh~=2.19.0
 - circus~=0.18.0
 - click-spinner~=0.1.8
 - click~=8.1
diff --git a/pyproject.toml b/pyproject.toml
index 971f338e8..92afa8986 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -20,7 +20,7 @@ classifiers = [
 dependencies = [
   'alembic~=1.2',
   'archive-path~=0.4.2',
-  "asyncssh",
+  "asyncssh~=2.19.0",
   'circus~=0.18.0',
   'click-spinner~=0.1.8',
   'click~=8.1',
diff --git a/src/aiida/transports/plugins/ssh_async.py b/src/aiida/transports/plugins/ssh_async.py
index c0d875b8a..307cffbbf 100644
--- a/src/aiida/transports/plugins/ssh_async.py
+++ b/src/aiida/transports/plugins/ssh_async.py
@@ -64,6 +64,8 @@ async def attempt_connection():
 class AsyncSshTransport(AsyncTransport):
     """Transport plugin via SSH, asynchronously."""
 
+    _DEFAULT_max_io_allowed = 8
+
     # note, I intentionally wanted to keep connection parameters as simple as possible.
     _valid_auth_options = [
         (
@@ -84,7 +86,7 @@ class AsyncSshTransport(AsyncTransport):
             'max_io_allowed',
             {
                 'type': int,
-                'default': 8,
+                'default': _DEFAULT_max_io_allowed,
                 'prompt': 'Maximum number of concurrent I/O operations.',
                 'help': 'Depends on various factors, such as your network bandwidth, the server load, etc.'
                 ' (An experimental number)',
@@ -117,23 +119,30 @@ def _get_machine_suggestion_string(cls, computer):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.machine = kwargs.pop('machine_or_host')
-        self._max_io_allowed = kwargs.pop('max_io_allowed')
+        # the machine is passed as `machine=computer.hostname` in the codebase
+        # 'machine' is immutable.
+        # 'machine_or_host' is mutable, so it can be changed via command:
+        # 'verdi computer configure core.ssh_async <LABEL>'.
+        # by default, 'machine_or_host' is set to 'machine' in the __init__ method, if not provided.
+        # NOTE: to guarantee a connection,
+        # a computer with core.ssh_async transport plugin should be configured before any instantiation.
+        self.machine = kwargs.pop('machine_or_host', kwargs.pop('machine'))
+        self._max_io_allowed = kwargs.pop('max_io_allowed', self._DEFAULT_max_io_allowed)
         self.script_before = kwargs.pop('script_before', 'None')
 
-        self._councurrent_io = 0
+        self._concurrent_io = 0
 
     @property
     def max_io_allowed(self):
         return self._max_io_allowed
 
     async def _lock(self, sleep_time=0.5):
-        while self._councurrent_io >= self.max_io_allowed:
+        while self._concurrent_io >= self.max_io_allowed:
             await asyncio.sleep(sleep_time)
-        self._councurrent_io += 1
+        self._concurrent_io += 1
 
     async def _unlock(self):
-        self._councurrent_io -= 1
+        self._concurrent_io -= 1
 
     async def open_async(self):
         """Open the transport.
diff --git a/tests/cmdline/commands/test_computer.py b/tests/cmdline/commands/test_computer.py
index 9959ea82b..3c98b7d1d 100644
--- a/tests/cmdline/commands/test_computer.py
+++ b/tests/cmdline/commands/test_computer.py
@@ -998,6 +998,14 @@ def test_computer_ssh_async(run_cli_command, aiida_computer):
 
     # It is important that 'ssh localhost' is functional in your test environment.
     # It should connect without asking for a password.
-    options = ['core.ssh_async', computer.uuid, '--non-interactive', '--safe-interval', '0', '--machine', 'localhost']
+    options = [
+        'core.ssh_async',
+        computer.uuid,
+        '--non-interactive',
+        '--safe-interval',
+        '0',
+        '--machine-or-host',
+        'localhost',
+    ]
     run_cli_command(computer_configure, options, use_subprocess=False)
     assert computer.is_configured
diff --git a/tests/transports/test_all_plugins.py b/tests/transports/test_all_plugins.py
index 6923a31a4..bc8c438e5 100644
--- a/tests/transports/test_all_plugins.py
+++ b/tests/transports/test_all_plugins.py
@@ -65,7 +65,9 @@ def custom_transport(request, tmp_path_factory, monkeypatch) -> Union['Transport
         if not filepath_config.exists():
             filepath_config.write_text('Host localhost')
     elif request.param == 'core.ssh_async':
-        kwargs = {'machine_or_host': 'localhost', 'max_io_allowed': 8}
+        kwargs = {
+            'machine': 'localhost',
+        }
     else:
         kwargs = {}