fix async sudo askpass helper setup

Author: JakkuSakura

src/pyinfra/connectors/util.py

@@ -353,16 +353,28 @@ def extract_control_arguments(arguments: "ConnectorArguments") -> "ConnectorArgu
 def _ensure_sudo_askpass_set_for_host(host: "Host"):
     if host.connector_data.get("sudo_askpass_path"):
         return
-    _, output = host.run_shell_command(
+
+    status, output = host.run_shell_command(
         SUDO_ASKPASS_COMMAND.format(host.get_temp_dir_config(), SUDO_ASKPASS_ENV_VAR)
     )
+
+    if status is False or not output.stdout_lines or not output.stdout_lines[0]:
+        raise RuntimeError("Failed to create sudo askpass helper")
+
     host.connector_data["sudo_askpass_path"] = shlex.quote(output.stdout_lines[0])
 
 
 async def _ensure_sudo_askpass_set_for_host_async(host: "Host") -> None:
     if host.connector_data.get("sudo_askpass_path"):
         return
-    _, output = await host.run_shell_command_async(SUDO_ASKPASS_COMMAND)
+
+    status, output = await host.run_shell_command_async(
+        SUDO_ASKPASS_COMMAND.format(host.get_temp_dir_config(), SUDO_ASKPASS_ENV_VAR)
+    )
+
+    if status is False or not output.stdout_lines or not output.stdout_lines[0]:
+        raise RuntimeError("Failed to create sudo askpass helper")
+
     host.connector_data["sudo_askpass_path"] = shlex.quote(output.stdout_lines[0])
 
 

tests/test_connectors/test_util.py

@@ -2,8 +2,16 @@
 
 from unittest import TestCase
 
+import asyncio
+
 from pyinfra.api import Config, State
-from pyinfra.connectors.util import make_unix_command, make_unix_command_for_host
+from pyinfra.connectors.util import (
+    CommandOutput,
+    OutputLine,
+    async_make_unix_command_for_host,
+    make_unix_command,
+    make_unix_command_for_host,
+)
 
 from ..util import make_inventory
 
@@ -100,6 +108,30 @@ def test_mixed_command(self):
             "&& echo hi'\"'\"''"  # command bit
         )
 
+    def test_async_sudo_password_sets_askpass(self):
+        state = State(make_inventory(), Config())
+        host = state.inventory.get_host("somehost")
+
+        async def fake_run_shell_command_async(*_args, **_kwargs):
+            return True, CommandOutput([OutputLine("stdout", "/tmp/pyinfra-askpass")])
+
+        host.run_shell_command_async = fake_run_shell_command_async
+
+        command = asyncio.run(
+            async_make_unix_command_for_host(
+                state,
+                host,
+                "uptime",
+                _sudo=True,
+                _sudo_password="password",
+            )
+        )
+
+        raw = command.get_raw_value()
+        assert "SUDO_ASKPASS=/tmp/pyinfra-askpass" in raw
+        assert "PYINFRA_SUDO_PASSWORD=password" in raw
+        assert "sudo -H -A -k" in raw
+
     def test_command_exists_su_config_only(self):
         """
         This tests covers a bug that appeared when `make_unix_command` is called