Content-Length: 9946 | pFad | http://github.com/postgrespro/testgres/pull/272.diff
thub.com diff --git a/testgres/operations/local_ops.py b/testgres/operations/local_ops.py index e8cac62..103ee4d 100644 --- a/testgres/operations/local_ops.py +++ b/testgres/operations/local_ops.py @@ -67,8 +67,13 @@ def _process_output(encoding, temp_file_path): output = output.decode(encoding) return output, None # In Windows stderr writing in stdout - def _run_command__nt(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, exec_env=None): + def _run_command__nt( + self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, + exec_env: typing.Optional[dict], + cwd: typing.Optional[str], + ): assert exec_env is None or type(exec_env) == dict # noqa: E721 + assert cwd is None or type(cwd) == str # noqa: E721 # TODO: why don't we use the data from input? @@ -104,6 +109,7 @@ def _run_command__nt(self, cmd, shell, input, stdin, stdout, stderr, get_process stdin=stdin or subprocess.PIPE if input is not None else None, stdout=stdout, stderr=stderr, + cwd=cwd, **extParams, ) if get_process: @@ -116,8 +122,13 @@ def _run_command__nt(self, cmd, shell, input, stdin, stdout, stderr, get_process output, error = self._process_output(encoding, temp_file_path) return process, output, error - def _run_command__generic(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, exec_env=None): + def _run_command__generic( + self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, + exec_env: typing.Optional[dict], + cwd: typing.Optional[str], + ): assert exec_env is None or type(exec_env) == dict # noqa: E721 + assert cwd is None or type(cwd) == str # noqa: E721 input_prepared = None if not get_process: @@ -154,6 +165,7 @@ def _run_command__generic(self, cmd, shell, input, stdin, stdout, stderr, get_pr stdin=stdin or subprocess.PIPE if input is not None else None, stdout=stdout or subprocess.PIPE, stderr=stderr or subprocess.PIPE, + cwd=cwd, **extParams ) assert not (process is None) @@ -173,26 +185,44 @@ def _run_command__generic(self, cmd, shell, input, stdin, stdout, stderr, get_pr error = error.decode(encoding) return process, output, error - def _run_command(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, exec_env=None): + def _run_command( + self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, + exec_env: typing.Optional[dict], + cwd: typing.Optional[str], + ): """Execute a command and return the process and its output.""" + + assert exec_env is None or type(exec_env) == dict # noqa: E721 + assert cwd is None or type(cwd) == str # noqa: E721 + if os.name == 'nt' and stdout is None: # Windows method = __class__._run_command__nt else: # Other OS method = __class__._run_command__generic - return method(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, exec_env=exec_env) + return method(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, exec_env, cwd) - def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False, encoding=None, shell=False, - text=False, input=None, stdin=None, stdout=None, stderr=None, get_process=False, timeout=None, - ignore_errors=False, exec_env=None): + def exec_command( + self, cmd, wait_exit=False, verbose=False, expect_error=False, encoding=None, shell=False, + text=False, input=None, stdin=None, stdout=None, stderr=None, get_process=False, timeout=None, + ignore_errors=False, + exec_env: typing.Optional[dict] = None, + cwd: typing.Optional[str] = None + ): """ Execute a command in a subprocess and handle the output based on the provided parameters. """ assert type(expect_error) == bool # noqa: E721 assert type(ignore_errors) == bool # noqa: E721 assert exec_env is None or type(exec_env) == dict # noqa: E721 + assert cwd is None or type(cwd) == str # noqa: E721 + + process, output, error = self._run_command( + cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, + exec_env, + cwd + ) - process, output, error = self._run_command(cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, exec_env=exec_env) if get_process: return process diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py index e137b8e..8b973a1 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -63,10 +63,13 @@ def __init__(self, conn_params: ConnectionParams): def __enter__(self): return self - def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False, - encoding=None, shell=True, text=False, input=None, stdin=None, stdout=None, - stderr=None, get_process=None, timeout=None, ignore_errors=False, - exec_env=None): + def exec_command( + self, cmd, wait_exit=False, verbose=False, expect_error=False, + encoding=None, shell=True, text=False, input=None, stdin=None, stdout=None, + stderr=None, get_process=None, timeout=None, ignore_errors=False, + exec_env: typing.Optional[dict] = None, + cwd: typing.Optional[str] = None + ): """ Execute a command in the SSH session. Args: @@ -75,6 +78,7 @@ def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False, assert type(expect_error) == bool # noqa: E721 assert type(ignore_errors) == bool # noqa: E721 assert exec_env is None or type(exec_env) == dict # noqa: E721 + assert cwd is None or type(cwd) == str # noqa: E721 input_prepared = None if not get_process: @@ -82,21 +86,21 @@ def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False, assert input_prepared is None or (type(input_prepared) == bytes) # noqa: E721 - if type(cmd) == str: # noqa: E721 - cmd_s = cmd - elif type(cmd) == list: # noqa: E721 - cmd_s = subprocess.list2cmdline(cmd) - else: - raise ValueError("Invalid 'cmd' argument type - {0}".format(type(cmd).__name__)) + cmds = [] - assert type(cmd_s) == str # noqa: E721 + if cwd is not None: + assert type(cwd) == str # noqa: E721 + cmds.append(__class__._build_cmdline(["cd", cwd])) + + cmds.append(__class__._build_cmdline(cmd, exec_env)) - cmd_items = __class__._make_exec_env_list(exec_env=exec_env) - cmd_items.append(cmd_s) + assert len(cmds) >= 1 - env_cmd_s = ';'.join(cmd_items) + cmdline = ";".join(cmds) + assert type(cmdline) == str # noqa: E721 + assert cmdline != "" - ssh_cmd = ['ssh', self.ssh_dest] + self.ssh_args + [env_cmd_s] + ssh_cmd = ['ssh', self.ssh_dest] + self.ssh_args + [cmdline] process = subprocess.Popen(ssh_cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE) assert not (process is None) @@ -718,7 +722,31 @@ def _is_port_free__process_1(error: str) -> bool: return True @staticmethod - def _make_exec_env_list(exec_env: typing.Dict) -> typing.List[str]: + def _build_cmdline(cmd, exec_env: typing.Dict = None) -> str: + cmd_items = __class__._create_exec_env_list(exec_env) + + assert type(cmd_items) == list # noqa: E721 + + cmd_items.append(__class__._ensure_cmdline(cmd)) + + cmdline = ';'.join(cmd_items) + assert type(cmdline) == str # noqa: E721 + return cmdline + + @staticmethod + def _ensure_cmdline(cmd) -> typing.List[str]: + if type(cmd) == str: # noqa: E721 + cmd_s = cmd + elif type(cmd) == list: # noqa: E721 + cmd_s = subprocess.list2cmdline(cmd) + else: + raise ValueError("Invalid 'cmd' argument type - {0}".format(type(cmd).__name__)) + + assert type(cmd_s) == str # noqa: E721 + return cmd_s + + @staticmethod + def _create_exec_env_list(exec_env: typing.Dict) -> typing.List[str]: env: typing.Dict[str, str] = dict() # ---------------------------------- SYSTEM ENV diff --git a/tests/test_os_ops_common.py b/tests/test_os_ops_common.py index 5ae3a61..5b135a4 100644 --- a/tests/test_os_ops_common.py +++ b/tests/test_os_ops_common.py @@ -118,6 +118,23 @@ def test_exec_command_with_exec_env(self, os_ops: OsOperations): assert type(response) == bytes # noqa: E721 assert response == b'\n' + def test_exec_command_with_cwd(self, os_ops: OsOperations): + assert isinstance(os_ops, OsOperations) + + RunConditions.skip_if_windows() + + cmd = ["pwd"] + + response = os_ops.exec_command(cmd, cwd="/tmp") + assert response is not None + assert type(response) == bytes # noqa: E721 + assert response == b'/tmp\n' + + response = os_ops.exec_command(cmd) + assert response is not None + assert type(response) == bytes # noqa: E721 + assert response != b'/tmp\n' + def test_exec_command__test_unset(self, os_ops: OsOperations): assert isinstance(os_ops, OsOperations)Fetched URL: http://github.com/postgrespro/testgres/pull/272.diff
Alternative Proxies: