Dev-iL commented on code in PR #76528:
URL: https://github.com/apache/airflow/pull/76528#discussion_r4478394900
##########
providers/ssh/tests/unit/ssh/hooks/test_ssh_async.py:
##########
@@ -197,6 +198,58 @@ async def test_get_conn_builds_config(self):
assert call_kwargs["username"] == "testuser"
assert result == mock_ssh_client
+ @mock.patch("asyncssh.connect", new_callable=mock.AsyncMock)
+ @mock.patch(
+ "airflow.providers.ssh.hooks.ssh.get_async_connection",
+ new_callable=mock.AsyncMock,
+ )
+ @pytest.mark.asyncio
+ async def test_get_conn_reads_extra_asynchronously(self,
mock_get_connection, mock_connect):
+ """``extra_dejson`` masks synchronously, so the async hook reads the
extra through common.compat."""
+ mock_get_connection.return_value = _connection({"no_host_key_check":
"true"})
+ hook = SSHHookAsync(ssh_conn_id="test_conn")
+
+ with (
+ mock.patch(
+ "airflow.providers.ssh.hooks.ssh.get_async_extra_dejson",
+ new_callable=mock.AsyncMock,
+ return_value={"no_host_key_check": "true"},
+ ) as mock_get_extra,
+ mock.patch.object(
+ Connection, "extra_dejson", new_callable=mock.PropertyMock,
side_effect=AssertionError
+ ) as mock_extra_dejson,
+ ):
+ await hook._get_conn()
+
+
mock_get_extra.assert_awaited_once_with(mock_get_connection.return_value)
+ mock_extra_dejson.assert_not_called()
+ assert mock_connect.call_args.kwargs["known_hosts"] is None
Review Comment:
The connection's own extra and the mocked `get_async_extra_dejson` return
value are both `{"no_host_key_check": "true"}`, so the `known_hosts is None`
assertion holds whichever source `_get_conn` parses. An implementation that
awaits the helper and then parses `json.loads(conn.extra)` still passes.
Building the connection with `_connection({})` makes the assertion hold only if
the helper's result reaches `_parse_extras`.
AGENTS.md also prefers `@mock.patch` decorators over `with mock.patch(...)`,
and both patches here take fixed values. The new `AsyncMock` patches have no
spec; `autospec=True` on `get_async_connection` and `get_async_extra_dejson`
gives an async mock with the real signature (the same applies to
`get_async_connection` in `test_get_conn_compresses_like_sync_hook`).
```suggestion
@mock.patch.object(
Connection, "extra_dejson", new_callable=mock.PropertyMock,
side_effect=AssertionError
)
@mock.patch(
"airflow.providers.ssh.hooks.ssh.get_async_extra_dejson",
autospec=True,
return_value={"no_host_key_check": "true"},
)
@mock.patch("asyncssh.connect", new_callable=mock.AsyncMock)
@mock.patch("airflow.providers.ssh.hooks.ssh.get_async_connection",
autospec=True)
@pytest.mark.asyncio
async def test_get_conn_reads_extra_asynchronously(
self, mock_get_connection, mock_connect, mock_get_extra,
mock_extra_dejson
):
"""``extra_dejson`` masks synchronously, so the async hook reads the
extra through common.compat."""
mock_get_connection.return_value = _connection({})
await SSHHookAsync(ssh_conn_id="test_conn")._get_conn()
mock_get_extra.assert_awaited_once_with(mock_get_connection.return_value)
mock_extra_dejson.assert_not_called()
assert mock_connect.call_args.kwargs["known_hosts"] is None
```
<!-- manifest-dev:review-pr -->
##########
providers/ssh/src/airflow/providers/ssh/hooks/ssh.py:
##########
@@ -722,6 +732,8 @@ def _get_value(self_val, conn_val, default=None):
"port": _get_value(self.port, conn.port, SSH_PORT),
"username": _get_value(self.username, conn.login),
"password": _get_value(self.password, conn.password),
+ # asyncssh prefers no compression; offer zlib first as paramiko
does for SSHHook.
+ "compression_algs": ["[email protected]", "zlib", "none"] if
self.compress else ["none"],
Review Comment:
Two side effects of passing `compression_algs` explicitly are worth a line
under `compress` in `docs/connections/ssh.rst:53`. asyncssh allocates zlib
state at key exchange, about 0.3 MB per connection with the default
`compressobj()` settings, and `SSHRemoteJobTrigger` holds its connection for
the whole job, so a triggerer running many of them grows by that much per
trigger. asyncssh also ignores a `Compression` setting from `~/.ssh/config`
once `compression_algs` is given (`_select_algs` in asyncssh's
`connection.py`), so a host-level `Compression no` stops applying to the async
hook. For example:
```rst
* ``compress`` - ``true`` to ask the remote client/server to compress
traffic; ``false`` to refuse compression. Default is ``true``.
The async hook ignores ``Compression`` from ``~/.ssh/config``, and
each compressed connection holds about 0.3 MB of zlib state;
set ``false`` on triggerers running many SSH remote job triggers.
```
<!-- manifest-dev:review-pr -->
##########
providers/ssh/src/airflow/providers/ssh/hooks/ssh.py:
##########
@@ -658,6 +664,8 @@ def _parse_extras(self, conn: Any) -> None:
)
if "private_key" in extra_options:
self.private_key = extra_options["private_key"]
+ if str(extra_options.get("compress", "true")).lower() == "false":
Review Comment:
`SSHHook` reads the same extra at `ssh.py:225` with a different expression.
The two behave the same; spelling them identically makes it obvious from either
hook that they agree.
```suggestion
if "compress" in extra_options and
str(extra_options["compress"]).lower() == "false":
```
<!-- manifest-dev:review-pr -->
##########
providers/ssh/tests/unit/ssh/hooks/test_ssh_async.py:
##########
@@ -32,6 +32,11 @@ def _connection(extra: dict, host: str = "test.host") ->
Connection:
return Connection(conn_id="test_conn", conn_type="ssh", host=host,
extra=json.dumps(extra))
+def _parse_extras(hook: SSHHookAsync, extra: dict, host: str = "test.host") ->
None:
Review Comment:
The module-level helper shares its name with the method it wraps, so
`_parse_extras(hook, {...})` and `hook._parse_extras(conn, {...})` sit side by
side in this file and differ only in call shape. Renaming it to
`_parse_extra_dict` here and at its call sites makes each call say which one
runs.
<!-- manifest-dev:review-pr -->
##########
providers/ssh/tests/unit/ssh/hooks/test_ssh_async.py:
##########
@@ -66,7 +71,7 @@ def test_parse_extras_key_file(self):
mock_conn.extra_dejson = {"key_file": "/path/to/key"}
mock_conn.host = "test.host"
- hook._parse_extras(mock_conn)
+ hook._parse_extras(mock_conn, mock_conn.extra_dejson)
Review Comment:
This test and five others (lines 84, 132, 143, 156, 166) still build a
`MagicMock` with an `extra_dejson` attribute only to pass it back in as the
second argument, while the tests on the `_parse_extras` helper at line 35 build
a real `Connection`. `extra_dejson` on the mock no longer means anything to
`_parse_extras`. Moving these onto the helper puts every `_parse_extras` test
on one setup. For this one, and the same shape for the other five:
```python
def test_parse_extras_key_file(self):
"""Test parsing key_file from connection extras."""
hook = SSHHookAsync(ssh_conn_id="test_conn")
_parse_extras(hook, {"key_file": "/path/to/key"})
assert hook.key_file == "/path/to/key"
```
<!-- manifest-dev:review-pr -->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
