Skip to content

Commit

Permalink
Make CommandResult a dataclass (#722)
Browse files Browse the repository at this point in the history
  • Loading branch information
amaslenn authored Nov 13, 2023
1 parent 76170a4 commit dc48cd9
Show file tree
Hide file tree
Showing 4 changed files with 32 additions and 50 deletions.
8 changes: 3 additions & 5 deletions testinfra/backend/ansible.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,11 +58,9 @@ def run(self, command: str, *args: str, **kwargs: Any) -> base.CommandResult:
out = self.run_ansible("shell", module_args=command, check=False)
return self.result(
out["rc"],
command,
stdout_bytes=None,
stderr_bytes=None,
stdout=out["stdout"],
stderr=out["stderr"],
self.encode(command),
out["stdout"],
out["stderr"],
)

def run_ansible(
Expand Down
68 changes: 27 additions & 41 deletions testinfra/backend/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
import shlex
import subprocess
import urllib.parse
from typing import TYPE_CHECKING, Any, Optional
from typing import TYPE_CHECKING, Any, Optional, Union

if TYPE_CHECKING:
import testinfra.host
Expand All @@ -33,6 +33,7 @@ class HostSpec:
password: Optional[str]


@dataclasses.dataclass
class CommandResult:
"""Object that encapsulates all returned details of the command execution.
Expand All @@ -51,24 +52,11 @@ class CommandResult:
False
"""

def __init__(
self,
backend: "BaseBackend",
exit_status: int,
command: bytes,
stdout_bytes: bytes,
stderr_bytes: bytes,
stdout: Optional[str] = None,
stderr: Optional[str] = None,
):
self.exit_status = exit_status
self._stdout_bytes = stdout_bytes
self._stderr_bytes = stderr_bytes
self._stdout = stdout
self._stderr = stderr
self.command = command
self._backend = backend
super().__init__()
backend: "BaseBackend"
exit_status: int
command: bytes
_stdout: Union[str, bytes]
_stderr: Union[str, bytes]

@property
def succeeded(self) -> bool:
Expand Down Expand Up @@ -104,8 +92,8 @@ def stdout(self) -> str:
>>> host.run("mkdir -v new_directory").stdout
mkdir: created directory 'new_directory'
"""
if self._stdout is None:
self._stdout = self._backend.decode(self._stdout_bytes)
if isinstance(self._stdout, bytes):
return self.backend.decode(self._stdout)
return self._stdout

@property
Expand All @@ -115,8 +103,8 @@ def stderr(self) -> str:
>>> host.run("mkdir new_directory").stderr
mkdir: cannot create directory 'new_directory': File exists
"""
if self._stderr is None:
self._stderr = self._backend.decode(self._stderr_bytes)
if isinstance(self._stderr, bytes):
return self.backend.decode(self._stderr)
return self._stderr

@property
Expand All @@ -126,9 +114,9 @@ def stdout_bytes(self) -> bytes:
>>> host.run("mkdir -v new_directory").stdout_bytes
b"mkdir: created directory 'new_directory'"
"""
if self._stdout_bytes is None:
self._stdout_bytes = self._backend.encode(self._stdout)
return self._stdout_bytes
if isinstance(self._stdout, str):
return self.backend.encode(self._stdout)
return self._stdout

@property
def stderr_bytes(self) -> bytes:
Expand All @@ -137,19 +125,9 @@ def stderr_bytes(self) -> bytes:
>>> host.run("mkdir new_directory").stderr_bytes
b"mkdir: cannot create directory 'new_directory': File exists"
"""
if self._stderr_bytes is None:
self._stderr_bytes = self._backend.encode(self._stderr)
return self._stderr_bytes

def __repr__(self) -> str:
return (
"CommandResult(command={!r}, exit_status={}, stdout={!r}, " "stderr={!r})"
).format(
self.command,
self.exit_status,
self._stdout_bytes or self._stdout,
self._stderr_bytes or self._stderr,
)
if isinstance(self._stderr, str):
return self.backend.encode(self._stderr)
return self._stderr


class BaseBackend(metaclass=abc.ABCMeta):
Expand Down Expand Up @@ -337,7 +315,15 @@ def encode(self, data: str) -> bytes:
except UnicodeEncodeError:
return data.encode(self.encoding)

def result(self, *args: Any, **kwargs: Any) -> CommandResult:
result = CommandResult(self, *args, **kwargs)
def result(
self, rc: int, cmd: bytes, stdout: Union[str, bytes], stderr: Union[str, bytes]
) -> CommandResult:
result = CommandResult(
backend=self,
exit_status=rc,
command=cmd,
_stdout=stdout,
_stderr=stderr,
)
logger.debug("RUN %s", result)
return result
4 changes: 1 addition & 3 deletions testinfra/backend/salt.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,9 +40,7 @@ def run(self, command: str, *args: str, **kwargs: Any) -> base.CommandResult:
out = self.run_salt("cmd.run_all", [command])
return self.result(
out["retcode"],
command,
out["stdout"].encode("utf8"),
out["stderr"].encode("utf8"),
self.encode(command),
stdout=out["stdout"],
stderr=out["stderr"],
)
Expand Down
2 changes: 1 addition & 1 deletion testinfra/backend/winrm.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,7 @@ def run_winrm(self, command: str, *args: str) -> base.CommandResult:
stdout, stderr, rc = p.get_command_output(shell_id, command_id)
p.cleanup_command(shell_id, command_id)
p.close_shell(shell_id)
return self.result(rc, command, stdout, stderr)
return self.result(rc, self.encode(command), stdout, stderr)

@staticmethod
def quote(command: str, *args: str) -> str:
Expand Down

0 comments on commit dc48cd9

Please sign in to comment.