From ee9a31b217c3de906c3e34e5f38365b28609a732 Mon Sep 17 00:00:00 2001 From: "d.kovalenko" Date: Mon, 14 Apr 2025 19:17:23 +0300 Subject: [PATCH 1/2] [New] OsOps::execute_command supports a transfer of environment variables (exec_env) New feature allows to pass environment variables to an executed program. If variable in exec_env has None value, then this variable will be unset. PostgresNode::start and PostgresNode::slow_start supports exec_env. This feature is required for internal test projects... --- testgres/node.py | 12 +++--- testgres/operations/local_ops.py | 66 ++++++++++++++++++++++++++++--- testgres/operations/remote_ops.py | 44 ++++++++++++++++----- testgres/utils.py | 13 +++++- tests/test_os_ops_common.py | 64 ++++++++++++++++++++++++++++++ 5 files changed, 177 insertions(+), 22 deletions(-) diff --git a/testgres/node.py b/testgres/node.py index 5039fc43..3a294044 100644 --- a/testgres/node.py +++ b/testgres/node.py @@ -1020,7 +1020,7 @@ def get_control_data(self): return out_dict - def slow_start(self, replica=False, dbname='template1', username=None, max_attempts=0): + def slow_start(self, replica=False, dbname='template1', username=None, max_attempts=0, exec_env=None): """ Starts the PostgreSQL instance and then polls the instance until it reaches the expected state (primary or replica). The state is checked @@ -1033,7 +1033,9 @@ def slow_start(self, replica=False, dbname='template1', username=None, max_attem If False, waits for the instance to be in primary mode. Default is False. max_attempts: """ - self.start() + assert exec_env is None or type(exec_env) == dict # noqa: E721 + + self.start(exec_env=exec_env) if replica: query = 'SELECT pg_is_in_recovery()' @@ -1065,7 +1067,7 @@ def _detect_port_conflict(self, log_files0, log_files1): return True return False - def start(self, params=[], wait=True): + def start(self, params=[], wait=True, exec_env=None): """ Starts the PostgreSQL node using pg_ctl if node has not been started. By default, it waits for the operation to complete before returning. @@ -1079,7 +1081,7 @@ def start(self, params=[], wait=True): Returns: This instance of :class:`.PostgresNode`. """ - + assert exec_env is None or type(exec_env) == dict # noqa: E721 assert __class__._C_MAX_START_ATEMPTS > 1 if self.is_started: @@ -1098,7 +1100,7 @@ def start(self, params=[], wait=True): def LOCAL__start_node(): # 'error' will be None on Windows - _, _, error = execute_utility2(self.os_ops, _params, self.utils_log_file, verbose=True) + _, _, error = execute_utility2(self.os_ops, _params, self.utils_log_file, verbose=True, exec_env=exec_env) assert error is None or type(error) == str # noqa: E721 if error and 'does not exist' in error: raise Exception(error) diff --git a/testgres/operations/local_ops.py b/testgres/operations/local_ops.py index 9785d462..74323bb8 100644 --- a/testgres/operations/local_ops.py +++ b/testgres/operations/local_ops.py @@ -9,6 +9,7 @@ import socket import psutil +import typing from ..exceptions import ExecUtilException from ..exceptions import InvalidOperationException @@ -46,9 +47,34 @@ 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): + def _run_command__nt(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, exec_env=None): + assert exec_env is None or type(exec_env) == dict # noqa: E721 + # TODO: why don't we use the data from input? + extParams: typing.Dict[str, str] = dict() + + if exec_env is None: + pass + elif len(exec_env) == 0: + pass + else: + env = os.environ.copy() + assert type(env) == dict # noqa: E721 + for v in exec_env.items(): + assert type(v) == tuple # noqa: E721 + assert len(v) == 2 + assert type(v[0]) == str # noqa: E721 + assert v[0] != "" + + if v[1] is None: + env.pop(v[0], None) + else: + assert type(v[1]) == str # noqa: E721 + env[v[0]] = v[1] + + extParams["env"] = env + with tempfile.NamedTemporaryFile(mode='w+b', delete=False) as temp_file: stdout = temp_file stderr = subprocess.STDOUT @@ -58,6 +84,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, + **extParams, ) if get_process: return process, None, None @@ -69,19 +96,45 @@ 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): + def _run_command__generic(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, exec_env=None): + assert exec_env is None or type(exec_env) == dict # noqa: E721 + input_prepared = None if not get_process: input_prepared = Helpers.PrepareProcessInput(input, encoding) # throw assert input_prepared is None or (type(input_prepared) == bytes) # noqa: E721 + extParams: typing.Dict[str, str] = dict() + + if exec_env is None: + pass + elif len(exec_env) == 0: + pass + else: + env = os.environ.copy() + assert type(env) == dict # noqa: E721 + for v in exec_env.items(): + assert type(v) == tuple # noqa: E721 + assert len(v) == 2 + assert type(v[0]) == str # noqa: E721 + assert v[0] != "" + + if v[1] is None: + env.pop(v[0], None) + else: + assert type(v[1]) == str # noqa: E721 + env[v[0]] = v[1] + + extParams["env"] = env + process = subprocess.Popen( cmd, shell=shell, stdin=stdin or subprocess.PIPE if input is not None else None, stdout=stdout or subprocess.PIPE, stderr=stderr or subprocess.PIPE, + **extParams ) assert not (process is None) if get_process: @@ -100,25 +153,26 @@ 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): + def _run_command(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, exec_env=None): """Execute a command and return the process and its output.""" 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) + return method(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, exec_env=exec_env) 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): + ignore_errors=False, exec_env=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 - process, output, error = self._run_command(cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding) + 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 33b61ac2..e722a2cb 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -64,7 +64,8 @@ def __enter__(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): + stderr=None, get_process=None, timeout=None, ignore_errors=False, + exec_env=None): """ Execute a command in the SSH session. Args: @@ -72,6 +73,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 input_prepared = None if not get_process: @@ -88,7 +90,7 @@ def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False, assert type(cmd_s) == str # noqa: E721 - cmd_items = __class__._make_exec_env_list() + cmd_items = __class__._make_exec_env_list(exec_env=exec_env) cmd_items.append(cmd_s) env_cmd_s = ';'.join(cmd_items) @@ -670,14 +672,38 @@ def _is_port_free__process_1(error: str) -> bool: return True @staticmethod - def _make_exec_env_list() -> typing.List[str]: - result: typing.List[str] = list() + def _make_exec_env_list(exec_env: typing.Dict) -> typing.List[str]: + env: typing.Dict[str, str] = dict() + + # ---------------------------------- SYSTEM ENV for envvar in os.environ.items(): - if not __class__._does_put_envvar_into_exec_cmd(envvar[0]): - continue - qvalue = __class__._quote_envvar(envvar[1]) - assert type(qvalue) == str # noqa: E721 - result.append(envvar[0] + "=" + qvalue) + if __class__._does_put_envvar_into_exec_cmd(envvar[0]): + env[envvar[0]] = envvar[1] + + # ---------------------------------- EXEC (LOCAL) ENV + if exec_env is None: + pass + else: + for envvar in exec_env.items(): + assert type(envvar) == tuple # noqa: E721 + assert len(envvar) == 2 + assert type(envvar[0]) == str # noqa: E721 + env[envvar[0]] = envvar[1] + + # ---------------------------------- FINAL BUILD + result: typing.List[str] = list() + for envvar in env.items(): + assert type(envvar) == tuple # noqa: E721 + assert len(envvar) == 2 + assert type(envvar[0]) == str # noqa: E721 + + if envvar[1] is None: + result.append("unset " + envvar[0]) + else: + assert type(envvar[1]) == str # noqa: E721 + qvalue = __class__._quote_envvar(envvar[1]) + assert type(qvalue) == str # noqa: E721 + result.append(envvar[0] + "=" + qvalue) continue return result diff --git a/testgres/utils.py b/testgres/utils.py index 10ae81b6..2ff6f2a0 100644 --- a/testgres/utils.py +++ b/testgres/utils.py @@ -96,17 +96,26 @@ def execute_utility(args, logfile=None, verbose=False): return execute_utility2(tconf.os_ops, args, logfile, verbose) -def execute_utility2(os_ops: OsOperations, args, logfile=None, verbose=False, ignore_errors=False): +def execute_utility2( + os_ops: OsOperations, + args, + logfile=None, + verbose=False, + ignore_errors=False, + exec_env=None, +): assert os_ops is not None assert isinstance(os_ops, OsOperations) assert type(verbose) == bool # noqa: E721 assert type(ignore_errors) == bool # noqa: E721 + assert exec_env is None or type(exec_env) == dict # noqa: E721 exit_status, out, error = os_ops.exec_command( args, verbose=True, ignore_errors=ignore_errors, - encoding=OsHelpers.GetDefaultEncoding()) + encoding=OsHelpers.GetDefaultEncoding(), + exec_env=exec_env) out = '' if not out else out diff --git a/tests/test_os_ops_common.py b/tests/test_os_ops_common.py index ecfff5b2..7bec13d6 100644 --- a/tests/test_os_ops_common.py +++ b/tests/test_os_ops_common.py @@ -93,6 +93,70 @@ def test_exec_command_failure__expect_error(self, os_ops: OsOperations): assert b"nonexistent_command" in error assert b"not found" in error + def test_exec_command_with_exec_env(self, os_ops: OsOperations): + assert isinstance(os_ops, OsOperations) + + RunConditions.skip_if_windows() + + C_ENV_NAME = "EXEC_TEST_TESTGRES_ENV_1975" + + cmd = ["sh", "-c", "echo ${}".format(C_ENV_NAME)] + + exec_env = {C_ENV_NAME: "Hello!"} + + response = os_ops.exec_command(cmd, exec_env=exec_env) + assert response is not None + assert type(response) == bytes # noqa: E721 + assert response == b'Hello!\n' + + response = os_ops.exec_command(cmd) + assert response is not None + assert type(response) == bytes # noqa: E721 + assert response == b'\n' + + def test_exec_command__test_unset(self, os_ops: OsOperations): + assert isinstance(os_ops, OsOperations) + + RunConditions.skip_if_windows() + + C_ENV_NAME = "LANG" + + cmd = ["sh", "-c", "echo ${}".format(C_ENV_NAME)] + + response1 = os_ops.exec_command(cmd) + assert response1 is not None + assert type(response1) == bytes # noqa: E721 + + if response1 == b'\n': + logging.warning("Environment variable {} is not defined.".format(C_ENV_NAME)) + return + + exec_env = {C_ENV_NAME: None} + response2 = os_ops.exec_command(cmd, exec_env=exec_env) + assert response2 is not None + assert type(response2) == bytes # noqa: E721 + assert response2 == b'\n' + + response3 = os_ops.exec_command(cmd) + assert response3 is not None + assert type(response3) == bytes # noqa: E721 + assert response3 == response1 + + def test_exec_command__test_unset_dummy_var(self, os_ops: OsOperations): + assert isinstance(os_ops, OsOperations) + + RunConditions.skip_if_windows() + + C_ENV_NAME = "IT_IS_A_TEST_DUMMY_VAR" + + cmd = ["sh", "-c", "echo ${}".format(C_ENV_NAME)] + + exec_env = {C_ENV_NAME: None} + response2 = os_ops.exec_command(cmd, exec_env=exec_env) + assert response2 is not None + assert type(response2) == bytes # noqa: E721 + assert response2 == b'\n' + def test_is_executable_true(self, os_ops: OsOperations): """ Test is_executable for an existing executable. From 117165f5c8b275057c3df857b6f23366a0c0eb39 Mon Sep 17 00:00:00 2001 From: "d.kovalenko" Date: Tue, 15 Apr 2025 08:01:46 +0300 Subject: [PATCH 2/2] TestOsOpsCommon is refactored (test env names) --- tests/test_os_ops_common.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_os_ops_common.py b/tests/test_os_ops_common.py index 7bec13d6..17c3151c 100644 --- a/tests/test_os_ops_common.py +++ b/tests/test_os_ops_common.py @@ -98,7 +98,7 @@ def test_exec_command_with_exec_env(self, os_ops: OsOperations): RunConditions.skip_if_windows() - C_ENV_NAME = "EXEC_TEST_TESTGRES_ENV_1975" + C_ENV_NAME = "TESTGRES_TEST__EXEC_ENV_20250414" cmd = ["sh", "-c", "echo ${}".format(C_ENV_NAME)] @@ -147,7 +147,7 @@ def test_exec_command__test_unset_dummy_var(self, os_ops: OsOperations): RunConditions.skip_if_windows() - C_ENV_NAME = "IT_IS_A_TEST_DUMMY_VAR" + C_ENV_NAME = "TESTGRES_TEST__DUMMY_VAR_20250414" cmd = ["sh", "-c", "echo ${}".format(C_ENV_NAME)]