Share async shell command core in command_line utils (#182352)

homeassistant/components/command_line/notify.py

@@ -1,7 +1,5 @@
 """Support for command line notification services."""
 
-import asyncio
-from contextlib import suppress
 from typing import Any, override
 
 from homeassistant.components.notify import (
@@ -14,7 +12,11 @@ from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 
 from .const import CONF_COMMAND_TIMEOUT, DOMAIN, LOGGER
-from .utils import create_platform_yaml_not_supported_issue, render_template_args
+from .utils import (
+    async_run_shell_command,
+    create_platform_yaml_not_supported_issue,
+    render_template_args,
+)
 
 
 async def async_get_service(
@@ -51,11 +53,17 @@ class CommandLineNotificationService(BaseNotificationService):
         LOGGER.debug("Running with message: %s", message)
 
         try:
-            proc = await asyncio.create_subprocess_shell(  # shell by design
-                command,
-                stdin=asyncio.subprocess.PIPE,
-                close_fds=False,  # required for posix_spawn
+            proc, _ = await async_run_shell_command(
+                command, self._timeout, stdin=message.encode()
             )
+        except TimeoutError as err:
+            # TimeoutError subclasses OSError, so it must be caught first.
+            LOGGER.debug("Timeout for command: %s", command)
+            raise HomeAssistantError(
+                translation_domain=DOMAIN,
+                translation_key="timeout_error",
+                translation_placeholders={"command": command},
+            ) from err
         except OSError as err:
             LOGGER.debug("Error trying to exec command: %s", command)
             raise HomeAssistantError(
@@ -64,34 +72,6 @@ class CommandLineNotificationService(BaseNotificationService):
                 translation_placeholders={"command": command, "error": str(err)},
             ) from err
 
-        try:
-            async with asyncio.timeout(self._timeout):
-                await proc.communicate(input=message.encode())
-        except TimeoutError as err:
-            LOGGER.debug("Timeout for command: %s", command)
-            with suppress(ProcessLookupError):
-                # The command may have exited between the timeout and the kill.
-                proc.kill()
-            if (stdin := proc.stdin) is not None and (
-                not stdin.is_closing() or stdin.transport.get_write_buffer_size()
-            ):
-                # A still connected stdin pipe keeps proc.wait() pending forever,
-                # see https://bugs.python.org/issue43884.
-                stdin.transport.abort()
-            await proc.wait()
-            raise HomeAssistantError(
-                translation_domain=DOMAIN,
-                translation_key="timeout_error",
-                translation_placeholders={"command": command},
-            ) from err
-        except asyncio.CancelledError:
-            # Kill synchronously so the child isn't orphaned; the event loop
-            # reaps it without awaiting wait(), which cancellation would
-            # interrupt anyway.
-            with suppress(ProcessLookupError):
-                proc.kill()
-            raise
-
         if proc.returncode != 0:
             LOGGER.error(
                 "Command failed (with return code %s): %s",

homeassistant/components/command_line/utils.py

@@ -1,6 +1,8 @@
 """The command_line component utils."""
 
 import asyncio
+from contextlib import suppress
+from typing import Literal, overload
 
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import TemplateError
@@ -14,6 +16,72 @@ from .const import DOMAIN, LOGGER
 _EXEC_FAILED_CODE = 127
 
 
+@overload
+async def async_run_shell_command(
+    command: str,
+    timeout: int,
+    *,
+    stdin: bytes | None = ...,
+    capture_stdout: Literal[False] = ...,
+) -> tuple[asyncio.subprocess.Process, None]: ...
+
+
+@overload
+async def async_run_shell_command(
+    command: str,
+    timeout: int,
+    *,
+    stdin: bytes | None = ...,
+    capture_stdout: Literal[True],
+) -> tuple[asyncio.subprocess.Process, bytes]: ...
+
+
+async def async_run_shell_command(
+    command: str,
+    timeout: int,
+    *,
+    stdin: bytes | None = None,
+    capture_stdout: bool = False,
+) -> tuple[asyncio.subprocess.Process, bytes | None]:
+    """Run a shell command with a timeout and return the process and stdout.
+
+    The returned stdout is the captured bytes when capture_stdout is set, else None.
+    An OSError from spawning propagates; TimeoutError propagates after stdin cleanup
+    when stdin is provided.
+    """
+    proc = await asyncio.create_subprocess_shell(  # shell by design
+        command,
+        stdin=asyncio.subprocess.PIPE if stdin is not None else None,
+        stdout=asyncio.subprocess.PIPE if capture_stdout else None,
+        close_fds=False,  # required for posix_spawn
+    )
+    try:
+        async with asyncio.timeout(timeout):
+            stdout, _ = await proc.communicate(input=stdin)
+    except TimeoutError:
+        if stdin is not None:
+            with suppress(ProcessLookupError):
+                # The command may have exited between the timeout and the kill.
+                proc.kill()
+            if (proc_stdin := proc.stdin) is not None and (
+                not proc_stdin.is_closing()
+                or proc_stdin.transport.get_write_buffer_size()
+            ):
+                # A still connected stdin pipe keeps proc.wait() pending forever,
+                # see https://bugs.python.org/issue43884.
+                proc_stdin.transport.abort()
+            await proc.wait()
+        raise
+    except asyncio.CancelledError:
+        # Kill synchronously so the child isn't orphaned; the event loop
+        # reaps it without awaiting wait(), which cancellation would
+        # interrupt anyway.
+        with suppress(ProcessLookupError):
+            proc.kill()
+        raise
+    return proc, stdout
+
+
 async def async_call_shell_with_timeout(
     command: str, timeout: int, *, log_return_code: bool = True
 ) -> int:
@@ -22,14 +90,9 @@ async def async_call_shell_with_timeout(
     If log_return_code is set to False, it will not print an error if a non-zero
     return code is returned.
     """
+    LOGGER.debug("Running command: %s", command)
     try:
-        LOGGER.debug("Running command: %s", command)
-        proc = await asyncio.create_subprocess_shell(  # shell by design
-            command,
-            close_fds=False,  # required for posix_spawn
-        )
-        async with asyncio.timeout(timeout):
-            await proc.communicate()
+        proc, _ = await async_run_shell_command(command, timeout)
     except TimeoutError:
         LOGGER.error("Timeout for command: %s", command)
         return -1
@@ -49,24 +112,19 @@ async def async_call_shell_with_timeout(
 async def async_check_output_or_log(command: str, timeout: int) -> str | None:
     """Run a shell command with a timeout and return the output."""
     try:
-        proc = await asyncio.create_subprocess_shell(  # shell by design
-            command,
-            close_fds=False,  # required for posix_spawn
-            stdout=asyncio.subprocess.PIPE,
+        proc, stdout = await async_run_shell_command(
+            command, timeout, capture_stdout=True
         )
-        async with asyncio.timeout(timeout):
-            stdout, _ = await proc.communicate()
-
-        if proc.returncode != 0:
-            LOGGER.error(
-                "Command failed (with return code %s): %s", proc.returncode, command
-            )
-        else:
-            return stdout.strip().decode("utf-8")
     except TimeoutError:
         LOGGER.error("Timeout for command: %s", command)
+        return None
 
-    return None
+    if proc.returncode != 0:
+        LOGGER.error(
+            "Command failed (with return code %s): %s", proc.returncode, command
+        )
+        return None
+    return stdout.strip().decode("utf-8")
 
 
 def render_template_args(hass: HomeAssistant, command: str) -> str | None:

tests/components/command_line/__init__.py

@@ -16,7 +16,7 @@ def mock_asyncio_subprocess_run(
         def returncode(self):
             return returncode
 
-        async def communicate(self):
+        async def communicate(self, input=None):
             if exception:
                 raise exception
             return response, b""

tests/components/command_line/test_cover.py

@@ -80,6 +80,7 @@ async def test_poll_when_cover_has_command_state(
         await hass.async_block_till_done()
         mock_subprocess_run.assert_called_once_with(
             "echo state",
+            stdin=None,
             close_fds=False,
             stdout=-1,
         )

tests/components/command_line/test_notify.py

@@ -295,7 +295,7 @@ async def test_spawn_error(
 
     with (
         patch(
-            "homeassistant.components.command_line.notify.asyncio.create_subprocess_shell",
+            "homeassistant.components.command_line.utils.asyncio.create_subprocess_shell",
             side_effect=OSError("exec failed"),
         ),
         pytest.raises(HomeAssistantError) as exc_info,
@@ -356,7 +356,7 @@ async def test_timeout_cleanup(
 
     with (
         patch(
-            "homeassistant.components.command_line.notify.asyncio.create_subprocess_shell",
+            "homeassistant.components.command_line.utils.asyncio.create_subprocess_shell",
             return_value=mock_proc,
         ),
         pytest.raises(HomeAssistantError) as exc_info,
@@ -397,7 +397,7 @@ async def test_cancelled_kills_process(
 
     with (
         patch(
-            "homeassistant.components.command_line.notify.asyncio.create_subprocess_shell",
+            "homeassistant.components.command_line.utils.asyncio.create_subprocess_shell",
             return_value=mock_proc,
         ),
         pytest.raises(asyncio.CancelledError),

tests/components/command_line/test_sensor.py

@@ -145,6 +145,7 @@ async def test_template_render_with_quote(hass: HomeAssistant) -> None:
         assert len(mock_subprocess_run.mock_calls) == 1
         mock_subprocess_run.assert_called_with(
             'echo "sensor_value" "3 4"',
+            stdin=None,
             stdout=-1,
             close_fds=False,
         )