From a750dfd1489763d2966cf0baba6840d84745bcb6 Mon Sep 17 00:00:00 2001 From: Vladimir Domnich Date: Wed, 5 Oct 2022 20:41:47 +0400 Subject: [PATCH] [#9] Implement hosting package Package defines interface for host management and provides implementation for docker host (local or remote). Other hosts can be added via plugins. Signed-off-by: Vladimir Domnich --- README.md | 11 + pyproject.toml | 7 +- requirements.txt | 2 + src/neofs_testlib/hosting/__init__.py | 3 + src/neofs_testlib/hosting/config.py | 70 +++++++ src/neofs_testlib/hosting/docker_host.py | 193 ++++++++++++++++++ src/neofs_testlib/hosting/hosting.py | 107 ++++++++++ src/neofs_testlib/hosting/interfaces.py | 124 +++++++++++ src/neofs_testlib/reporter/allure_handler.py | 4 +- src/neofs_testlib/reporter/reporter.py | 2 +- src/neofs_testlib/shell/__init__.py | 2 +- src/neofs_testlib/shell/command_inspectors.py | 13 ++ src/neofs_testlib/shell/interfaces.py | 35 +++- src/neofs_testlib/shell/local_shell.py | 9 +- src/neofs_testlib/shell/ssh_shell.py | 19 +- 15 files changed, 582 insertions(+), 19 deletions(-) create mode 100644 src/neofs_testlib/hosting/__init__.py create mode 100644 src/neofs_testlib/hosting/config.py create mode 100644 src/neofs_testlib/hosting/docker_host.py create mode 100644 src/neofs_testlib/hosting/hosting.py create mode 100644 src/neofs_testlib/hosting/interfaces.py create mode 100644 src/neofs_testlib/shell/command_inspectors.py diff --git a/README.md b/README.md index cd4593c..ed28dfc 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,16 @@ Alternative approach for registering handlers is to use method `configure`. It i get_reporter().configure({ "handlers": [{"plugin_name": "allure"}] }) ``` +### Hosting Configuration +Hosting component is a class that represents infrastructure (machines/containers/services) where neoFS is hosted. Interaction with specific infrastructure instance (host) is encapsulated in classes that implement interface `neofs_testlib.hosting.Host`. To pass information about hosts to the `Hosting` class in runtime we use method `configure`: + +```python +from neofs_testlib.hosting import Hosting + +hosting = Hosting() +hosting.configure({ "hosts": [{ "address": "localhost", "plugin_name": "docker" ... }]}) +``` + ## Plugins Testlib uses [entrypoint specification](https://docs.python.org/3/library/importlib.metadata.html) for plugins. Testlib supports the following entrypoint groups for plugins: - `neofs.testlib.reporter` - group for reporter handler plugins. Plugin should be a class that implements interface `neofs_testlib.reporter.interfaces.ReporterHandler`. @@ -74,6 +84,7 @@ Detailed information about registering entrypoints can be found at [setuptools d ## Library structure The library provides the following primary components: * `cli` - wrappers on top of neoFS command-line tools. These wrappers execute on a shell and provide type-safe interface for interacting with the tools. + * `hosting` - management of infrastructure (docker, virtual machines, services where neoFS is hosted). The library provides host implementation for docker environment (when neoFS services are running as docker containers). Support for other hosts is provided via plugins. * `reporter` - abstraction on top of test reporting tool like Allure. Components of the library will report their steps and attach artifacts to the configured reporter instance. * `shell` - shells that can be used to execute commands. Currently library provides local shell (on machine that runs the code) or SSH shell that connects to a remote machine via SSH. diff --git a/pyproject.toml b/pyproject.toml index d4b3eec..78c88f5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["setuptools>=65.0.0", "wheel"] +requires = ["setuptools>=65.0.0", "wheel"] build-backend = "setuptools.build_meta" [project] @@ -17,9 +17,11 @@ classifiers = [ keywords = ["neofs", "test"] dependencies = [ "allure-python-commons>=2.9.45", + "docker>=4.4.0", "importlib_metadata>=5.0; python_version < '3.10'", "paramiko>=2.10.3", "pexpect>=4.8.0", + "requests>=2.28.0", ] requires-python = ">=3.9" @@ -32,6 +34,9 @@ Homepage = "https://github.com/nspcc-dev/neofs-testlib" [project.entry-points."neofs.testlib.reporter"] allure = "neofs_testlib.reporter.allure_handler:AllureHandler" +[project.entry-points."neofs.testlib.hosting"] +docker = "neofs_testlib.hosting.docker_host:DockerHost" + [tool.isort] profile = "black" src_paths = ["src", "tests"] diff --git a/requirements.txt b/requirements.txt index 39b6bd3..294f406 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,9 @@ allure-python-commons==2.9.45 +docker==4.4.0 importlib_metadata==5.0.0 paramiko==2.10.3 pexpect==4.8.0 +requests==2.28.1 # Dev dependencies black==22.8.0 diff --git a/src/neofs_testlib/hosting/__init__.py b/src/neofs_testlib/hosting/__init__.py new file mode 100644 index 0000000..d3f1f8f --- /dev/null +++ b/src/neofs_testlib/hosting/__init__.py @@ -0,0 +1,3 @@ +from neofs_testlib.hosting.config import CLIConfig, HostConfig, ServiceConfig +from neofs_testlib.hosting.hosting import Hosting +from neofs_testlib.hosting.interfaces import Host diff --git a/src/neofs_testlib/hosting/config.py b/src/neofs_testlib/hosting/config.py new file mode 100644 index 0000000..febc848 --- /dev/null +++ b/src/neofs_testlib/hosting/config.py @@ -0,0 +1,70 @@ +from dataclasses import dataclass, field, fields +from typing import Any + + +@dataclass +class ParsedAttributes: + """Base class for data structures representing parsed attributes from configs.""" + + @classmethod + def parse(cls, attributes: dict[str, Any]): + # Pick attributes supported by the class + field_names = set(field.name for field in fields(cls)) + supported_attributes = { + key: value for key, value in attributes.items() if key in field_names + } + return cls(**supported_attributes) + + +@dataclass +class CLIConfig: + """Describes CLI tool on some host. + + Attributes: + name: Name of the tool. + exec_path: Path to executable file of the tool. + attributes: Dict with extra information about the tool. + """ + + name: str + exec_path: str + attributes: dict[str, str] = field(default_factory=dict) + + +@dataclass +class ServiceConfig: + """Describes neoFS service on some host. + + Attributes: + name: Name of the service that uniquely identifies it across all hosts. + attributes: Dict with extra information about the service. For example, we can store + name of docker container (or name of systemd service), endpoints, path to wallet, + path to configuration file, etc. + """ + + name: str + attributes: dict[str, str] = field(default_factory=dict) + + +@dataclass +class HostConfig: + """Describes machine that hosts neoFS services. + + Attributes: + plugin_name: Name of plugin that should be used to manage the host. + address: Address of the machine (IP or DNS name). + services: List of services hosted on the machine. + clis: List of CLI tools available on the machine. + attributes: Dict with extra information about the host. For example, we can store + connection parameters in this dict. + """ + + plugin_name: str + address: str + services: list[ServiceConfig] = field(default_factory=list) + clis: list[CLIConfig] = field(default_factory=list) + attributes: dict[str, str] = field(default_factory=dict) + + def __post_init__(self) -> None: + self.services = [ServiceConfig(**service) for service in self.services or []] + self.clis = [CLIConfig(**cli) for cli in self.clis or []] diff --git a/src/neofs_testlib/hosting/docker_host.py b/src/neofs_testlib/hosting/docker_host.py new file mode 100644 index 0000000..34ebb87 --- /dev/null +++ b/src/neofs_testlib/hosting/docker_host.py @@ -0,0 +1,193 @@ +import json +import logging +import os +import time +from dataclasses import dataclass +from datetime import datetime +from typing import Any, Optional + +import docker +from requests import HTTPError + +from neofs_testlib.hosting.config import ParsedAttributes +from neofs_testlib.hosting.interfaces import Host +from neofs_testlib.shell import LocalShell, Shell, SSHShell +from neofs_testlib.shell.command_inspectors import SudoInspector + +logger = logging.getLogger("neofs.testlib.hosting") + + +@dataclass +class HostAttributes(ParsedAttributes): + """Represents attributes of host where Docker with neoFS runs. + + Attributes: + sudo_shell: Specifies whether shell commands should be auto-prefixed with sudo. + docker_endpoint: Protocol, address and port of docker where neoFS runs. Recommended format + is tcp socket (https://docs.docker.com/engine/reference/commandline/dockerd/#daemon-socket-option), + for example: tcp://{address}:2375 (where 2375 is default docker port). + ssh_login: Login for SSH connection to the machine where docker runs. + ssh_password: Password for SSH connection. + ssh_private_key_path: Path to private key for SSH connection. + ssh_private_key_passphrase: Passphrase for the private key. + """ + + sudo_shell: bool = False + docker_endpoint: Optional[str] = None + ssh_login: Optional[str] = None + ssh_password: Optional[str] = None + ssh_private_key_path: Optional[str] = None + ssh_private_key_passphrase: Optional[str] = None + + +@dataclass +class ServiceAttributes(ParsedAttributes): + """Represents attributes of service running as Docker container. + + Attributes: + container_name: Name of Docker container where the service runs. + volume_name: Name of volume where storage node service stores the data. + start_timeout: Timeout (in seconds) for service to start. + stop_timeout: Timeout (in seconds) for service to stop. + """ + + container_name: str + volume_name: Optional[str] = None + start_timeout: int = 60 + stop_timeout: int = 60 + + +class DockerHost(Host): + """Manages services hosted in Docker containers running on a local or remote machine.""" + + def get_shell(self) -> Shell: + host_attributes = HostAttributes.parse(self._config.attributes) + command_inspectors = [] + if host_attributes.sudo_shell: + command_inspectors.append(SudoInspector()) + + if not host_attributes.ssh_login: + # If there is no SSH connection to the host, use local shell + return LocalShell(command_inspectors) + + # If there is SSH connection to the host, use SSH shell + return SSHShell( + host=self._config.address, + login=host_attributes.ssh_login, + password=host_attributes.ssh_password, + private_key_path=host_attributes.ssh_private_key_path, + private_key_passphrase=host_attributes.ssh_private_key_passphrase, + command_inspectors=command_inspectors, + ) + + def start_host(self) -> None: + # We emulate starting machine by starting all services + # As an alternative we can probably try to stop docker service... + for service_config in self._config.services: + self.start_service(service_config.name) + + def stop_host(self) -> None: + # We emulate stopping machine by stopping all services + # As an alternative we can probably try to stop docker service... + for service_config in self._config.services: + self.stop_service(service_config.name) + + def start_service(self, service_name: str) -> None: + service_attributes = self._get_service_attributes(service_name) + + client = self._get_docker_client() + client.start(service_attributes.container_name) + + self._wait_for_container_to_be_in_state( + container_name=service_attributes.container_name, + expected_state="running", + timeout=service_attributes.start_timeout, + ) + + def stop_service(self, service_name: str) -> None: + service_attributes = self._get_service_attributes(service_name) + + client = self._get_docker_client() + client.stop(service_attributes.container_name) + + self._wait_for_container_to_be_in_state( + container_name=service_attributes.container_name, + expected_state="exited", + timeout=service_attributes.stop_timeout, + ) + + def delete_storage_node_data(self, service_name: str) -> None: + service_attributes = self._get_service_attributes(service_name) + + client = self._get_docker_client() + volume_info = client.inspect_volume(service_attributes.volume_name) + volume_path = volume_info["Mountpoint"] + + shell = self.get_shell() + shell.exec(f"rm -rf {volume_path}/*") + + def dump_logs( + self, + directory_path: str, + since: Optional[datetime] = None, + until: Optional[datetime] = None, + ) -> None: + client = self._get_docker_client() + for service_config in self._config.services: + container_name = self._get_service_attributes(service_config.name).container_name + try: + logs = client.logs(container_name, since=since, until=until) + except HTTPError as exc: + logger.info(f"Got exception while dumping logs of '{container_name}': {exc}") + continue + + # Save logs to the directory + file_path = os.path.join( + directory_path, + f"{self._config.address}-{container_name}-log.txt", + ) + with open(file_path, "wb") as file: + file.write(logs) + + def _get_service_attributes(self, service_name) -> ServiceAttributes: + service_config = self.get_service_config(service_name) + return ServiceAttributes.parse(service_config.attributes) + + def _get_docker_client(self) -> docker.APIClient: + docker_endpoint = HostAttributes.parse(self._config.attributes).docker_endpoint + + if not docker_endpoint: + # Use default docker client that talks to unix socket + return docker.APIClient() + + # Otherwise use docker client that talks to specified endpoint + return docker.APIClient(base_url=docker_endpoint) + + def _get_container_by_name(self, container_name: str) -> dict[str, Any]: + client = self._get_docker_client() + containers = client.containers(all=True) + + for container in containers: + # Names in local docker environment are prefixed with / + clean_names = set(name.strip("/") for name in container["Names"]) + if container_name in clean_names: + return container + return None + + def _wait_for_container_to_be_in_state( + self, container_name: str, expected_state: str, timeout: int + ) -> None: + iterations = 10 + iteration_wait_time = timeout / iterations + + # To speed things up, we break timeout in smaller iterations and check container state + # several times. This way waiting stops as soon as container reaches the expected state + for _ in range(iterations): + container = self._get_container_by_name(container_name) + logger.debug(f"Current container state\n:{json.dumps(container, indent=2)}") + + if container and container["State"] == expected_state: + return + time.sleep(iteration_wait_time) + + raise RuntimeError(f"Container {container_name} is not in {expected_state} state.") diff --git a/src/neofs_testlib/hosting/hosting.py b/src/neofs_testlib/hosting/hosting.py new file mode 100644 index 0000000..d127f25 --- /dev/null +++ b/src/neofs_testlib/hosting/hosting.py @@ -0,0 +1,107 @@ +import re +from typing import Any + +from neofs_testlib.hosting.config import HostConfig, ServiceConfig +from neofs_testlib.hosting.interfaces import Host +from neofs_testlib.plugins import load_plugin + + +class Hosting: + """Hosting manages infrastructure where neoFS runs (machines and neoFS services).""" + + _hosts: list[Host] + _host_by_address: dict[str, Host] + _host_by_service_name: dict[str, Host] + + @property + def hosts(self) -> list[Host]: + """Returns all hosts registered in the hosting. + + Returns: + List of hosts. + """ + return self._hosts + + def configure(self, config: dict[str, Any]) -> None: + """Configures hosts from specified config. + + All existing hosts will be removed from the hosting. + + Args: + config: Dictionary with hosting configuration. + """ + hosts = [] + host_by_address = {} + host_by_service_name = {} + + host_configs = [HostConfig(**host_config) for host_config in config["hosts"]] + for host_config in host_configs: + host_class = load_plugin("neofs.testlib.hosting", host_config.plugin_name) + host = host_class(host_config) + + hosts.append(host) + host_by_address[host_config.address] = host + + for service_config in host_config.services: + host_by_service_name[service_config.name] = host + + self._hosts = hosts + self._host_by_address = host_by_address + self._host_by_service_name = host_by_service_name + + def get_host_by_address(self, host_address: str) -> Host: + """Returns host with specified address. + + Args: + host_address: Address of the host. + + Returns: + Host that manages machine with specified address. + """ + host = self._host_by_address.get(host_address) + if host is None: + raise ValueError(f"Unknown host address: '{host_address}'") + return host + + def get_host_by_service(self, service_name: str) -> Host: + """Returns host where service with specified name is located. + + Args: + service_name: Name of the service. + + Returns: + Host that manages machine where service is located. + """ + host = self._host_by_service_name.get(service_name) + if host is None: + raise ValueError(f"Unknown service name: '{service_name}'") + return host + + def get_service_config(self, service_name: str) -> ServiceConfig: + """Returns config of service with specified name. + + Args: + service_name: Name of the service. + + Returns: + Config of the service. + """ + host = self.get_host_by_service(service_name) + return host.get_service_config(service_name) + + def find_service_configs(self, service_name_pattern: str) -> list[ServiceConfig]: + """Finds configs of services where service name matches specified regular expression. + + Args: + service_name_pattern - regular expression for service names. + + Returns: + List of service configs matched with the regular expression. + """ + service_configs = [ + service_config + for host in self.hosts + for service_config in host.config.services + if re.match(service_name_pattern, service_config.name) + ] + return service_configs diff --git a/src/neofs_testlib/hosting/interfaces.py b/src/neofs_testlib/hosting/interfaces.py new file mode 100644 index 0000000..b004689 --- /dev/null +++ b/src/neofs_testlib/hosting/interfaces.py @@ -0,0 +1,124 @@ +from abc import ABC, abstractmethod +from datetime import datetime +from typing import Optional + +from neofs_testlib.hosting.config import CLIConfig, HostConfig, ServiceConfig +from neofs_testlib.shell.interfaces import Shell + + +class Host(ABC): + """Interface of a host machine where neoFS services are running. + + Allows to manage the machine and neoFS services that are hosted on it. + """ + + def __init__(self, config: HostConfig) -> None: + self._config = config + self._service_config_by_name = { + service_config.name: service_config for service_config in config.services + } + self._cli_config_by_name = {cli_config.name: cli_config for cli_config in config.clis} + + @property + def config(self) -> HostConfig: + """Returns config of the host. + + Returns: + Config of this host. + """ + return self._config + + def get_service_config(self, service_name: str) -> ServiceConfig: + """Returns config of service with specified name. + + The service must be hosted on this host. + + Args: + service_name: Name of the service. + + Returns: + Config of the service. + """ + service_config = self._service_config_by_name.get(service_name) + if service_config is None: + raise ValueError(f"Unknown service name: '{service_name}'") + return service_config + + def get_cli_config(self, cli_name: str) -> CLIConfig: + """Returns config of CLI tool with specified name. + + The CLI must be located on this host. + + Args: + cli_name: Name of the CLI tool. + + Returns: + Config of the CLI tool. + """ + cli_config = self._cli_config_by_name.get(cli_name) + if cli_config is None: + raise ValueError(f"Unknown CLI name: '{cli_name}'") + return cli_config + + @abstractmethod + def get_shell(self) -> Shell: + """Returns shell to this host. + + Returns: + Shell that executes commands on this host. + """ + + @abstractmethod + def start_host(self) -> None: + """Starts the host machine.""" + + @abstractmethod + def stop_host(self, mode: str) -> None: + """Stops the host machine. + + Args: + mode: Specifies mode how host should be stopped. Mode might be host-specific. + """ + + @abstractmethod + def start_service(self, service_name: str) -> None: + """Starts the service with specified name and waits until it starts. + + The service must be hosted on this host. + + Args: + service_name: Name of the service to start. + """ + + @abstractmethod + def stop_service(self, service_name: str) -> None: + """Stops the service with specified name and waits until it stops. + + The service must be hosted on this host. + + Args: + service_name: Name of the service to stop. + """ + + @abstractmethod + def delete_storage_node_data(self, service_name: str) -> None: + """Erases all data of the storage node with specified name. + + Args: + service_name: Name of storage node service. + """ + + @abstractmethod + def dump_logs( + self, + directory_path: str, + since: Optional[datetime] = None, + until: Optional[datetime] = None, + ) -> None: + """Dumps logs of all services on the host to specified directory. + + Args: + directory_path: Path to the directory where logs should be stored. + since: If set, limits the time from which logs should be collected. Must be in UTC. + until: If set, limits the time until which logs should be collected. Must be in UTC. + """ diff --git a/src/neofs_testlib/reporter/allure_handler.py b/src/neofs_testlib/reporter/allure_handler.py index 9c7f978..0fceffb 100644 --- a/src/neofs_testlib/reporter/allure_handler.py +++ b/src/neofs_testlib/reporter/allure_handler.py @@ -20,7 +20,7 @@ class AllureHandler(ReporterHandler): attachment_name, extension = os.path.splitext(file_name) attachment_type = self._resolve_attachment_type(extension) - allure.attach(body, attachment_name, attachment_type) + allure.attach(body, attachment_name, attachment_type, extension) def _resolve_attachment_type(self, extension: str) -> attachment_type: """Try to find matching Allure attachment type by extension. @@ -30,5 +30,5 @@ class AllureHandler(ReporterHandler): extension = extension.lower() return next( (allure_type for allure_type in attachment_type if allure_type.extension == extension), - attachment_type.TXT, + attachment_type.TEXT, ) diff --git a/src/neofs_testlib/reporter/reporter.py b/src/neofs_testlib/reporter/reporter.py index 3e9e394..d12cb05 100644 --- a/src/neofs_testlib/reporter/reporter.py +++ b/src/neofs_testlib/reporter/reporter.py @@ -34,7 +34,7 @@ class Reporter: All existing handlers will be removed from the reporter. Args: - config: dictionary with reporter configuration. + config: Dictionary with reporter configuration. """ # Reset current configuration self.handlers = [] diff --git a/src/neofs_testlib/shell/__init__.py b/src/neofs_testlib/shell/__init__.py index c51f3b9..3fd63bd 100644 --- a/src/neofs_testlib/shell/__init__.py +++ b/src/neofs_testlib/shell/__init__.py @@ -1,3 +1,3 @@ -from neofs_testlib.shell.interfaces import CommandResult, Shell +from neofs_testlib.shell.interfaces import CommandOptions, CommandResult, Shell from neofs_testlib.shell.local_shell import LocalShell from neofs_testlib.shell.ssh_shell import SSHShell diff --git a/src/neofs_testlib/shell/command_inspectors.py b/src/neofs_testlib/shell/command_inspectors.py new file mode 100644 index 0000000..9537549 --- /dev/null +++ b/src/neofs_testlib/shell/command_inspectors.py @@ -0,0 +1,13 @@ +from neofs_testlib.shell.interfaces import CommandInspector + + +class SudoInspector(CommandInspector): + """Prepends command with sudo. + + If command is already prepended with sudo, then has no effect. + """ + + def inspect(self, command: str) -> str: + if not command.startswith("sudo"): + return f"sudo {command}" + return command diff --git a/src/neofs_testlib/shell/interfaces.py b/src/neofs_testlib/shell/interfaces.py index 4d6e8ac..52e77a3 100644 --- a/src/neofs_testlib/shell/interfaces.py +++ b/src/neofs_testlib/shell/interfaces.py @@ -8,27 +8,46 @@ class InteractiveInput: """Interactive input for a shell command. Attributes: - prompt_pattern: regular expression that defines expected prompt from the command. - input: user input that should be supplied to the command in response to the prompt. + prompt_pattern: Regular expression that defines expected prompt from the command. + input: User input that should be supplied to the command in response to the prompt. """ prompt_pattern: str input: str +class CommandInspector(ABC): + """Interface of inspector that processes command text before execution.""" + + @abstractmethod + def inspect(self, command: str) -> str: + """Transforms command text and returns modified command. + + Args: + command: Command to transform with this inspector. + + Returns: + Transformed command text. + """ + + @dataclass class CommandOptions: """Options that control command execution. Attributes: - interactive_inputs: user inputs that should be interactively supplied to + interactive_inputs: User inputs that should be interactively supplied to the command during execution. - timeout: timeout for command execution (in seconds). - check: controls whether to check return code of the command. Set to False to + close_stdin: Controls whether stdin stream should be closed after feeding interactive + inputs or after requesting non-interactive command. If shell implementation does not + support this functionality, it should ignore this flag without raising an error. + timeout: Timeout for command execution (in seconds). + check: Controls whether to check return code of the command. Set to False to ignore non-zero return codes. """ interactive_inputs: Optional[list[InteractiveInput]] = None + close_stdin: bool = False timeout: int = 30 check: bool = True @@ -38,9 +57,9 @@ class CommandResult: """Represents a result of a command executed via shell. Attributes: - stdout: complete content of stdout stream. - stderr: complete content of stderr stream. - return_code: return code (or exit code) of the command's process. + stdout: Complete content of stdout stream. + stderr: Complete content of stderr stream. + return_code: Return code (or exit code) of the command's process. """ stdout: str diff --git a/src/neofs_testlib/shell/local_shell.py b/src/neofs_testlib/shell/local_shell.py index a329990..b20988c 100644 --- a/src/neofs_testlib/shell/local_shell.py +++ b/src/neofs_testlib/shell/local_shell.py @@ -7,7 +7,7 @@ from typing import IO, Optional import pexpect from neofs_testlib.reporter import get_reporter -from neofs_testlib.shell.interfaces import CommandOptions, CommandResult, Shell +from neofs_testlib.shell.interfaces import CommandInspector, CommandOptions, CommandResult, Shell logger = logging.getLogger("neofs.testlib.shell") reporter = get_reporter() @@ -16,10 +16,17 @@ reporter = get_reporter() class LocalShell(Shell): """Implements command shell on a local machine.""" + def __init__(self, command_inspectors: Optional[list[CommandInspector]] = None) -> None: + super().__init__() + self.command_inspectors = command_inspectors or [] + def exec(self, command: str, options: Optional[CommandOptions] = None) -> CommandResult: # If no options were provided, use default options options = options or CommandOptions() + for inspector in self.command_inspectors: + command = inspector.inspect(command) + logger.info(f"Executing command: {command}") if options.interactive_inputs: return self._exec_interactive(command, options) diff --git a/src/neofs_testlib/shell/ssh_shell.py b/src/neofs_testlib/shell/ssh_shell.py index 967fbbf..d56e2c4 100644 --- a/src/neofs_testlib/shell/ssh_shell.py +++ b/src/neofs_testlib/shell/ssh_shell.py @@ -19,7 +19,7 @@ from paramiko import ( from paramiko.ssh_exception import AuthenticationException from neofs_testlib.reporter import get_reporter -from neofs_testlib.shell.interfaces import CommandOptions, CommandResult, Shell +from neofs_testlib.shell.interfaces import CommandInspector, CommandOptions, CommandResult, Shell logger = logging.getLogger("neofs.testlib.shell") reporter = get_reporter() @@ -97,13 +97,16 @@ class SSHShell(Shell): private_key_path: Optional[str] = None, private_key_passphrase: Optional[str] = None, port: str = "22", + command_inspectors: Optional[list[CommandInspector]] = None, ) -> None: + super().__init__() self.host = host self.port = port self.login = login self.password = password self.private_key_path = private_key_path self.private_key_passphrase = private_key_passphrase + self.command_inspectors = command_inspectors or [] self.__connection: Optional[SSHClient] = None @property @@ -118,6 +121,9 @@ class SSHShell(Shell): def exec(self, command: str, options: Optional[CommandOptions] = None) -> CommandResult: options = options or CommandOptions() + for inspector in self.command_inspectors: + command = inspector.inspect(command) + if options.interactive_inputs: result = self._exec_interactive(command, options) else: @@ -125,8 +131,7 @@ class SSHShell(Shell): if options.check and result.return_code != 0: raise RuntimeError( - f"Command: {command}\nreturn code: {result.return_code}" - f"\nOutput: {result.stdout}" + f"Command: {command}\nreturn code: {result.return_code}\nOutput: {result.stdout}" ) return result @@ -141,7 +146,8 @@ class SSHShell(Shell): stdin.write(input) except OSError: logger.exception(f"Error while feeding {input} into command {command}") - # stdin.close() + if options.close_stdin: + stdin.close() # Wait for command to complete and flush its buffer before we attempt to read output sleep(self.DELAY_AFTER_EXIT) @@ -158,7 +164,10 @@ class SSHShell(Shell): @log_command def _exec_non_interactive(self, command: str, options: CommandOptions) -> CommandResult: try: - _, stdout, stderr = self._connection.exec_command(command, timeout=options.timeout) + stdin, stdout, stderr = self._connection.exec_command(command, timeout=options.timeout) + + if options.close_stdin: + stdin.close() # Wait for command to complete and flush its buffer before we attempt to read output return_code = stdout.channel.recv_exit_status()