201 lines
6.5 KiB
Python
201 lines
6.5 KiB
Python
import contextlib
|
|
import dataclasses
|
|
import json
|
|
import logging
|
|
import subprocess
|
|
from pathlib import Path
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from testlib.fixtures.env import ManagedEnv
|
|
from testlib.terminal_code_eater import eat_terminal_codes
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class CommandResult:
|
|
cmd: list[str]
|
|
"""Command arguments which were run"""
|
|
rc: int
|
|
"""Return code"""
|
|
stderr: bytes | None
|
|
"""Outputted stderr"""
|
|
stdout: bytes
|
|
"""Outputted stdout"""
|
|
|
|
def ok(self) -> "CommandResult":
|
|
"""
|
|
assumes a return code of 0
|
|
:raises CalledProcessError: if the return code wasn't 0 and logs the processes stdout and stderr
|
|
"""
|
|
__tracebackhide__ = True
|
|
return self.expect(0)
|
|
|
|
def expect(self, rc: int) -> "CommandResult":
|
|
"""
|
|
assumes a return code of `rc`
|
|
:param rc: The expected return code
|
|
:raises CalledProcessError: if the return code wasn't `rc` and logs the processes stdout and stderr
|
|
"""
|
|
__tracebackhide__ = True
|
|
if self.rc != rc:
|
|
logger.error("stdout: %s", self.stdout_s)
|
|
logger.error("stderr: %s", self.stderr_s)
|
|
exc = subprocess.CalledProcessError(
|
|
returncode=self.rc, cmd=self.cmd, stderr=self.stderr, output=self.stdout
|
|
)
|
|
raise exc
|
|
return self
|
|
|
|
@property
|
|
def stdout_s(self) -> str:
|
|
"""Command stdout as str"""
|
|
return self.stdout.decode("utf-8", errors="replace")
|
|
|
|
@property
|
|
def stderr_s(self) -> str:
|
|
"""Command stderr as str"""
|
|
return self.stderr.decode("utf-8", errors="replace") if self.stderr is not None else ""
|
|
|
|
@property
|
|
def stdout_plain(self) -> str:
|
|
"""Command stderr as str with terminal escape sequences eaten and whitespace stripped"""
|
|
return eat_terminal_codes(self.stdout).decode("utf-8", errors="replace").strip()
|
|
|
|
@property
|
|
def stderr_plain(self) -> str:
|
|
"""Command stderr as str with terminal escape sequences eaten and whitespace stripped"""
|
|
return eat_terminal_codes(self.stderr).decode("utf-8", errors="replace").strip()
|
|
|
|
def json(self) -> Any:
|
|
"""
|
|
Assumes an ok() result and returns the Commands stdout parsed as json
|
|
:return: A parsed json object
|
|
"""
|
|
__tracebackhide__ = True
|
|
self.ok()
|
|
return json.loads(self.stdout)
|
|
|
|
|
|
class RunningCommand(contextlib.AbstractContextManager):
|
|
argv: list[str]
|
|
stdin: bytes | None = None
|
|
_proc: subprocess.Popen
|
|
|
|
def __init__(self, argv: list[str], stdin: bytes | None, proc: subprocess.Popen):
|
|
self.argv = argv
|
|
self.stdin = stdin
|
|
self._proc = proc
|
|
|
|
def __exit__(self, exc_type, exc, tb): # noqa: ANN001
|
|
self.kill()
|
|
|
|
def kill(self) -> CommandResult | None:
|
|
"""
|
|
Kill the process immediately without waiting for it to exit cleanly.
|
|
:return: `None` if the process was already dead, else the process result.
|
|
"""
|
|
if self._proc is not None:
|
|
self._proc.kill()
|
|
# wait forever. killing must never fail, so we'd rather timeout than not notice errors here
|
|
return self.wait()
|
|
return None
|
|
|
|
def terminate(self, timeout: float | None = None) -> CommandResult | None:
|
|
"""
|
|
Send a termination signal to the process and waits for it to exit. `None` timeouts are treated as infinite.
|
|
:return: `None` if timeout expired before the process exited, else the process result.
|
|
"""
|
|
self._proc.terminate()
|
|
return self.wait(timeout)
|
|
|
|
def wait(self, timeout: float | None = None) -> CommandResult | None:
|
|
"""
|
|
Waits for the process to exit. `None` timeouts are treated as infinite.
|
|
:return: `None` if timeout expired before the process exited, else the process result.
|
|
"""
|
|
try:
|
|
stdout, stderr = self._proc.communicate(input=self.stdin, timeout=timeout)
|
|
rc = self._proc.returncode
|
|
self._proc = None
|
|
return CommandResult(cmd=self.argv, rc=rc, stdout=stdout, stderr=stderr)
|
|
except subprocess.TimeoutExpired:
|
|
return None
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Command:
|
|
argv: list[str]
|
|
_env: ManagedEnv
|
|
exe: Path | None = None
|
|
stdin: bytes | None = None
|
|
err_to_out: bool = False
|
|
cwd: Path = dataclasses.field(default=None)
|
|
_logger: logging.Logger = dataclasses.field(default=logger, init=False)
|
|
_run: bool = dataclasses.field(default=False, init=False)
|
|
|
|
def __post_init__(self):
|
|
self._env._registered_commands.append(self)
|
|
if self.cwd is None:
|
|
self.cwd = self._env.dirs.home
|
|
|
|
def with_stdin(self, stdin: bytes) -> "Command":
|
|
self.stdin = stdin
|
|
return self
|
|
|
|
def set_args(self, *argv: str) -> "Command":
|
|
self.argv = list(argv)
|
|
return self
|
|
|
|
def with_wrapper(self, cmd: str | Path, *args: list[str]) -> "Command":
|
|
"""
|
|
Wraps the current command in the given wrapper.
|
|
e.g: Command(["nix", "eval", "1+1"]).with_wrapper("strace") => Command(["strace", "nix", "eval", "1+1"])
|
|
"""
|
|
self.argv = [cmd, *args, *self.argv]
|
|
if self.exe:
|
|
self.exe = None if isinstance(cmd, str) else cmd
|
|
return self
|
|
|
|
def run(self) -> CommandResult:
|
|
"""
|
|
Runs the configured command
|
|
:return: Information about the Result of the execution
|
|
"""
|
|
return self.start().wait()
|
|
|
|
def start(self) -> RunningCommand:
|
|
"""
|
|
Starts the configured command
|
|
:return: Handle to the running command for interaction or waiting
|
|
"""
|
|
self._run = True
|
|
self._logger.debug("Running Command with args: %s; stdin: %s", self.argv, self.stdin)
|
|
proc = subprocess.Popen(
|
|
self.argv,
|
|
executable=self.exe,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT if self.err_to_out else subprocess.PIPE,
|
|
stdin=subprocess.PIPE if self.stdin else subprocess.DEVNULL,
|
|
cwd=self.cwd,
|
|
env=self._env.to_env(),
|
|
)
|
|
return RunningCommand(self.argv, self.stdin, proc)
|
|
|
|
def discard(self):
|
|
"""
|
|
Discards the command, suppressing the "Command not run" error, which would otherwise be raised
|
|
"""
|
|
self._run = True
|
|
|
|
|
|
@pytest.fixture
|
|
def command(env: ManagedEnv) -> Callable[..., Command]:
|
|
def wrapper(*args, **kwargs) -> Command:
|
|
return Command(_env=env, *args, **kwargs)
|
|
|
|
return wrapper
|