Fix flaky test_command_timeout_fail in SSH provider

providers/ssh/tests/unit/ssh/hooks/test_ssh.py

@@ -929,26 +929,47 @@ def test_command_timeout_success(self):
             assert ret == (0, b"airflow\n", b"")
 
     def test_command_timeout_fail(self):
-        # cmd_timeout is forwarded to paramiko's exec_command, which uses it both to open the
-        # channel and to read the command output. It must therefore be large enough to reliably
-        # open the channel (otherwise a loaded runner raises "Timeout opening channel." instead of
-        # the AirflowException we expect) while staying smaller than the command runtime so the
-        # read loop times out. A sub-millisecond timeout makes channel opening flaky.
         hook = SSHHook(
             ssh_conn_id="ssh_default",
             conn_timeout=30,
-            cmd_timeout=0.5,
+            cmd_timeout=0.001,
             banner_timeout=100,
         )
 
-        with hook.get_conn() as client:
-            with pytest.raises(AirflowException):
-                hook.exec_ssh_client_command(
-                    client,
-                    "sleep 5",
-                    False,
-                    None,
-                )
+        mock_channel = mock.MagicMock(spec=paramiko.Channel)
+        type(mock_channel).closed = mock.PropertyMock(return_value=False)
+        mock_channel.recv_ready.return_value = False
+        mock_channel.recv_stderr_ready.return_value = False
+        mock_channel.exit_status_ready.return_value = False
+        mock_channel.in_buffer = b""
+        mock_channel.in_stderr_buffer = b""
+
+        mock_stdout = mock.MagicMock(spec=paramiko.ChannelFile)
+        mock_stdout.channel = mock_channel
+        mock_stdin = mock.MagicMock(spec=paramiko.ChannelStdinFile)
+        mock_stderr = mock.MagicMock(spec=paramiko.ChannelStderrFile)
+        mock_stderr.channel = mock_channel
+
+        mock_client = mock.MagicMock(spec=paramiko.SSHClient)
+        mock_client.exec_command.return_value = (mock_stdin, mock_stdout, mock_stderr)
+
+        def fake_select(rlist, wlist, xlist, timeout=None):
+            assert timeout == pytest.approx(0.001), f"Expected cmd_timeout passed to select, got {timeout}"
+            return [], [], []
+
+        with mock.patch("airflow.providers.ssh.hooks.ssh.select", side_effect=fake_select):
+            with pytest.raises(AirflowException, match="SSH command timed out"):
+                hook.exec_ssh_client_command(mock_client, "sleep 1", False, None)
+
+        assert mock_client.exec_command.call_args_list == [
+            mock.call(command="sleep 1", get_pty=False, timeout=0.001, environment=None)
+        ]
+        assert mock.call() in mock_stdin.close.call_args_list
+        assert mock.call() in mock_channel.shutdown_write.call_args_list
+        assert mock.call() in mock_channel.shutdown_read.call_args_list
+        assert mock.call() in mock_channel.close.call_args_list
+        assert mock.call() in mock_stdout.close.call_args_list
+        assert mock.call() in mock_stderr.close.call_args_list
 
     def test_command_timeout_not_set(self, monkeypatch):
         hook = SSHHook(